add vectorstore.py

vectorstore.py

@@ -1,26 +1,83 @@
+"""
+Vector store utilities for ChromaDB with Ollama embeddings.
+
+Functions:
+- create_vectorstore(persist_directory="./chroma_db") -> chromadb.Chroma
+- load_documents(directory: str, vectorstore) -> None
+"""
 import os
 from pathlib import Path
 from typing import List
 
 from langchain_ollama import OllamaEmbeddings
-from langchain_chroma import Chroma
 from langchain_text_splitters import RecursiveCharacterTextSplitter
+from langchain_chroma import Chroma
 from langchain.docstore.document import Document
 
-def create_vectorstore(persist_directory: str = "./chroma_db"):
-    if not os.path.exists(persist_directory):
-        os.makedirs(persist_directory, exist_ok=True)
+
+def create_vectorstore(persist_directory: str = "./chroma_db") -> Chroma:
+    """Create or load a Chroma vector store.
+
+    Parameters
+    ----------
+    persist_directory : str, optional
+        Directory where the Chroma database is stored. Defaults to ``./chroma_db``.
+
+    Returns
+    -------
+    chromadb.Chroma
+        A Chroma instance backed by Ollama embeddings.
+    """
+    # Ensure directory exists
+    Path(persist_directory).mkdir(parents=True, exist_ok=True)
     embeddings = OllamaEmbeddings(model="nomic-embed-text")
     return Chroma(persist_directory=persist_directory, embedding_function=embeddings)
 
-def load_documents(directory: str, vectorstore) -> None:
-    splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
+
+def _load_text_files(directory: str) -> List[Document]:
+    """Load all .txt and .md files from *directory* into Documents.
+
+    Parameters
+    ----------
+    directory : str
+        Path to the folder containing documents.
+
+    Returns
+    -------
+    list[langchain.docstore.document.Document]
+        List of Document objects with ``page_content`` set to file text and
+        ``metadata`` containing the source path.
+    """
     docs: List[Document] = []
-    for path in Path(directory).rglob("*.txt"):
-        text = path.read_text(encoding="utf-8")
-        docs.extend(splitter.split_documents([Document(page_content=text)]))
-    for path in Path(directory).rglob("*.md"):
-        text = path.read_text(encoding="utf-8")
-        docs.extend(splitter.split_documents([Document(page_content=text)]))
-    if docs:
-        vectorstore.add_documents(docs)
+    for root, _, files in os.walk(directory):
+        for fname in files:
+            if not fname.lower().endswith(('.txt', '.md')):
+                continue
+            full_path = Path(root) / fname
+            text = full_path.read_text(encoding="utf-8")
+            docs.append(Document(page_content=text, metadata={"source": str(full_path)}))
+    return docs
+
+
+def load_documents(directory: str, vectorstore: Chroma) -> None:
+    """Chunk documents from *directory* and add them to *vectorstore*.
+
+    Parameters
+    ----------
+    directory : str
+        Folder with .txt/.md files.
+    vectorstore : chromadb.Chroma
+        The vector store instance returned by ``create_vectorstore``.
+    """
+    docs = _load_text_files(directory)
+    if not docs:
+        return
+    splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
+    chunks = splitter.split_documents(docs)
+    vectorstore.add_documents(chunks)
+
+# Example usage (uncomment for manual testing):
+# if __name__ == "__main__":
+#     vs = create_vectorstore()
+#     load_documents("documents", vs)
+"""