26 lines
890 B
Python
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))
|