88 lines
3.2 KiB
Python
88 lines
3.2 KiB
Python
#!/usr/bin/env python3
|
|
# -*- coding: utf-8 -*-
|
|
"""
|
|
Скрипт для создания коллекции в Qdrant на основе конфигурации профиля бота.
|
|
Использует BotConfig для чтения параметров подключения к Qdrant.
|
|
"""
|
|
|
|
import sys
|
|
import argparse
|
|
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.http import models
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(
|
|
description="Создать коллекцию в Qdrant для профиля бота"
|
|
)
|
|
parser.add_argument(
|
|
"collection_name",
|
|
nargs="?",
|
|
help="Имя коллекции (переопределяет значение из конфига)"
|
|
)
|
|
parser.add_argument(
|
|
"--profile-dir",
|
|
required=True,
|
|
help="Путь к директории профиля (содержит bot.conf)"
|
|
)
|
|
parser.add_argument(
|
|
"--force",
|
|
action="store_true",
|
|
help="Удалить существующую коллекцию, если она есть, и создать заново"
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
# Загружаем конфигурацию профиля
|
|
try:
|
|
config = BotConfig(args.profile_dir)
|
|
except Exception as e:
|
|
print(f"Ошибка загрузки конфигурации: {e}")
|
|
sys.exit(1)
|
|
|
|
# Определяем имя коллекции
|
|
collection = args.collection_name or config.qdrant_collection
|
|
if not collection:
|
|
print("Ошибка: не указано имя коллекции и нет qdrant_collection в конфиге")
|
|
sys.exit(1)
|
|
|
|
# Параметры подключения к Qdrant
|
|
host = config.qdrant_host
|
|
port = config.qdrant_port
|
|
vector_size = config.qdrant_vector_size
|
|
distance = config.qdrant_distance
|
|
|
|
print(f"Подключение к Qdrant: {host}:{port}")
|
|
client = QdrantClient(host=host, port=port)
|
|
|
|
# Проверяем существование коллекции
|
|
if client.collection_exists(collection):
|
|
if args.force:
|
|
print(f"Коллекция '{collection}' существует. Удаляем...")
|
|
client.delete_collection(collection)
|
|
else:
|
|
print(f"Коллекция '{collection}' уже существует. Используйте --force для пересоздания.")
|
|
sys.exit(0)
|
|
|
|
# Создаём коллекцию
|
|
try:
|
|
client.create_collection(
|
|
collection_name=collection,
|
|
vectors_config=models.VectorParams(
|
|
size=vector_size,
|
|
distance=models.Distance[distance.upper()]
|
|
)
|
|
)
|
|
print(f"Коллекция '{collection}' успешно создана (size={vector_size}, distance={distance})")
|
|
except Exception as e:
|
|
print(f"Ошибка создания коллекции: {e}")
|
|
sys.exit(1)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main() |