add vectorstore.py

This commit is contained in:
2026-06-04 16:27:19 +00:00
parent c8d6e36c80
commit 5e29ed8e20
+71
View File
@@ -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)