diff --git a/vectorstore.py b/vectorstore.py new file mode 100644 index 0000000..336ad3d --- /dev/null +++ b/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)