tech/create_qdrant_collection.py

#!/usr/bin/env python3
"""
Скрипт для создания коллекции в Qdrant для хранения эмбеддингов книг.

Структура данных основана на epub-parser/app.py:
- Каждая точка содержит эмбеддинг чанка текста
- Payload содержит метаданные книги, главы и чанка
"""
import requests
import json
from typing import Dict, Any

# Конфигурация Qdrant
QDRANT_URL = "http://localhost:6333"
COLLECTION_NAME = "epub_books"
VECTOR_SIZE = 768  # Размерность эмбеддингов nomic-embed-text


def create_collection() -> Dict[str, Any]:
    """Создает коллекцию в Qdrant для хранения эмбеддингов книг.
    
    Returns:
        Ответ от Qdrant API с результатом создания коллекции.
    """
    url = f"{QDRANT_URL}/collections/{COLLECTION_NAME}"
    
    payload = {
        "vectors": {
            "size": VECTOR_SIZE,
            "distance": "Cosine"  # Косинусное расстояние для семантического поиска
        },
        "optimizers_config": {
            "default_segment_number": 2
        },
        "replication_factor": 1,
        "write_consistency_factor": 1
    }
    
    response = requests.put(url, json=payload)
    response.raise_for_status()
    return response.json()


def create_payload_indexes() -> Dict[str, Any]:
    """Создает индексы для полей payload для быстрого поиска.
    
    Returns:
        Список ответов от Qdrant API.
    """
    results = []
    
    # Индексы для фильтрации и поиска
    indexes = [
        {
            "field_name": "bookId",
            "field_schema": "keyword"
        },
        {
            "field_name": "title",
            "field_schema": "keyword"
        },
        {
            "field_name": "author",
            "field_schema": "keyword"
        },
        {
            "field_name": "chapterNumber",
            "field_schema": "integer"
        },
        {
            "field_name": "chapterTitle",
            "field_schema": "text"
        },
        {
            "field_name": "chapterId",
            "field_schema": "keyword"
        },
        {
            "field_name": "chunkIndex",
            "field_schema": "integer"
        }
    ]
    
    for index in indexes:
        url = f"{QDRANT_URL}/collections/{COLLECTION_NAME}/index"
        payload = {
            "field_name": index["field_name"],
            "field_schema": {
                "type": index["field_schema"]
            }
        }
        try:
            response = requests.post(url, json=payload)
            response.raise_for_status()
            results.append({
                "field": index["field_name"],
                "status": "created",
                "response": response.json()
            })
        except requests.exceptions.HTTPError as e:
            if e.response.status_code == 400:
                # Индекс уже существует
                results.append({
                    "field": index["field_name"],
                    "status": "already_exists"
                })
            else:
                raise
    
    return results


def check_collection_exists() -> bool:
    """Проверяет, существует ли коллекция.
    
    Returns:
        True если коллекция существует, False иначе.
    """
    url = f"{QDRANT_URL}/collections/{COLLECTION_NAME}"
    try:
        response = requests.get(url)
        return response.status_code == 200
    except:
        return False


def get_collection_info() -> Dict[str, Any]:
    """Получает информацию о коллекции.
    
    Returns:
        Информация о коллекции от Qdrant API.
    """
    url = f"{QDRANT_URL}/collections/{COLLECTION_NAME}"
    response = requests.get(url)
    response.raise_for_status()
    return response.json()


def main():
    """Основная функция для создания коллекции и индексов."""
    print(f"Создание коллекции '{COLLECTION_NAME}' в Qdrant...")
    print(f"URL: {QDRANT_URL}")
    print(f"Размерность векторов: {VECTOR_SIZE}")
    print()
    
    # Проверяем, существует ли коллекция
    if check_collection_exists():
        print(f"⚠️  Коллекция '{COLLECTION_NAME}' уже существует.")
        response = input("Удалить и создать заново? (y/n): ")
        if response.lower() == 'y':
            url = f"{QDRANT_URL}/collections/{COLLECTION_NAME}"
            requests.delete(url)
            print("✅ Коллекция удалена.")
        else:
            print("Используем существующую коллекцию.")
            info = get_collection_info()
            print(json.dumps(info, indent=2, ensure_ascii=False))
            return
    
    # Создаем коллекцию
    try:
        result = create_collection()
        print("✅ Коллекция успешно создана!")
        print(json.dumps(result, indent=2, ensure_ascii=False))
        print()
        
        # Создаем индексы
        print("Создание индексов для payload...")
        index_results = create_payload_indexes()
        for idx_result in index_results:
            status = idx_result.get("status", "unknown")
            field = idx_result.get("field", "unknown")
            if status == "created":
                print(f"  ✅ Индекс для '{field}' создан")
            elif status == "already_exists":
                print(f"  ⚠️  Индекс для '{field}' уже существует")
        print()
        
        # Получаем информацию о коллекции
        print("Информация о коллекции:")
        info = get_collection_info()
        print(json.dumps(info, indent=2, ensure_ascii=False))
        
    except requests.exceptions.ConnectionError:
        print("❌ Ошибка: Не удалось подключиться к Qdrant.")
        print(f"   Убедитесь, что Qdrant запущен на {QDRANT_URL}")
    except requests.exceptions.HTTPError as e:
        print(f"❌ Ошибка HTTP: {e}")
        if e.response.text:
            print(f"   Ответ сервера: {e.response.text}")
    except Exception as e:
        print(f"❌ Неожиданная ошибка: {e}")


if __name__ == "__main__":
    main()