""" Vector store management for document retrieval. """ import chromadb from chromadb.config import Settings from typing import List, Dict, Optional, Any import uuid from pathlib import Path import json class VectorStore: """ ChromaDB-based vector store for medical knowledge. """ def __init__( self, collection_name: str = "medical_knowledge", persist_directory: str = "data/knowledge_base", embedding_function = None ): self.persist_directory = Path(persist_directory) self.persist_directory.mkdir(parents=True, exist_ok=True) # Version-agnostic chromadb client initialization try: # Try new API (chromadb >= 1.0) self.client = chromadb.PersistentClient( path=str(self.persist_directory) ) except (AttributeError, TypeError): try: # Try 0.4+ API with Settings self.client = chromadb.PersistentClient( path=str(self.persist_directory), settings=Settings( anonymized_telemetry=False, allow_reset=True ) ) except AttributeError: # Fall back to old API (chromadb < 0.4) self.client = chromadb.Client(Settings( chroma_db_impl="duckdb+parquet", persist_directory=str(self.persist_directory), anonymized_telemetry=False )) self.collection_name = collection_name # FIXED: Force use of SentenceTransformer to match build script # Default ChromaDB uses ONNX which might be incompatible with our build if embedding_function is None: from chromadb.utils import embedding_functions self.embedding_function = embedding_functions.SentenceTransformerEmbeddingFunction( model_name="all-MiniLM-L6-v2" ) else: self.embedding_function = embedding_function # Get or create collection w/ explicit embedding function self.collection = self.client.get_or_create_collection( name=collection_name, embedding_function=self.embedding_function, metadata={"hnsw:space": "cosine"} ) print(f"✅ Vector store initialized: {collection_name}") print(f"📊 Documents in collection: {self.collection.count()}") def add_documents( self, documents: List[str], embeddings: List[List[float]], metadatas: Optional[List[Dict[str, Any]]] = None, ids: Optional[List[str]] = None ) -> List[str]: """Add documents with embeddings to the store.""" if ids is None: ids = [str(uuid.uuid4()) for _ in documents] if metadatas is None: metadatas = [{}] * len(documents) # Ensure metadata values are valid types (str, int, float, bool) clean_metadatas = [] for meta in metadatas: clean_meta = {} for k, v in meta.items(): if isinstance(v, (str, int, float, bool)): clean_meta[k] = v elif v is None: clean_meta[k] = "" else: clean_meta[k] = str(v) clean_metadatas.append(clean_meta) self.collection.add( ids=ids, embeddings=embeddings, documents=documents, metadatas=clean_metadatas ) return ids def search( self, query_embedding: List[float], n_results: int = 10, filter_metadata: Optional[Dict] = None ) -> Dict: """Search for similar documents.""" results = self.collection.query( query_embeddings=[query_embedding], n_results=n_results, where=filter_metadata, include=["documents", "metadatas", "distances"] ) return results def mmr_search( self, query_embedding: List[float], k: int = 10, fetch_k: int = 20, lambda_mult: float = 0.5, filter_metadata: Optional[Dict] = None ) -> Dict: """ Maximal Marginal Relevance search for diverse results. MMR balances relevance to query with diversity among results. Args: query_embedding: Query vector k: Number of documents to return fetch_k: Number of documents to fetch before MMR lambda_mult: Balance between relevance and diversity (0 = max diversity, 1 = max relevance) filter_metadata: Optional metadata filter Returns: Dict with documents, metadatas, and distances """ import numpy as np # Fetch more documents than needed results = self.collection.query( query_embeddings=[query_embedding], n_results=fetch_k, where=filter_metadata, include=["documents", "metadatas", "distances", "embeddings"] ) if not results["documents"] or not results["documents"][0]: return {"documents": [[]], "metadatas": [[]], "distances": [[]]} documents = results["documents"][0] metadatas = results["metadatas"][0] distances = results["distances"][0] embeddings = results.get("embeddings", [[]])[0] if not embeddings or len(embeddings) == 0: # Fallback to regular search if embeddings not available return self.search(query_embedding, n_results=k, filter_metadata=filter_metadata) # Convert query to numpy query_vec = np.array(query_embedding) doc_embeddings = np.array(embeddings) # Calculate similarity to query (convert distance to similarity) query_similarities = 1 - np.array(distances) # MMR selection selected_indices = [] remaining_indices = list(range(len(documents))) while len(selected_indices) < k and remaining_indices: mmr_scores = [] for idx in remaining_indices: # Relevance to query relevance = query_similarities[idx] # Max similarity to already selected docs if selected_indices: selected_embeddings = doc_embeddings[selected_indices] similarities = np.dot(selected_embeddings, doc_embeddings[idx]) max_sim_to_selected = np.max(similarities) else: max_sim_to_selected = 0 # MMR formula: λ * relevance - (1 - λ) * max_similarity mmr_score = (lambda_mult * relevance - (1 - lambda_mult) * max_sim_to_selected) mmr_scores.append((idx, mmr_score)) # Select highest MMR score best_idx = max(mmr_scores, key=lambda x: x[1])[0] selected_indices.append(best_idx) remaining_indices.remove(best_idx) # Build result return { "documents": [[documents[i] for i in selected_indices]], "metadatas": [[metadatas[i] for i in selected_indices]], "distances": [[distances[i] for i in selected_indices]] } def get_stats(self) -> Dict: """Get statistics about the collection.""" return { "name": self.collection_name, "count": self.collection.count() } self.client.delete_collection(self.collection_name) print(f"🗑️ Deleted collection: {self.collection_name}") class QdrantVectorStore: def __init__(self, url, api_key, collection_name="medical_knowledge"): from qdrant_client import QdrantClient self.client = QdrantClient(url=url, api_key=api_key) self.collection_name = collection_name print(f"✅ Context: Connected to Qdrant Cloud: {collection_name}") def search(self, query_embedding, n_results=5, filter_metadata=None): # Qdrant expects query_vector response = self.client.query_points( collection_name=self.collection_name, query=query_embedding, limit=n_results ) # Convert to standard format docs = [] metadatas = [] distances = [] for res in response.points: docs.append(res.payload.get("page_content", "")) metadatas.append({k:v for k,v in res.payload.items() if k != "page_content"}) # Chroma returns distance (lower is better), Qdrant returns score (higher is better). # Retriever expects distance and calculates 1-distance. # So we return 1-score matches Chroma distance (approx). distances.append(1 - res.score) return { "documents": [docs], "metadatas": [metadatas], "distances": [distances] } def get_stats(self): try: count = self.client.count(self.collection_name).count return {"name": self.collection_name, "count": count} except: return {"name": self.collection_name, "count": 0} def get_vector_store(): # Factory import os if os.getenv("VECTOR_DB_TYPE") == "qdrant": return QdrantVectorStore( url=os.getenv("QDRANT_URL"), api_key=os.getenv("QDRANT_API_KEY") ) return VectorStore()