diff --git a/init_client.py b/init_client.py new file mode 100644 index 0000000..bdaf678 --- /dev/null +++ b/init_client.py @@ -0,0 +1,69 @@ +#!/usr/bin/env python3 +"""Initialization client - loads documents from a directory into the vector store.""" + +import os +from typing import List + +from langchain_core.documents import Document +from langchain_community.documenloaders import TextLoader + + from vector_store import get_vector_store, add_documents_to_store + + +def load_documents_from_dir(dir_path: str) -> List[Document]: + """Load all .txt files from a directory as LangChain Documents. + + Args: + dir_path: Path to directory containing .txt files. + + Returns: + List of Document objects. + """ + documents = [] + for filename in os.listdir(dir_path): + if filename.endswith(".txt"): + file_path = os.path.join(dir_path, filename) + loader = TextLoader(file_path, encoding="utf-8") + docs = loader.load() + # Add title metadata from filename + title = os.path.splitext(filename)[0] + for doc in docs: + doc.metadata["title"] = title + documents.extend(docs) + return documents + + +def init_knowledge_base(dir_path: str) -> int: + """Initialize the knowledge base by loading documents from a directory. + + Args: + dir_path: Path to directory containing .txt files. + + Returns: + Number of chunks added to the store. + """ + store = get_vector_store() + documents = load_documents_from_dir(dir_path) + + if not documents: + print(f"No .txt files found in {dir_path}") + return 0 + + ids = add_documents_to_store(store, documents) + print(f"Loaded {len(documents)} document(s) into knowledge base ({len(ids)} chunks).") + return len(ids) + + +if __name__ == "__main__": + import sys + + if len(sys.argv) < 2: + print("Usage: python init_client.py ") + sys.exit(1) + + dir_path = sys.argv[1] + if not os.path.isdir(dir_path): + print(f"Error: {dir_path} is not a directory") + sys.exit(1) + + init_knowledge_base(dir_path)