diff --git a/load_documents.py b/load_documents.py new file mode 100644 index 0000000..f5ba73a --- /dev/null +++ b/load_documents.py @@ -0,0 +1,66 @@ +"""Utility script to load documents from a directory into the Qdrant vector store. + +The script walks through the specified directory, reads all .txt files, splits them into chunks using +`RecursiveCharacterTextSplitter`, embeds the chunks with `OllamaEmbeddings`, and stores them in the +local Qdrant collection via the helper functions defined in :mod:`main`. + +Usage: + python load_documents.py /path/to/docs + +The script prints the number of documents added. +""" + +import os +import sys +from pathlib import Path +from uuid import uuid4 + +from main import get_vector_store, get_embeddings, chunk_document + + +def load_documents_from_dir(directory: str) -> int: + """Load all .txt files from *directory* into the vector store. + + Parameters + ---------- + directory: str + Path to the directory containing text files. + + Returns + ------- + int + Number of documents (chunks) added to the store. + """ + directory_path = Path(directory) + if not directory_path.is_dir(): + raise ValueError(f"{directory!r} is not a directory") + + vector_store = get_vector_store() + embedding = get_embeddings() + + total_chunks = 0 + for file_path in directory_path.rglob("*.txt"): + try: + content = file_path.read_text(encoding="utf-8") + except Exception as exc: + print(f"Failed to read {file_path}: {exc}", file=sys.stderr) + continue + title = file_path.stem + documents = chunk_document(content, title) + ids = [str(uuid4()) for _ in documents] + vector_store.add_documents(documents, ids=ids) + total_chunks += len(documents) + return total_chunks + + +if __name__ == "__main__": + if len(sys.argv) != 2: + print("Usage: python load_documents.py ") + sys.exit(1) + dir_path = sys.argv[1] + try: + added = load_documents_from_dir(dir_path) + print(f"Added {added} chunk(s) to the knowledge base.") + except Exception as exc: + print(f"Error: {exc}", file=sys.stderr) + sys.exit(1)