Update vectorstore.py

vectorstore.py

@@ -1,53 +1,53 @@
-"""Module for creating and loading a Chroma vector store with Ollama embeddings."""
+"""
+Vector store utilities for ChromaDB.
+"""
 
 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:
+def create_vectorstore(persist_directory: str = "./chroma_db"):
-    """Create a Chroma vector store with Ollama embeddings.
+    """Create or load a Chroma vector store.
 
     Parameters
     ----------
     persist_directory: str
-        Directory where the vector store will be persisted.
+        Directory where the Chroma database is persisted.
 
     Returns
     -------
     Chroma
-        The created Chroma vector store.
+        The 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:
+def load_documents(directory: str, vectorstore: Chroma, chunk_size: int = 1000, chunk_overlap: int = 200):
-    """Load .txt and .md files from a directory, split them into chunks, and add to the vector store.
+    """Read all .txt and .md files from a directory, split them into chunks and add to the vectorstore.
 
     Parameters
     ----------
     directory: str
-        Path to the directory containing documents.
+        Path to the folder containing documents.
     vectorstore: Chroma
-        The vector store to add documents to.
+        The vector store to which documents will be added.
     chunk_size: int
-        Maximum size of each chunk.
+        Maximum number of characters per chunk.
     chunk_overlap: int
-        Number of characters to overlap between chunks.
+        Number of overlapping characters between consecutive chunks.
     """
     splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
-    documents: List[Document] = []
+    docs = []
-    for file_path in Path(directory).glob("**/*"):
+    for file_path in Path(directory).rglob("*.txt"):
-        if file_path.suffix.lower() in {".txt", ".md"}:
+        docs.append(file_path.read_text(encoding="utf-8"))
-            text = file_path.read_text(encoding="utf-8")
+    for file_path in Path(directory).rglob("*.md"):
-            chunks = splitter.split_text(text)
+        docs.append(file_path.read_text(encoding="utf-8"))
-            for chunk in chunks:
+    if not docs:
-                documents.append(Document(page_content=chunk, metadata={"source": str(file_path)}))
+        return
-    if documents:
+    texts = splitter.split_text("\n\n".join(docs))
-        vectorstore.add_documents(documents)
+    vectorstore.add_texts(texts)
-        vectorstore.persist()
+    vectorstore.persist()