utils/vector_db.py

import logging
import os
import threading
import uuid
import warnings
from typing import Dict, List, Optional

import chromadb
from chromadb.config import Settings
from sentence_transformers import CrossEncoder, SentenceTransformer
import torch
import torch.nn.functional as F
from transformers import AutoTokenizer, AutoModel

import config

# Suppress unnecessary warnings
warnings.filterwarnings("ignore", category=FutureWarning)
logging.getLogger("chromadb").setLevel(logging.ERROR)
logging.getLogger("sentence_transformers").setLevel(logging.WARNING)


class VectorDatabase:
    """Manage vector database for document embeddings using ChromaDB."""

    _embedding_model = None
    _embedding_model_name = None
    _embedding_model_lock = threading.Lock()
    _tokenizer = None

    _reranker_model = None
    _reranker_model_name = None
    _reranker_lock = threading.Lock()

    @staticmethod
    def _empty_query_result() -> Dict:
        return {
            "documents": [[]],
            "metadatas": [[]],
            "distances": [[]],
            "ids": [[]],
        }

    @staticmethod
    def _is_dimension_mismatch_error(err: Exception) -> bool:
        msg = str(err).lower()
        return (
            ("expecting embedding with dimension" in msg and "got" in msg)
            or ("does not match index dimensionality" in msg)
            or ("dimensionality of" in msg and "index dimensionality" in msg)
        )

    @staticmethod
    def _is_index_not_found_error(err: Exception) -> bool:
        msg = str(err).lower()
        return "index not found" in msg or "create an instance before querying" in msg

    def _recreate_collection(self) -> None:
        """Recreate collection to recover from stale/missing index internals."""
        collection_name = self.collection.name
        try:
            self.client.delete_collection(name=collection_name)
        except Exception:
            pass

        self.collection = self.client.get_or_create_collection(
            name=collection_name,
            metadata={"hnsw:space": "cosine"},
        )

    @staticmethod
    def _resolve_torch_device(env_var_name: str) -> str:
        """Resolve target device with optional env override."""
        import torch

        preference = os.getenv(env_var_name, "auto").strip().lower()
        if preference == "cpu":
            return "cpu"
        if preference == "cuda":
            return "cuda" if torch.cuda.is_available() else "cpu"
        return "cuda" if torch.cuda.is_available() else "cpu"

    def __init__(self, collection_name: str = "documents", persist_directory: str = None):
        if persist_directory is None:
            persist_directory = str(config.VECTOR_DB_DIR)

        self.client = self._create_chroma_client(persist_directory)
        self.collection = self.client.get_or_create_collection(
            name=collection_name,
            metadata={"hnsw:space": "cosine"},
        )

        self.embedding_model = self._get_or_create_embedding_model()

    def _create_chroma_client(self, persist_directory: str):
        """Create a Chroma client compatible with both legacy and modern APIs."""
        if hasattr(chromadb, "PersistentClient"):
            try:
                return chromadb.PersistentClient(
                    path=persist_directory,
                    settings=Settings(anonymized_telemetry=False),
                )
            except TypeError:
                return chromadb.PersistentClient(path=persist_directory)
            except Exception:
                pass

        return chromadb.Client(
            Settings(
                chroma_db_impl="duckdb+parquet",
                persist_directory=persist_directory,
                anonymized_telemetry=False,
            )
        )

    def ensure_collection_embedding_compatibility(self) -> bool:
        """
        Ensure persisted collection dimensionality matches current embedding model.

        Returns:
            True if collection was reset due to mismatch, else False.
        """
        try:
            count = self.collection.count()
        except Exception:
            return False

        try:
            if count == 0:
                # Query probing on empty indexes may not trigger dimension validation.
                probe_id = f"_dim_probe_{uuid.uuid4().hex[:12]}"
                probe_text = "embedding-dimension-probe"
                probe_embedding = self._encode([probe_text])
                self.collection.add(
                    embeddings=probe_embedding,
                    documents=[probe_text],
                    metadatas=[{"_probe": True}],
                    ids=[probe_id],
                )
                self.collection.delete(ids=[probe_id])
                return False

            probe = self._encode(["embedding-dimension-probe"])
            self.collection.query(query_embeddings=probe, n_results=1)
            return False
        except Exception as e:
            if not (self._is_dimension_mismatch_error(e) or self._is_index_not_found_error(e)):
                raise

            collection_name = self.collection.name
            reason = "missing_index" if self._is_index_not_found_error(e) else "dimension_mismatch"
            print(
                "[VECTOR_DB] collection_incompatible "
                f"reason={reason} collection={collection_name} model={config.EMBEDDING_MODEL}; resetting index"
            )
            self._recreate_collection()
            return True

    @classmethod
    def _get_or_create_embedding_model(cls):
        """Create embedding model once and reuse it for all vector DB instances."""
        with cls._embedding_model_lock:
            if cls._embedding_model is None or cls._embedding_model_name != config.EMBEDDING_MODEL:
                device = cls._resolve_torch_device("EMBEDDING_DEVICE")
                model_name = config.EMBEDDING_MODEL

                print(f"Loading embedding model ({model_name}) on {device}...")
                cls._tokenizer = None
                cls._embedding_model = SentenceTransformer(
                    model_name,
                    device=device,
                    model_kwargs={"torch_dtype": torch.float16},
                    trust_remote_code=True
                )
                cls._embedding_model_name = model_name
                print(f"Embedding model {model_name} loaded on {device}.")

            return cls._embedding_model

    @classmethod
    def _get_or_create_reranker(cls):
        """Create reranker model once and reuse it for all vector DB instances."""
        target_model = getattr(config, "RERANKER_MODEL", "cross-encoder/ms-marco-MiniLM-L-12-v2")

        with cls._reranker_lock:
            if cls._reranker_model is None or cls._reranker_model_name != target_model:
                device = cls._resolve_torch_device("RERANKER_DEVICE")
                print(f"Loading reranker model ({target_model}) on {device}...")
                cls._reranker_model = CrossEncoder(target_model, device=device, trust_remote_code=True, model_kwargs={"torch_dtype": torch.float16})
                
                # Qwen3-Reranker lacks a default pad token, which breaks batch sizes > 1
                if cls._reranker_model.tokenizer.pad_token is None:
                    cls._reranker_model.tokenizer.pad_token = cls._reranker_model.tokenizer.eos_token
                    
                cls._reranker_model_name = target_model
                print(f"Reranker model {target_model} loaded on {device}.")

            return cls._reranker_model

    @classmethod
    def clear_runtime_caches(cls, unload_embedding_model: bool = False):
        """Best-effort cleanup for GPU/CPU caches between large space workloads."""
        try:
            import gc

            gc.collect()
        except Exception:
            pass

        try:
            import torch

            if unload_embedding_model and cls._embedding_model is not None:
                try:
                    cls._embedding_model.to("cpu")
                except Exception:
                    pass

                try:
                    if cls._reranker_model is not None and hasattr(cls._reranker_model, "model"):
                        cls._reranker_model.model.to("cpu")
                except Exception:
                    pass

                cls._embedding_model = None
                cls._embedding_model_name = None
                cls._reranker_model = None
                cls._reranker_model_name = None

            if torch.cuda.is_available():
                torch.cuda.empty_cache()
                if hasattr(torch.cuda, "ipc_collect"):
                    torch.cuda.ipc_collect()
        except Exception:
            pass

    def _embedding_batch_size(self) -> int:
        """Embedding batch size tuned for low-VRAM GPUs with env override."""
        if getattr(config, "EMBEDDING_BATCH_SIZE", 0) > 0:
            return int(config.EMBEDDING_BATCH_SIZE)

        model_name = (config.EMBEDDING_MODEL or "").lower()
        if "nomic-embed-text" in model_name:
            return 32
        if "bge-m3" in model_name:
            return 16
        if "bge-large" in model_name:
            return 24
        return 64

    def _reranker_batch_size(self) -> int:
        if getattr(config, "RERANKER_BATCH_SIZE", 0) > 0:
            return int(config.RERANKER_BATCH_SIZE)
        return 16

    def _encode(self, texts: list, instruction: str = None) -> list:
        """Encode texts using SentenceTransformer."""
        if instruction:
            texts = [f"Instruct: {instruction}\nQuery: {t}" for t in texts]

        batch_size = self._embedding_batch_size()
        
        embeddings = self.embedding_model.encode(
            texts,
            batch_size=batch_size,
            show_progress_bar=False,
            convert_to_tensor=False,
            normalize_embeddings=True
        )
        return embeddings.tolist()

    def add_documents(self, texts: List[str], metadatas: List[Dict], ids: List[str]):
        """Embed documents — no instruction prefix for docs."""
        if not texts:
            return
        
        # Documents are encoded WITHOUT instruction
        embeddings = self._encode(texts)

        try:
            self.collection.add(
                embeddings=embeddings,
                documents=texts,
                metadatas=metadatas,
                ids=ids,
            )
        except Exception as e:
            if not (self._is_dimension_mismatch_error(e) or self._is_index_not_found_error(e)):
                raise

            # Auto-recover from stale dimensionality or missing collection index internals.
            collection_name = self.collection.name
            reason = "missing_index" if self._is_index_not_found_error(e) else "dimension_mismatch"
            print(
                "[VECTOR_DB] add_recover "
                f"reason={reason} collection={collection_name} model={config.EMBEDDING_MODEL}; resetting index and retrying"
            )
            self._recreate_collection()
            self.collection.add(
                embeddings=embeddings,
                documents=texts,
                metadatas=metadatas,
                ids=ids,
            )

        if hasattr(self.client, "persist"):
            self.client.persist()

    def query(self, query_text: str, n_results: int = 20, filter_dict: Optional[Dict] = None) -> Dict:
        """Embed query WITH instruction for better retrieval."""
        try:
            if self.collection.count() == 0:
                return self._empty_query_result()
        except Exception:
            pass

        instruction = (
            "Given a student's question, retrieve relevant passages "
            "from academic textbooks that answer the question"
        )
        query_embedding = self._encode([query_text], instruction=instruction)

        try:
            return self.collection.query(
                query_embeddings=query_embedding,
                n_results=n_results,
                where=filter_dict,
            )
        except Exception as e:
            if "index not found" in str(e).lower() or "create an instance" in str(e).lower():
                return self._empty_query_result()

            if self._is_dimension_mismatch_error(e):
                print(
                    "[VECTOR_DB] query_dimension_mismatch "
                    f"model={config.EMBEDDING_MODEL}; returning empty results"
                )
                return self._empty_query_result()
            raise

    def query_and_rerank(
        self,
        query_text: str,
        n_retrieve: int = 20,
        n_final: int = 6,
        filter_dict: Optional[Dict] = None,
    ) -> Dict:
        """Two-stage retrieval: dense retrieval in Chroma followed by reranking."""
        raw = self.query(query_text=query_text, n_results=n_retrieve, filter_dict=filter_dict)

        docs = raw.get("documents", [[]])[0]
        metadatas = raw.get("metadatas", [[]])[0]
        distances = raw.get("distances", [[]])[0]
        ids = raw.get("ids", [[]])[0]

        if not docs:
            return raw

        reranker = self._get_or_create_reranker()
        pairs = [[query_text, doc] for doc in docs]
        scores = reranker.predict(
            pairs,
            batch_size=self._reranker_batch_size(),
            show_progress_bar=False,
        )
        
        # Qwen3 outputs raw logits (e.g. -11 to +8). Convert them to 0-1 probabilities for the UI.
        scores = torch.sigmoid(torch.tensor(list(scores))).tolist()
        scores = [float(s) for s in scores]

        ranked = sorted(
            zip(scores, docs, metadatas, distances, ids),
            key=lambda x: x[0],
            reverse=True,
        )[: max(1, n_final)]

        scores_out, docs_out, metas_out, dists_out, ids_out = zip(*ranked)
        return {
            "documents": [list(docs_out)],
            "metadatas": [list(metas_out)],
            "distances": [list(dists_out)],
            "ids": [list(ids_out)],
            "reranker_scores": [list(scores_out)],
        }

    def delete_collection(self):
        """Delete the entire collection."""
        self.client.delete_collection(name=self.collection.name)

    def get_collection_count(self) -> int:
        """Get the number of documents in the collection."""
        return self.collection.count()

    def get_all_documents(self) -> tuple[List[str], List[Dict]]:
        """Get all documents and metadata from the collection."""
        count = self.collection.count()
        if count == 0:
            return [], []

        results = self.collection.get()
        return results.get("documents", []), results.get("metadatas", [])

    def create_space_collection(self, space_name: str):
        """Create a new collection for a specific subject space."""
        return self.client.get_or_create_collection(
            name=f"space_{space_name}",
            metadata={"hnsw:space": "cosine"},
        )