Solution ready for review: add load_documents.py

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