diff --git a/init_loader.py b/init_loader.py index 52d8233..a7b288f 100644 --- a/init_loader.py +++ b/init_loader.py @@ -1,10 +1,13 @@ -from pathlib import Path +import os from rag_tools import add_to_knowledge_base -def load_docs_from_dir(directory: str): - for file in Path(directory).glob("**/*.txt"): - text = file.read_text(encoding="utf-8") - add_to_knowledge_base(text, title=file.name) - if __name__ == "__main__": - load_docs_from_dir("docs") + dir_path = os.getenv("DOCS_DIR", ".") + for root, _, files in os.walk(dir_path): + for f in files: + if f.lower().endswith(('.txt', '.md')): + path = os.path.join(root, f) + with open(path, 'r', encoding='utf-8') as fp: + content = fp.read() + title = os.path.relpath(path, dir_path) + add_to_knowledge_base(content, title)