search_engine.py

"""
GCAS Search Engine – Hybrid search pipeline (v2)

Architecture
------------
Every query in the GCAS taxonomy is fundamentally a structured-filter
problem, not a semantic-search problem.  The data is an Excel DB with
well-defined columns; the challenge is NLP (entity extraction + normalisation),
not retrieval.

New two-tier pipeline
---------------------
  Tier 1 – Structured lookup  (fast, precise, ~50 ms)
      query_planner.build_query_plan()
          → entities: college, district, university, program, gender, category …
      structured_search.lookup()
          → O(n) scan of in-memory data_store, no embeddings needed
          → used whenever ≥1 entity is resolved

  Tier 2 – FAISS semantic fallback  (for truly vague/open queries)
      embeddings.embed_query()  +  indexer.search()
          → used only when Tier 1 returns 0 results

  LLM reranking  (explicit opt-in, NOT the default)
      use_llm_rerank=true in request body
      adds ~25 s; only suitable for async / non-voice contexts

Result shaping (unchanged from v1)
-----------------------------------
  dedup by college → smart result count → field filtering
"""
from __future__ import annotations

import logging
import time
from typing import List

from config import settings
from embeddings import embed_query
from field_filter import filter_fields, smart_result_count
from llm_client import rerank_with_llm
from models import EntityCorrection, SearchRequest, SearchResponse, SearchResult
import indexer
import structured_search
from query_planner import build_query_plan

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Dedup & smart cutoff (unchanged)
# ---------------------------------------------------------------------------

def _dedup_by_college(candidates: list, intent: str) -> list:
    """
    Keep highest-scoring row per college per table.

    EXCEPTION: fees, courses, and cutoff queries intentionally have
    multiple rows per college (one per program / category).  Deduping
    those would hide e.g. all fee rows except the top-scoring program.
    """
    if intent in ("fees", "courses", "cutoff"):
        return candidates   # each row is meaningfully distinct

    seen: dict = {}
    out: list = []
    for c in candidates:
        data = c.get("data", {})
        name = (
            data.get("CollegeName")
            or data.get("College")
            or data.get("UniversityName")
        )
        if not name:
            out.append(c)
            continue
        key = (c.get("table", ""), name.strip().lower())
        if key in seen:
            existing_idx = seen[key]
            if float(c.get("score", 0)) > float(out[existing_idx].get("score", 0)):
                out[existing_idx] = c
        else:
            seen[key] = len(out)
            out.append(c)
    return out


def _smart_cutoff(candidates: list, requested_k: int) -> list:
    """
    Gap-based pruning — don't pad to requested_k with weakly related rows.

    Rules (in order):
      1. Structured results (score ≥ 0.90, all similar) → return exactly
         requested_k (smart_result_count already set the right ceiling).
      2. FAISS results: find first score gap > 0.05 → cut there.
      3. Never exceed min(8, requested_k) — hard cap.
    """
    if not candidates:
        return candidates

    cap = min(8, requested_k)
    scores = [float(c.get("score", 0)) for c in candidates]
    best = scores[0]

    # Structured lookup: all scores are near-identical (0.90–0.99).
    # smart_result_count already computed the right count — just slice to it.
    if best >= 0.90:
        return candidates[:cap]

    # FAISS path: gap-based pruning (cap at 4 for FAISS results)
    faiss_cap = min(4, requested_k)
    cutoff = 1
    for i in range(1, min(len(scores), faiss_cap)):
        if scores[i - 1] - scores[i] > 0.05:
            break
        cutoff = i + 1

    return candidates[:cutoff]


# ---------------------------------------------------------------------------
# Main search
# ---------------------------------------------------------------------------

def search(request: SearchRequest) -> SearchResponse:
    """Execute a full hybrid search and return a rich SearchResponse."""
    t_start = time.perf_counter()

    # ------------------------------------------------------------------
    # Step 1 – Build structured query plan (NLP + entity resolution)
    # ------------------------------------------------------------------
    plan = build_query_plan(request.query)

    # ------------------------------------------------------------------
    # Step 2 – Tier 1: Structured in-memory lookup
    # ------------------------------------------------------------------
    candidates = structured_search.lookup(plan, indexer._data_store)
    used_structured = bool(candidates)

    # ------------------------------------------------------------------
    # Step 3 – Tier 2: FAISS fallback (when structured lookup found nothing)
    # ------------------------------------------------------------------
    if not candidates:
        logger.info("[search] Falling back to FAISS for query: %r", request.query)
        query_vec = embed_query(plan.corrected_query or request.query)

        # Pool size: enough candidates for post-processing
        pool_size = max(request.top_k * 5, 30)
        effective_tables = plan.preferred_tables or None

        candidates = indexer.search(
            query_embedding = query_vec,
            top_k           = pool_size,
            tables          = effective_tables,
        )

    reranked = False

    # ------------------------------------------------------------------
    # Step 4 – Optional LLM reranking  (explicit opt-in only)
    # ------------------------------------------------------------------
    if request.use_llm_rerank and candidates:
        llm_pool = candidates[:50]
        try:
            ranked_items = rerank_with_llm(
                query    = request.query,
                candidates = llm_pool,
                top_k    = request.top_k,
                provider = request.llm_provider,
                model    = request.llm_model,
                api_key  = request.api_key,
            )
            remapped: List[dict] = []
            for item in ranked_items:
                idx = item.get("index")
                if idx is None or not (0 <= idx < len(llm_pool)):
                    continue
                c = dict(llm_pool[idx])
                c["score"]      = float(item.get("score", c["score"]))
                c["llm_reason"] = item.get("reason", "")
                remapped.append(c)
            if remapped:
                candidates = remapped
                reranked = True
        except Exception:
            logger.exception("LLM reranking failed — using structured/embedding scores.")

    # ------------------------------------------------------------------
    # Step 5 – Dedup + smart count + cutoff
    # ------------------------------------------------------------------
    deduped = _dedup_by_college(candidates, plan.intent)

    if request.top_k <= 4 and not used_structured:
        # Caller explicitly asked for a small number AND we're on the FAISS path
        top = deduped[:request.top_k]
    else:
        ideal = smart_result_count(deduped, plan.intent)
        top   = _smart_cutoff(deduped, ideal)

    # ------------------------------------------------------------------
    # Step 6 – Build SearchResponse (intent + gender + category field filter)
    # ------------------------------------------------------------------
    results = [
        SearchResult(
            table      = c["table"],
            row_index  = c["row_index"],
            score      = round(float(c["score"]), 6),
            llm_reason = c.get("llm_reason"),
            data       = filter_fields(
                            c["data"], c["table"], plan.intent,
                            gender   = plan.gender,
                            category = plan.category,
                         ),
        )
        for c in top
    ]

    elapsed_ms = (time.perf_counter() - t_start) * 1000

    shown_correction = (
        plan.corrected_query
        if plan.corrected_query.strip().lower() != request.query.strip().lower()
        else None
    )

    entity_corrections = [EntityCorrection(**ec) for ec in plan.entity_corrections]

    # Confidence: high when structured lookup succeeded, medium/low otherwise
    if used_structured and results:
        confidence = "high"
    elif results and float(results[0].score) >= 0.70:
        confidence = "medium"
    else:
        confidence = "low"

    return SearchResponse(
        query             = request.query,
        total_results     = len(results),
        results           = results,
        search_time_ms    = round(elapsed_ms, 2),
        reranked          = reranked,
        detected_language = plan.detected_language,
        corrected_query   = shown_correction,
        entity_corrections = entity_corrections,
        confidence_level  = confidence,
        detected_intent   = plan.intent,
        did_you_mean      = [],
    )