diff --git a/init_documents.py b/init_documents.py new file mode 100644 index 0000000..09d65aa --- /dev/null +++ b/init_documents.py @@ -0,0 +1,25 @@ +""" +Script to load all .txt files from a directory into the knowledge base. +""" + +import argparse +from pathlib import Path +from agent import agent +from langchain_core.messages import HumanMessage + +async def init_docs(dir_path: Path): + for txt_file in dir_path.rglob("*.txt"): + content = txt_file.read_text(encoding="utf-8") + title = txt_file.stem + await agent.ainvoke( + {"messages": [HumanMessage(content=f"Add document: {title}")]}, + {"configurable": {"thread_id": f"init-{txt_file.name}"}}, + ) + print("Initialization complete.") + +if __name__ == "__main__": # pragma: no cover + parser = argparse.ArgumentParser(description="Initialize knowledge base from directory") + parser.add_argument("dir", type=Path, help="Directory with text files") + args = parser.parse_args() + import asyncio + asyncio.run(init_docs(args.dir))