backend/app/rag/retriever.py

"""
Hybrid Retriever — Combines vector search (ChromaDB) with keyword search (BM25).
Uses Reciprocal Rank Fusion to merge results from both retrieval methods.
"""

import os

import chromadb
import structlog
from rank_bm25 import BM25Okapi
from typing import Optional

from app.rag.schema_enricher import SchemaEnricher

logger = structlog.get_logger()


class NoopCollection:
    """Minimal Chroma collection stand-in for local startup without vector search."""

    def count(self) -> int:
        return 0

    def get(self) -> dict:
        return {"ids": []}

    def delete(self, ids: list[str]):
        return None

    def add(self, documents: list[str], metadatas: list[dict], ids: list[str]):
        return None

    def query(self, query_texts: list[str], n_results: int) -> dict:
        return {"documents": [[]]}


class HybridRetriever:
    """
    Production RAG retriever combining:
    1. ChromaDB vector similarity (semantic search)
    2. BM25 keyword matching (exact table/column name matching)
    3. Reciprocal Rank Fusion to merge results
    4. Optional cross-encoder reranking for precision
    """

    def __init__(self, db_pool, chroma_persist_dir: str = "./chroma_db"):
        self.db_pool = db_pool
        self.enricher = SchemaEnricher(db_pool)
        self.vector_enabled = os.getenv("DISABLE_VECTOR_RAG", "").lower() not in {"1", "true", "yes"}

        # Initialize ChromaDB unless local dev explicitly disables vector search.
        if self.vector_enabled:
            self.chroma_client = chromadb.PersistentClient(path=chroma_persist_dir)
            
            # Use Hugging Face Inference API embeddings if token is configured
            # to avoid loading heavy local PyTorch/ONNX models.
            embedding_fn = None
            hf_token = os.getenv("HUGGINGFACEHUB_API_TOKEN")
            if hf_token:
                try:
                    from chromadb.utils.embedding_functions import HuggingFaceEmbeddingFunction
                    embedding_fn = HuggingFaceEmbeddingFunction(
                        api_key=hf_token,
                        model_name="sentence-transformers/all-MiniLM-L6-v2"
                    )
                    logger.info("using_huggingface_embedding_function")
                except Exception as e:
                    logger.warning("hf_embedding_init_failed", error=str(e))

            self.collection = self.chroma_client.get_or_create_collection(
                name="schema_knowledge_v2",
                metadata={"hnsw:space": "cosine"},
                embedding_function=embedding_fn,
            )
        else:
            self.chroma_client = None
            self.collection = NoopCollection()
            logger.warning("vector_rag_disabled")

        # Document store for BM25
        self.documents: list[str] = []
        self.doc_ids: list[str] = []
        self.bm25: Optional[BM25Okapi] = None

        # Optional cross-encoder reranker — lazy-loaded on first use
        self._reranker = None
        self._reranker_loaded = False

        # Index on startup — file-locked to prevent SQLite race when
        # multiple Gunicorn workers start simultaneously.
        self._index_schema_safe(chroma_persist_dir)

    def _index_schema_safe(self, chroma_persist_dir: str):
        """Index schema with a file lock to prevent concurrent writes."""
        lock_path = os.path.join(chroma_persist_dir, ".index_lock")
        try:
            import filelock
            lock = filelock.FileLock(lock_path, timeout=30)
            with lock:
                self._index_schema()
        except ImportError:
            # filelock not installed — proceed without locking (single-worker is fine)
            logger.warning("filelock_not_installed", hint="pip install filelock for multi-worker safety")
            self._index_schema()
        except filelock.Timeout:
            logger.error("index_lock_timeout", lock_path=lock_path)
            self._index_schema()  # Proceed anyway

    def _index_schema(self):
        """Index all tables into both ChromaDB and BM25."""
        try:
            schema_file = os.path.join(self.chroma_client.persist_directory if self.chroma_client and hasattr(self.chroma_client, "persist_directory") else "./chroma_db", "enriched_schema.json")
            
            # Try to load from file cache first to avoid startup DB queries
            loaded_from_cache = False
            enriched_tables = []
            
            if os.path.exists(schema_file):
                try:
                    import json
                    with open(schema_file, "r", encoding="utf-8") as f:
                        enriched_tables = json.load(f)
                    loaded_from_cache = True
                    logger.info("loaded_schema_from_file_cache", path=schema_file)
                except Exception as e:
                    logger.warning("failed_to_load_schema_cache", error=str(e))
            
            if not loaded_from_cache:
                enriched_tables = self.enricher.enrich_all_tables()
                if enriched_tables:
                    try:
                        import json
                        os.makedirs(os.path.dirname(schema_file), exist_ok=True)
                        with open(schema_file, "w", encoding="utf-8") as f:
                            json.dump(enriched_tables, f, ensure_ascii=False, indent=2)
                        logger.info("saved_schema_to_file_cache", path=schema_file)
                    except Exception as e:
                        logger.warning("failed_to_save_schema_cache", error=str(e))

            if not enriched_tables:
                logger.warning("no_tables_to_index")
                return

            documents = []
            metadatas = []
            ids = []

            for item in enriched_tables:
                documents.append(item["document"])
                metadatas.append(item["metadata"])
                ids.append(item["table_name"])

            if self.vector_enabled:
                current_count = self.collection.count()
                if current_count == len(enriched_tables):
                    logger.info("schema_already_indexed_skipping_write", count=current_count)
                else:
                    # Clear existing ChromaDB data
                    if current_count > 0:
                        existing = self.collection.get()
                        if existing["ids"]:
                            self.collection.delete(ids=existing["ids"])

                    # Index into ChromaDB
                    self.collection.add(
                        documents=documents,
                        metadatas=metadatas,
                        ids=ids,
                    )

            # Build BM25 index
            self.documents = documents
            self.doc_ids = ids
            tokenized = [doc.lower().split() for doc in documents]
            self.bm25 = BM25Okapi(tokenized)

            logger.info(
                "schema_indexed",
                tables=len(enriched_tables),
                chroma_count=self.collection.count(),
            )

        except Exception as e:
            logger.error("schema_indexing_failed", error=str(e))

    def retrieve(self, query: str, top_k: int = 5) -> list[str]:
        """
        Retrieve relevant schema documents using hybrid search.
        Combines ChromaDB vector search with BM25 keyword search,
        then optionally reranks with a cross-encoder for precision.
        """
        if not self.documents:
            logger.warning("empty_index_fallback")
            return [self.db_pool.get_full_schema()]

        try:
            # ── Vector search (ChromaDB) ─────────────────
            vector_docs = self._vector_search(query, top_k)

            # ── Keyword search (BM25) ────────────────────
            bm25_docs = self._keyword_search(query, top_k)

            # ── Reciprocal Rank Fusion ───────────────────
            # Fetch more candidates than needed for reranking
            merge_k = top_k * 2 if self._get_reranker() else top_k
            merged = self._rrf_merge(vector_docs, bm25_docs, merge_k)

            if not merged:
                # Fallback: return all documents
                return self.documents

            # ── Cross-encoder reranking (optional) ───────
            reranked = False
            reranker = self._get_reranker()
            if reranker and len(merged) > top_k:
                merged = self._rerank(query, merged, top_k)
                reranked = True

            result = merged[:top_k]

            logger.info(
                "retrieval_complete",
                vector_count=len(vector_docs),
                bm25_count=len(bm25_docs),
                merged_count=len(result),
                reranked=reranked,
            )

            return result

        except Exception as e:
            logger.error("retrieval_failed", error=str(e))
            return [self.db_pool.get_full_schema()]

    def retrieve_expanded(self, query: str, entities: list[str] = None, top_k: int = 5) -> list[str]:
        """
        Retrieve with query expansion for better recall on multi-table queries.

        When a user says 'revenue by region', the system needs both the sales
        and customers tables. Query expansion generates multiple search queries
        to ensure all needed schemas are retrieved.

        Falls back to standard retrieve() if entities are empty.
        """
        if not entities or len(entities) <= 1:
            return self.retrieve(query, top_k)

        try:
            expanded_queries = self._expand_query(query, entities)
            all_docs = []
            seen_hashes = set()

            for q in expanded_queries:
                docs = self.retrieve(q, top_k=top_k)
                for doc in docs:
                    doc_hash = hash(doc[:200])  # Hash first 200 chars for dedup
                    if doc_hash not in seen_hashes:
                        seen_hashes.add(doc_hash)
                        all_docs.append(doc)

            # Rerank the combined set to select the best top_k
            reranker = self._get_reranker()
            if reranker and len(all_docs) > top_k:
                all_docs = self._rerank(query, all_docs, top_k)

            result = all_docs[:top_k]

            logger.info(
                "expanded_retrieval_complete",
                num_queries=len(expanded_queries),
                total_candidates=len(all_docs),
                returned=len(result),
            )

            return result

        except Exception as e:
            logger.warning("expanded_retrieval_failed", error=str(e))
            return self.retrieve(query, top_k)

    @staticmethod
    def _expand_query(query: str, entities: list[str]) -> list[str]:
        """
        Generate multiple search queries for better recall.

        Strategy:
        1. Original query (captures user intent)
        2. Per-entity queries (ensures each table's schema is searched)
        3. Relationship query (helps find JOIN paths between entities)
        """
        queries = [query]

        # Per-entity focused queries
        for entity in entities:
            queries.append(f"{entity} table schema columns relationships")

        # Cross-entity relationship query
        if len(entities) > 1:
            queries.append(f"relationship between {' and '.join(entities)} foreign key join")

        return queries

    def _vector_search(self, query: str, top_k: int) -> list[str]:
        """ChromaDB semantic similarity search."""
        if not self.vector_enabled:
            return []

        try:
            results = self.collection.query(
                query_texts=[query],
                n_results=min(top_k, self.collection.count()),
            )
            return results["documents"][0] if results["documents"] else []
        except Exception as e:
            logger.warning("vector_search_failed", error=str(e))
            return []

    def _keyword_search(self, query: str, top_k: int) -> list[str]:
        """BM25 keyword search for exact table/column name matching."""
        if not self.bm25:
            return []

        try:
            tokenized_query = query.lower().split()
            scores = self.bm25.get_scores(tokenized_query)

            # Get top-k indices sorted by score
            top_indices = sorted(
                range(len(scores)),
                key=lambda i: scores[i],
                reverse=True,
            )[:top_k]

            # Filter out zero-score results
            return [
                self.documents[i]
                for i in top_indices
                if scores[i] > 0
            ]
        except Exception as e:
            logger.warning("bm25_search_failed", error=str(e))
            return []

    @staticmethod
    def _rrf_merge(list_a: list[str], list_b: list[str], top_k: int, k: int = 60) -> list[str]:
        """
        Reciprocal Rank Fusion — merges two ranked lists.
        RRF score = Σ 1/(k + rank) for each list the document appears in.
        k=60 is the standard constant from the original RRF paper.
        """
        scores: dict[str, float] = {}

        for rank, doc in enumerate(list_a):
            scores[doc] = scores.get(doc, 0) + 1.0 / (k + rank + 1)

        for rank, doc in enumerate(list_b):
            scores[doc] = scores.get(doc, 0) + 1.0 / (k + rank + 1)

        sorted_docs = sorted(scores.items(), key=lambda x: x[1], reverse=True)
        return [doc for doc, _ in sorted_docs[:top_k]]

    def _get_reranker(self):
        """Lazy-load the cross-encoder reranker."""
        if os.environ.get("DISABLE_ML_INTENT", "false").lower() in ("true", "1", "yes") or not self.vector_enabled:
            return None

        if self._reranker_loaded:
            return self._reranker

        self._reranker_loaded = True
        try:
            from sentence_transformers import CrossEncoder
            self._reranker = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2")
            logger.info("cross_encoder_reranker_loaded")
        except ImportError:
            logger.info("reranker_unavailable", hint="pip install sentence-transformers for reranking")
        except Exception as e:
            logger.warning("reranker_load_failed", error=str(e))

        return self._reranker

    def _rerank(self, query: str, docs: list[str], top_k: int) -> list[str]:
        """Rerank documents using cross-encoder for precise relevance scoring."""
        try:
            pairs = [(query, doc) for doc in docs]
            scores = self._reranker.predict(pairs)
            ranked = sorted(zip(docs, scores), key=lambda x: x[1], reverse=True)
            return [doc for doc, _ in ranked[:top_k]]
        except Exception as e:
            logger.warning("reranking_failed", error=str(e))
            return docs[:top_k]

    def refresh_index(self):
        """Re-index the schema (call after schema changes)."""
        logger.info("reindexing_schema")
        if hasattr(self.db_pool, 'clear_schema_cache'):
            self.db_pool.clear_schema_cache()
            
        schema_file = os.path.join(self.chroma_client.persist_directory if self.chroma_client and hasattr(self.chroma_client, "persist_directory") else "./chroma_db", "enriched_schema.json")
        if os.path.exists(schema_file):
            try:
                os.remove(schema_file)
                logger.info("deleted_schema_file_cache")
            except Exception as e:
                logger.warning("failed_to_delete_schema_file_cache", error=str(e))
                
        self._index_schema()