add vectorstore.py

vectorstore.py

@@ -0,0 +1,60 @@
+"""ChromaDB + Ollama embeddings: создание хранилища и загрузка документов."""
+from __future__ import annotations
+
+from pathlib import Path
+
+from langchain_chroma import Chroma
+from langchain_core.documents import Document
+from langchain_ollama import OllamaEmbeddings
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+
+DEFAULT_PERSIST_DIR = "./chroma_db"
+DEFAULT_EMBED_MODEL = "nomic-embed-text"
+CHUNK_SIZE = 800
+CHUNK_OVERLAP = 120
+
+
+def create_vectorstore(persist_directory: str = DEFAULT_PERSIST_DIR) -> Chroma:
+    embeddings = OllamaEmbeddings(model=DEFAULT_EMBED_MODEL)
+    return Chroma(
+        collection_name="local_kb",
+        embedding_function=embeddings,
+        persist_directory=persist_directory,
+    )
+
+
+def _read_text_file(path: Path) -> str:
+    return path.read_text(encoding="utf-8")
+
+
+def load_documents(directory: str, vectorstore: Chroma) -> int:
+    """Читает .txt/.md из directory, чанкует и добавляет в ChromaDB."""
+    root = Path(directory)
+    if not root.exists():
+        return 0
+
+    files = sorted(root.glob("*.txt")) + sorted(root.glob("*.md"))
+    if not files:
+        return 0
+
+    splitter = RecursiveCharacterTextSplitter(
+        chunk_size=CHUNK_SIZE,
+        chunk_overlap=CHUNK_OVERLAP,
+    )
+    docs: list[Document] = []
+    for fp in files:
+        text = _read_text_file(fp)
+        chunks = splitter.split_text(text)
+        for i, chunk in enumerate(chunks):
+            docs.append(
+                Document(
+                    page_content=chunk,
+                    metadata={"source": fp.name, "chunk": i},
+                )
+            )
+
+    if not docs:
+        return 0
+
+    vectorstore.add_documents(docs)
+    return len(docs)