File size: 6,533 Bytes
c75184e 8a1fc38 c75184e 8a1fc38 c75184e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 | """
Vector Store Service — ChromaDB/FAISS vector database operations.
Extracted from main.py L276-L397. Manages embedding model loading,
vector store creation/loading, similarity search, and the full
knowledge base build pipeline.
"""
import logging
from pathlib import Path
from typing import Optional, List
from backend.src.config import settings
logger = logging.getLogger("math_assistant.vector")
_EMBEDDINGS_CACHE = {}
def get_embeddings():
"""Get or create the singleton embedding model instance."""
if "model" not in _EMBEDDINGS_CACHE:
from langchain_huggingface import HuggingFaceEmbeddings
logger.info(f"Loading embedding model: {settings.embedding_model}")
_EMBEDDINGS_CACHE["model"] = HuggingFaceEmbeddings(
model_name=settings.embedding_model,
model_kwargs={"device": "cpu"},
encode_kwargs={"normalize_embeddings": True},
)
return _EMBEDDINGS_CACHE["model"]
class MathVectorStore:
"""Vector store abstraction supporting ChromaDB (primary) and FAISS (fallback)."""
def __init__(self):
self.embeddings = get_embeddings()
self.vectorstore = None
self.db_type = settings.vector_db_type
self._load_existing()
def _load_existing(self):
if self.db_type == "chroma":
self._try_chroma()
else:
self._try_faiss()
def _try_chroma(self, documents=None):
try:
from langchain_community.vectorstores import Chroma
persist_path = Path(settings.chroma_persist_dir)
persist_path.mkdir(parents=True, exist_ok=True)
if documents:
self.vectorstore = Chroma.from_documents(
documents=documents, embedding=self.embeddings,
collection_name=settings.collection_name,
persist_directory=str(persist_path))
logger.info("ChromaDB created.")
elif list(persist_path.glob("*.sqlite3")):
self.vectorstore = Chroma(
collection_name=settings.collection_name,
embedding_function=self.embeddings,
persist_directory=str(persist_path))
logger.info(f"ChromaDB loaded ({self.vectorstore._collection.count()} docs)")
except Exception as e:
logger.warning(f"ChromaDB failed ({type(e).__name__}: {e}), switching to FAISS")
self.db_type = "faiss"
if documents:
self._try_faiss(documents)
def _try_faiss(self, documents=None):
try:
from langchain_community.vectorstores import FAISS
index_path = Path(settings.faiss_index_path)
if documents:
self.vectorstore = FAISS.from_documents(documents, self.embeddings)
index_path.mkdir(parents=True, exist_ok=True)
self.vectorstore.save_local(str(index_path))
logger.info(f"FAISS saved to {index_path}")
elif index_path.exists() and any(index_path.iterdir()):
self.vectorstore = FAISS.load_local(
str(index_path), self.embeddings,
allow_dangerous_deserialization=True)
logger.info("FAISS loaded.")
except Exception as e:
logger.error(f"FAISS failed: {e}")
def build_knowledge_base(self, documents):
"""Build the vector index from a list of Document chunks."""
logger.info(f"Building knowledge base with {len(documents)} chunks...")
if self.db_type == "chroma":
self._try_chroma(documents)
else:
self._try_faiss(documents)
logger.info("Knowledge base ready.")
def add_documents(self, documents):
"""Add documents to an existing vector store, or create one."""
if self.vectorstore is None:
self.build_knowledge_base(documents)
else:
self.vectorstore.add_documents(documents)
def similarity_search(self, query: str, k: int = None, filter_topic: str = None):
"""Search for similar documents by query string."""
k = k or settings.top_k_results
if self.vectorstore is None:
return []
try:
if filter_topic and self.db_type == "chroma":
return self.vectorstore.similarity_search(query, k=k, filter={"topic": filter_topic})
return self.vectorstore.similarity_search(query, k=k)
except Exception as e:
logger.error(f"Search failed: {e}")
return []
def as_retriever(self, k: int = None):
"""Return a LangChain retriever interface."""
k = k or settings.top_k_results
return self.vectorstore.as_retriever(search_kwargs={"k": k}) if self.vectorstore else None
def get_document_count(self) -> int:
"""Get the number of indexed documents."""
if self.vectorstore is None:
return 0
try:
return (self.vectorstore._collection.count() if self.db_type == "chroma"
else self.vectorstore.index.ntotal)
except Exception as e:
logger.error(f"get_document_count failed: {e}")
return 0
def is_ready(self) -> bool:
"""Check if the vector store is loaded and has documents."""
return self.vectorstore is not None and self.get_document_count() > 0
_PIPELINE_CACHE = {}
def build_pipeline(force_rebuild=False):
"""Build or load the knowledge base pipeline."""
from backend.src.config import settings
if "store" in _PIPELINE_CACHE and not force_rebuild:
return _PIPELINE_CACHE["store"]
if getattr(settings, "vector_db", "qdrant") == "qdrant":
from backend.src.services.qdrant_service import QdrantService
store = QdrantService()
if store.get_document_count() == 0 or force_rebuild:
from backend.src.math.knowledge_indexer import KnowledgeIndexer
indexer = KnowledgeIndexer()
indexer.index_all()
_PIPELINE_CACHE["store"] = store
return store
else:
# Fallback to Chroma (legacy)
store = MathVectorStore()
if store.is_ready() and not force_rebuild:
logger.info(f"Knowledge base already built ({store.get_document_count()} docs).")
else:
logger.warning("Chroma build from scratch is deprecated. Please use Qdrant.")
_PIPELINE_CACHE["store"] = store
return store
|