72 lines
2.3 KiB
Python
72 lines
2.3 KiB
Python
"""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)
|