sync from accepted repo: init_knowledge.py

init_knowledge.py

@@ -4,28 +4,30 @@ from __future__ import annotations
 import argparse
 from pathlib import Path
 
-try:
-    from knowledge_base import build_knowledge_base
-except ModuleNotFoundError:
-    from .knowledge_base import build_knowledge_base
+from knowledge_base import KnowledgeBase
 
 SUPPORTED_SUFFIXES = {".txt", ".md", ".markdown"}
 
 
-def load_directory(directory: str | Path) -> int:
-    root = Path(directory)
-    if not root.exists():
-        raise FileNotFoundError(f"Директория не найдена: {root}")
+def load_directory(knowledge_dir: Path, kb: KnowledgeBase | None = None) -> int:
+    kb = kb or KnowledgeBase()
+    if not knowledge_dir.exists():
+        knowledge_dir.mkdir(parents=True, exist_ok=True)
+        print(f"Создана пустая директория: {knowledge_dir}")
+        return 0
 
-    kb = build_knowledge_base()
     total_chunks = 0
-    for path in sorted(root.rglob("*")):
-        if not path.is_file() or path.suffix.lower() not in SUPPORTED_SUFFIXES:
-            continue
-        content = path.read_text(encoding="utf-8")
-        title = str(path.relative_to(root))
-        total_chunks += kb.add_document(content, title)
-        print(f"Загружен: {title} ({len(content)} символов)")
+    files = [p for p in knowledge_dir.rglob("*") if p.suffix.lower() in SUPPORTED_SUFFIXES]
+    if not files:
+        print(f"Нет файлов {SUPPORTED_SUFFIXES} в {knowledge_dir}")
+        return 0
+
+    for file_path in sorted(files):
+        chunks = kb.add_file(file_path)
+        total_chunks += chunks
+        print(f"  + {file_path.name}: {chunks} чанк(ов)")
+
+    print(f"Готово: {len(files)} файл(ов), {total_chunks} чанк(ов).")
     return total_chunks
 
 
@@ -33,12 +35,11 @@ def main() -> None:
     parser = argparse.ArgumentParser(description="Инициализация базы знаний из папки")
     parser.add_argument(
         "--dir",
-        default="./docs",
-        help="Путь к директории с .txt/.md файлами",
+        default="./knowledge",
+        help="Путь к директории с документами (.txt, .md)",
     )
     args = parser.parse_args()
-    chunks = load_directory(args.dir)
-    print(f"Готово. Всего чанков: {chunks}")
+    load_directory(Path(args.dir))
 
 
 if __name__ == "__main__":