diff --git a/src/loader.py b/src/loader.py index 28efef7..951c31f 100644 --- a/src/loader.py +++ b/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: - """Load all .txt files from *directory* into the knowledge base. +def load_directory(dir_path: pathlib.Path, kb: KnowledgeBase) -> None: + """Load each text file into the knowledge base. - Parameters - ---------- - directory: str - Path to the directory containing text files. + The file name (without extension) is used as the document title. """ - 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 text files into the knowledge base") - parser.add_argument("directory", type=str, help="Directory with .txt files") + parser = argparse.ArgumentParser(description="Load documents into the RAG knowledge base") + parser.add_argument("directory", type=str, help="Path to directory containing .txt files") args = parser.parse_args() - load_directory(args.directory) - print("Loading completed.") \ No newline at end of file + kb = KnowledgeBase() + load_directory(pathlib.Path(args.directory), kb) + print("Loading complete.") \ No newline at end of file