Files
task-6a02e23da6fe2e4ac16acf65/init_documents.py
T

26 lines
890 B
Python

"""
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))