diff --git a/vectorstore.py b/vectorstore.py new file mode 100644 index 0000000..ecf80bb --- /dev/null +++ b/vectorstore.py @@ -0,0 +1,71 @@ +"""ChromaDB vector store с Ollama embeddings для FAQ курса.""" +import os +from pathlib import Path +from typing import List + +from langchain_chroma import Chroma +from langchain_ollama import OllamaEmbeddings +from langchain_core.documents import Document +from langchain.text_splitter import RecursiveCharacterTextSplitter + +PERSIST_DIR = os.getenv("CHROMA_PERSIST_DIR", "./chroma_faq") +COLLECTION_NAME = "faq_collection" +DATA_DIR = Path(__file__).parent / "data" + +embeddings = OllamaEmbeddings(model="nomic-embed-text") + +_store: Chroma | None = None + + +def get_vector_store() -> Chroma: + """Возвращает singleton‑экземпляр ChromaDB (персистентный на диск).""" + global _store + if _store is None: + _store = Chroma( + collection_name=COLLECTION_NAME, + embedding_function=embeddings, + persist_directory=PERSIST_DIR, + ) + return _store + + +def load_faq_to_chroma() -> int: + """Загружает все .md файлы из data/ в ChromaDB с чанкингом. + + Returns: + количество проиндексированных чанков + """ + if not DATA_DIR.exists(): + print(f"[vectorstore] Папка data/ не найдена: {DATA_DIR}") + return 0 + + splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50) + all_docs: List[Document] = [] + + for md_file in sorted(DATA_DIR.glob("*.md")): + text = md_file.read_text(encoding="utf-8") + chunks = splitter.create_documents( + [text], + metadatas=[{"source": md_file.name}], + ) + all_docs.extend(chunks) + print(f"[vectorstore] {md_file.name}: {len(chunks)} чанков") + + if all_docs: + get_vector_store().add_documents(all_docs) + + print(f"[vectorstore] Итого проиндексировано: {len(all_docs)} чанков") + return len(all_docs) + + +def similarity_search(query: str, k: int = 3) -> List[Document]: + """Семантический поиск в ChromaDB. + + Args: + query: поисковый запрос + k: количество результатов + + Returns: + список наиболее релевантных документов + """ + return get_vector_store().similarity_search(query, k=k)