add vectorstore.py

vectorstore.py

@@ -0,0 +1,54 @@
+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.docstore.document import Document
+
+
+def create_vectorstore(persist_directory: str = "./chroma_db") -> Chroma:
+    """Create or load a Chroma vector store backed by Ollama embeddings.
+
+    Parameters
+    ----------
+    persist_directory: str
+        Directory where Chroma will persist its data.
+
+    Returns
+    -------
+    Chroma
+        A Chroma vector store instance.
+    """
+    embeddings = OllamaEmbeddings(model="nomic-embed-text")
+    return Chroma(
+        persist_directory=persist_directory,
+        embedding_function=embeddings,
+    )
+
+
+def load_documents(directory: str, vectorstore: Chroma, chunk_size: int = 1000, chunk_overlap: int = 200) -> None:
+    """Load all .txt and .md files from *directory* into *vectorstore*.
+
+    The function reads files, splits them into chunks using
+    ``RecursiveCharacterTextSplitter`` and adds the resulting
+    :class:`~langchain.docstore.document.Document` objects to the
+    vector store.
+    """
+    splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
+    docs: List[Document] = []
+    for path in Path(directory).glob("**/*"):
+        if path.suffix.lower() not in {".txt", ".md"}:
+            continue
+        text = path.read_text(encoding="utf-8")
+        docs.extend(splitter.split_text(text))
+    # Convert list of strings to Document objects
+    documents = [Document(page_content=chunk, metadata={"source": str(p)}) for chunk in docs]
+    vectorstore.add_documents(documents)
+
+# If this module is executed directly, load the default documents folder.
+if __name__ == "__main__":
+    vs = create_vectorstore()
+    load_documents("documents", vs)
+    print("Vector store populated.")