src/embeddings/vector_store.py

"""
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()