feat: Chroma + OllamaEmbeddings(nomic-embed-text)

vectorstore.py

@@ -1,54 +1,33 @@
-import os
+"""ChromaDB + OllamaEmbeddings (nomic-embed-text)."""
 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
+from langchain_core.documents import Document
 
 
 def create_vectorstore(persist_directory: str = "./chroma_db") -> Chroma:
-    """Create or load a Chroma vector store backed by Ollama embeddings.
+    """Создать/открыть ChromaDB с Ollama-эмбеддингами."""
-
-    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,
+        collection_name="rag_kb",
         embedding_function=embeddings,
+        persist_directory=persist_directory,
     )
 
 
-def load_documents(directory: str, vectorstore: Chroma, chunk_size: int = 1000, chunk_overlap: int = 200) -> None:
+def load_documents(directory: str, vectorstore: Chroma) -> int:
-    """Load all .txt and .md files from *directory* into *vectorstore*.
+    """Загрузить .txt/.md файлы из директории в ChromaDB с чанкингом."""
-
+    splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
-    The function reads files, splits them into chunks using
+    docs = []
-    ``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"}:
+        if path.suffix.lower() not in (".txt", ".md"):
             continue
-        text = path.read_text(encoding="utf-8")
+        text = path.read_text(encoding="utf-8", errors="replace")
-        docs.extend(splitter.split_text(text))
+        chunks = splitter.split_text(text)
-    # Convert list of strings to Document objects
+        docs.extend(
-    documents = [Document(page_content=chunk, metadata={"source": str(p)}) for chunk in docs]
+            [Document(page_content=c, metadata={"source": str(path)}) for c in chunks]
-    vectorstore.add_documents(documents)
+        )
-
+    if docs:
-# If this module is executed directly, load the default documents folder.
+        vectorstore.add_documents(docs)
-if __name__ == "__main__":
+    return len(docs)
-    vs = create_vectorstore()
-    load_documents("documents", vs)
-    print("Vector store populated.")