Редактировать create_qdrant_collection.py
This commit is contained in:
@@ -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)
|
||||||
|
|||||||
Reference in New Issue
Block a user