app/query.py

"""Query side: text or Hindi-audio -> ranked moments.

  text query  -> bge-m3 embed -> LanceDB (dense-dominant hybrid) -> bge-reranker -> top-k
  audio query -> ffmpeg -> faster-whisper(hi) -> Hindi text -> (same as above)

An English query matches Hindi passages natively in bge-m3's shared space, so there is no
query-time translation. The reranker reads (query, Hindi passage) pairs together and is
the main precision lever.
"""
from __future__ import annotations

import re
import tempfile
import threading
from pathlib import Path
from typing import List, Optional, Tuple

from app.config import Config, get_config
from app.models import SearchResult
from app.store import Store

# A query wrapped in quotes (straight or curly) means: exact phrase, FTS only.
_EXACT_RE = re.compile(r'^\s*["“](.+)["”]\s*$', re.S)


def exact_phrase(query: str) -> Optional[str]:
    """The phrase inside a fully-quoted query, or None for a semantic query."""
    m = _EXACT_RE.match(query or "")
    phrase = m.group(1).strip() if m else ""
    return phrase or None


class Searcher:
    def __init__(self, cfg: Optional[Config] = None):
        self.cfg = cfg or get_config()
        self.store = Store(self.cfg)
        self._embedder = None
        self._reranker = None
        self._transcriber = None
        # The server warms models on a background thread; without a lock a user
        # request racing the warm-up would load bge-m3 twice.
        self._load_lock = threading.Lock()

    @property
    def embedder(self):
        if self._embedder is None:
            with self._load_lock:
                if self._embedder is None:
                    from app.embed import Embedder

                    self._embedder = Embedder(self.cfg)
        return self._embedder

    @property
    def reranker(self):
        if self._reranker is None:
            from app.rerank import Reranker

            self._reranker = Reranker(self.cfg)
        return self._reranker

    @property
    def transcriber(self):
        if self._transcriber is None:
            from app.asr import Transcriber

            self._transcriber = Transcriber(self.cfg)
        return self._transcriber

    # ---- text search ----------------------------------------------------
    def search_text(
        self,
        query: str,
        top_k: Optional[int] = None,
        recording_id: Optional[str] = None,
    ) -> List[SearchResult]:
        query = (query or "").strip()
        if not query:
            return []
        top_k = top_k or self.cfg.search["top_k"]
        where = None
        if recording_id:
            where = "recording_id = '{}'".format(recording_id.replace("'", "''"))

        phrase = exact_phrase(query)
        if phrase:
            # Quoted query: exact FTS phrase match, no embedding involved. score=0
            # signals the UI to show an "exact" badge instead of a percentage.
            rows = self.store.search_fts(phrase, k=top_k, where=where, phrase=True)
            for r in rows:
                r["score"] = 0.0
            return [self._to_result(r) for r in rows[:top_k]]

        use_rerank = self.reranker.enabled
        pool = max(self.cfg.reranker["candidates"], top_k) if use_rerank else top_k

        qvec = self.embedder.embed_query(query)
        rows = self.store.search(
            qvec,
            query_text=query,
            k=pool,
            hybrid=bool(self.cfg.search["hybrid"]),
            vector_weight=float(self.cfg.search["vector_weight"]),
            candidate_pool=max(pool, 80),
            where=where,
        )
        if not rows:
            return []

        if use_rerank:
            scores = self.reranker.scores(query, [r["hindi_text"] for r in rows])
            for r, s in zip(rows, scores):
                r["rerank_score"] = s
            rows.sort(key=lambda r: r.get("rerank_score", 0.0), reverse=True)

        rows = rows[:top_k]
        return [self._to_result(r) for r in rows]

    # ---- audio search ---------------------------------------------------
    def search_audio(self, audio_path: str | Path, top_k: Optional[int] = None) -> Tuple[str, List[SearchResult]]:
        """Transcribe a Hindi clip, then search. Returns (recognized_text, results)."""
        from app.audio_utils import normalize_to_wav

        with tempfile.TemporaryDirectory() as tmp:
            wav = Path(tmp) / "query.wav"
            normalize_to_wav(audio_path, wav)
            text = self.transcriber.transcribe_query(wav)
        results = self.search_text(text, top_k=top_k) if text.strip() else []
        return text, results

    # ---- helpers --------------------------------------------------------
    @staticmethod
    def _to_result(r: dict) -> SearchResult:
        # Prefer the reranker score, else the dense cosine (nice 0..1 for display);
        # fused RRF score is only a last resort (tiny values, ranking-only).
        score = r.get("rerank_score")
        if score is None:
            score = r.get("score", r.get("fused_score", 0.0))
        return SearchResult(
            recording_id=r["recording_id"],
            source_file=r["source_file"],
            start_ms=int(r["start_ms"]),
            end_ms=int(r["end_ms"]),
            hindi_text=r.get("hindi_text", ""),
            english_gloss=r.get("english_gloss", ""),
            score=float(score),
            rerank_score=(float(r["rerank_score"]) if "rerank_score" in r else None),
        )