Обновить init_knowledge_base.py

init_knowledge_base.py

@@ -1,73 +1,62 @@
 """
-init_knowledge_base.py — скрипт инициализации базы знаний.
+init_knowledge_base.py
 
-Загружает все текстовые документы из указанной директории
-в векторное хранилище Qdrant.
+Script for loading documents from a directory into the vector store.
 
-Использование:
-    python init_knowledge_base.py [директория]
+Usage:
+    python init_knowledge_base.py <documents_directory>
 
-По умолчанию — папка ./docs
+Supported file types: .txt, .md
 """
 
 import sys
 import os
 from pathlib import Path
+from vector_store import add_documents
 
-from vector_store import init_collection, add_documents
-
-# Поддерживаемые расширения
-TEXT_EXTENSIONS = {".txt", ".md", ".rst", ".csv", ".json"}
+SUPPORTED_EXTENSIONS = {".txt", ".md"}
 
 
-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}")
+def load_documents_from_directory(directory: str) -> None:
+    dir_path = Path(directory)
+    if not dir_path.exists() or not dir_path.is_dir():
+        print(f"Error: '{directory}' is not a valid directory.")
+        sys.exit(1)
 
-    files = [f for f in base.rglob("*") if f.suffix.lower() in TEXT_EXTENSIONS]
+    files = [
+        f for f in dir_path.rglob("*")
+        if f.is_file() and f.suffix.lower() in SUPPORTED_EXTENSIONS
+    ]
 
     if not files:
-        print(f"[init] В директории «{docs_dir}» нет поддерживаемых файлов.")
+        print(f"No supported files found in '{directory}'.")
         return
 
-    print(f"[init] Найдено файлов: {len(files)}")
-    print("[init] Инициализация коллекции Qdrant...")
-    init_collection()
-
+    print(f"Found {len(files)} file(s) to load.\n")
     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}")
+    for file_path in files:
+        try:
+            content = file_path.read_text(encoding="utf-8")
+            title = file_path.stem
+            num_chunks = add_documents(content=content, title=title)
+            total_chunks += num_chunks
+            print(f"  [OK] '{file_path.name}' -> {num_chunks} chunk(s) added.")
+        except Exception as e:
+            print(f"  [ERROR] '{file_path.name}': {e}")
+
+    print(f"\nDone. Total chunks added: {total_chunks}")
+
+
+def main():
+    if len(sys.argv) < 2:
+        print("Usage: python init_knowledge_base.py <documents_directory>")
+        sys.exit(1)
+
+    directory = sys.argv[1]
+    print(f"Loading documents from: {directory}\n")
+    load_documents_from_directory(directory)
 
 
 if __name__ == "__main__":
-    directory = sys.argv[1] if len(sys.argv) > 1 else "./docs"
-    load_directory(directory)
+    main()