diff --git a/init_knowledge.py b/init_knowledge.py new file mode 100644 index 0000000..58663f3 --- /dev/null +++ b/init_knowledge.py @@ -0,0 +1,46 @@ +"""Загрузка документов из директории в векторное хранилище.""" +from __future__ import annotations + +import argparse +from pathlib import Path + +from knowledge_base import KnowledgeBase + +SUPPORTED_SUFFIXES = {".txt", ".md", ".markdown"} + + +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 + + total_chunks = 0 + 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 + + +def main() -> None: + parser = argparse.ArgumentParser(description="Инициализация базы знаний из папки") + parser.add_argument( + "--dir", + default="./knowledge", + help="Путь к директории с документами (.txt, .md)", + ) + args = parser.parse_args() + load_directory(Path(args.dir)) + + +if __name__ == "__main__": + main()