diff --git a/init_documents.py b/init_documents.py index 09d65aa..ccd1165 100644 --- a/init_documents.py +++ b/init_documents.py @@ -1,25 +1,20 @@ """ -Script to load all .txt files from a directory into the knowledge base. +Script to load all .txt files from a directory into the vector store. """ - -import argparse +import os from pathlib import Path -from agent import agent -from langchain_core.messages import HumanMessage +from chunker import load_and_split +from vector_store import vector_store -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)) +if __name__ == "__main__": + docs_dir = Path("docs") + if not docs_dir.exists(): + print("Docs directory not found.") + exit(1) + chunks = load_and_split(docs_dir) + # Each chunk is (text, metadata) + documents = [] + for idx, (chunk, meta) in enumerate(chunks): + documents.append({"content": chunk, "metadata": meta}) + vector_store.add_documents(documents) + print(f"Loaded {len(documents)} chunks into the vector store.")