diff --git a/load_docs.py b/load_docs.py index d690a62..2d631bf 100644 --- a/load_docs.py +++ b/load_docs.py @@ -1,19 +1,21 @@ import sys from pathlib import Path +from rag_agent import add_to_knowledge_base -from utils import splitter, vector_store -from langchain_core.documents import Document - -def load_documents(dir_path: str): - for file_path in Path(dir_path).rglob("*.txt"): - content = file_path.read_text(encoding="utf-8") - chunks = splitter.split_text(content) - docs = [Document(page_content=chunk, metadata={"title": file_path.name}) for chunk in chunks] - vector_store.add_documents(docs) - print(f"Added {len(docs)} chunks from {file_path}") +def load_docs_from_dir(directory: str): + for file_path in Path(directory).rglob("*"): + if file_path.is_file() and file_path.suffix.lower() in {".txt", ".md", ".pdf"}: + try: + content = file_path.read_text(encoding="utf-8", errors="ignore") + except Exception as e: + print(f"Failed to read {file_path}: {e}") + continue + title = file_path.stem + result = add_to_knowledge_base(content, title) + print(f"Loaded {file_path}: {result}") if __name__ == "__main__": - if len(sys.argv) != 2: + if len(sys.argv) < 2: print("Usage: python load_docs.py ") sys.exit(1) - load_documents(sys.argv[1]) \ No newline at end of file + load_docs_from_dir(sys.argv[1]) \ No newline at end of file