Update src/loader.py

src/loader.py

@@ -1,29 +1,37 @@
-"""Utility to load all .txt files from a directory into the knowledge base.
+"""Utility to load all text files from a directory into the knowledge base.
 
-This module is useful for initializing the vector store before interacting with the agent.
+Usage:
+    python -m src.loader /path/to/docs
 """
 
-from pathlib import Path
+from __future__ import annotations
+
+import argparse
+import pathlib
+from typing import Iterable
 
 from .vector_store import KnowledgeBase
 
-kb = KnowledgeBase()
+def _iter_text_files(dir_path: pathlib.Path) -> Iterable[pathlib.Path]:
+    """Yield all .txt files in *dir_path* recursively."""
+    for p in dir_path.rglob("*.txt"):
+        yield p
 
-def load_directory(directory: str) -> None:
+def load_directory(dir_path: pathlib.Path, kb: KnowledgeBase) -> None:
-    """Load all .txt files from *directory* into the knowledge base.
+    """Load each text file into the knowledge base.
 
-    Parameters
+    The file name (without extension) is used as the document title.
-    ----------
-    directory: str
-        Path to the directory containing text files.
     """
-    kb.load_from_directory(directory)
+    for file_path in _iter_text_files(dir_path):
+        title = file_path.stem
+        content = file_path.read_text(encoding="utf-8")
+        kb.add_document(title=title, content=content)
+        print(f"Loaded {file_path} as '{title}'")
 
 if __name__ == "__main__":
-    import argparse
+    parser = argparse.ArgumentParser(description="Load documents into the RAG knowledge base")
-
+    parser.add_argument("directory", type=str, help="Path to directory containing .txt files")
-    parser = argparse.ArgumentParser(description="Load text files into the knowledge base")
-    parser.add_argument("directory", type=str, help="Directory with .txt files")
     args = parser.parse_args()
-    load_directory(args.directory)
+    kb = KnowledgeBase()
-    print("Loading completed.")
+    load_directory(pathlib.Path(args.directory), kb)
+    print("Loading complete.")