init

create_qdrant_collection.py

@@ -0,0 +1,196 @@
+#!/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()