Редактировать create_qdrant_collection.py

This commit is contained in:
Markov Andrey
2026-06-30 08:50:15 +00:00
parent e3cc5c9f03
commit e0d4a29efa

View File

@@ -1,35 +1,37 @@
#!/usr/bin/env python3 #!/usr/bin/env python3
# -*- coding: utf-8 -*- # -*- coding: utf-8 -*-
""" """
Скрипт для создания коллекции в Qdrant на основе конфигурации профиля бота. Скрипт для создания коллекции в Qdrant на основе core.conf.
Использует BotConfig для чтения параметров подключения к Qdrant. Используется для единой коллекции всех ботов.
""" """
import sys import sys
import argparse import argparse
import yaml
from pathlib import Path from pathlib import Path
# Добавляем путь к core (чтобы импортировать BotConfig)
sys.path.insert(0, str(Path(__file__).parent.parent))
from core.utils.config_loader import BotConfig
from qdrant_client import QdrantClient from qdrant_client import QdrantClient
from qdrant_client.http import models from qdrant_client.http import models
def load_config(config_path: str) -> dict:
"""Загружает core.conf как YAML."""
with open(config_path, 'r', encoding='utf-8') as f:
return yaml.safe_load(f)
def main(): def main():
parser = argparse.ArgumentParser( parser = argparse.ArgumentParser(
description="Создать коллекцию в Qdrant для профиля бота" description="Создать коллекцию в Qdrant для RAG-сервера"
) )
parser.add_argument( parser.add_argument(
"collection_name", "--config",
nargs="?", default="core/core.conf",
help="Имя коллекции (переопределяет значение из конфига)" help="Путь к core.conf (по умолчанию core/core.conf)"
) )
parser.add_argument( parser.add_argument(
"--profile-dir", "--collection",
required=True, help="Имя коллекции (переопределяет значение из core.conf)"
help="Путь к директории профиля (содержит bot.conf)"
) )
parser.add_argument( parser.add_argument(
"--force", "--force",
@@ -38,24 +40,33 @@ def main():
) )
args = parser.parse_args() args = parser.parse_args()
# Загружаем конфигурацию профиля # Загружаем конфигурацию
config_path = Path(args.config)
if not config_path.exists():
print(f"Ошибка: файл {config_path} не найден")
sys.exit(1)
try: try:
config = BotConfig(args.profile_dir) config = load_config(config_path)
except Exception as e: except Exception as e:
print(f"Ошибка загрузки конфигурации: {e}") print(f"Ошибка загрузки конфигурации: {e}")
sys.exit(1) sys.exit(1)
# Определяем имя коллекции # Извлекаем параметры Qdrant
collection = args.collection_name or config.qdrant_collection qdrant_cfg = config.get('qdrant')
if not collection: if not qdrant_cfg:
print("Ошибка: не указано имя коллекции и нет qdrant_collection в конфиге") print("Ошибка: в конфиге отсутствует секция 'qdrant'")
sys.exit(1) sys.exit(1)
# Параметры подключения к Qdrant host = qdrant_cfg.get('host', 'localhost')
host = config.qdrant_host port = qdrant_cfg.get('port', 6333)
port = config.qdrant_port vector_size = qdrant_cfg.get('vector_size', 1024)
vector_size = config.qdrant_vector_size distance = qdrant_cfg.get('distance', 'Cosine')
distance = config.qdrant_distance collection = args.collection or qdrant_cfg.get('collection')
if not collection:
print("Ошибка: не указано имя коллекции (в core.conf или через --collection)")
sys.exit(1)
print(f"Подключение к Qdrant: {host}:{port}") print(f"Подключение к Qdrant: {host}:{port}")
client = QdrantClient(host=host, port=port) client = QdrantClient(host=host, port=port)