Добавить init_knowledge_base.py

init_knowledge_base.py

@@ -0,0 +1,73 @@
+"""
+init_knowledge_base.py — скрипт инициализации базы знаний.
+
+Загружает все текстовые документы из указанной директории
+в векторное хранилище Qdrant.
+
+Использование:
+    python init_knowledge_base.py [директория]
+
+По умолчанию — папка ./docs
+"""
+
+import sys
+import os
+from pathlib import Path
+
+from vector_store import init_collection, add_documents
+
+# Поддерживаемые расширения
+TEXT_EXTENSIONS = {".txt", ".md", ".rst", ".csv", ".json"}
+
+
+def load_directory(docs_dir: str) -> None:
+    """
+    Рекурсивно обходит директорию и загружает все текстовые файлы в Qdrant.
+    """
+    base = Path(docs_dir)
+    if not base.exists():
+        print(f"[init] Директория «{docs_dir}» не найдена. Создаю пример.")
+        base.mkdir(parents=True, exist_ok=True)
+        # Создаём демонстрационный файл
+        sample = base / "sample.txt"
+        sample.write_text(
+            "LangChain — это фреймворк для построения приложений на основе больших языковых моделей. "
+            "Он предоставляет инструменты для цепочек, агентов, памяти и интеграции с внешними источниками данных.\n\n"
+            "Qdrant — это высокопроизводительная векторная база данных с открытым исходным кодом, "
+            "разработанная специально для хранения и поиска плотных векторов. "
+            "Поддерживает фильтрацию по метаданным и масштабируется горизонтально.\n\n"
+            "Ollama позволяет запускать большие языковые модели локально на вашем компьютере "
+            "без подключения к внешним API. Поддерживает модели llama3, mistral, gemma и другие.",
+            encoding="utf-8",
+        )
+        print(f"[init] Создан демонстрационный файл: {sample}")
+
+    files = [f for f in base.rglob("*") if f.suffix.lower() in TEXT_EXTENSIONS]
+
+    if not files:
+        print(f"[init] В директории «{docs_dir}» нет поддерживаемых файлов.")
+        return
+
+    print(f"[init] Найдено файлов: {len(files)}")
+    print("[init] Инициализация коллекции Qdrant...")
+    init_collection()
+
+    total_chunks = 0
+    for file in files:
+        try:
+            content = file.read_text(encoding="utf-8", errors="ignore")
+            if not content.strip():
+                print(f"  ⚠  Пропущен пустой файл: {file.name}")
+                continue
+            n = add_documents(content, title=file.name)
+            total_chunks += n
+            print(f"  ✓  {file.name} → {n} чанков")
+        except Exception as e:
+            print(f"  ✗  Ошибка при обработке {file.name}: {e}")
+
+    print(f"\n[init] Загрузка завершена. Всего добавлено чанков: {total_chunks}")
+
+
+if __name__ == "__main__":
+    directory = sys.argv[1] if len(sys.argv) > 1 else "./docs"
+    load_directory(directory)