pdf-quiz-generator/backend/app/services/search_service.py

"""Hybrid question retrieval: Postgres full text fused with pgvector similarity.

Why not OpenSearch/Elasticsearch: the lexical half of this problem is ordinary
ranked text matching that Postgres already does with `tsvector`/`ts_rank_cd`,
and the semantic half already runs on pgvector with embeddings that are stored
and kept current. A search cluster would add a second datastore to keep in sync,
a JVM's memory footprint on this host, and a new failure mode, to replace an
index Postgres maintains for free inside the same transaction.

The two rankers are combined with Reciprocal Rank Fusion rather than a weighted
score, because a BM25-style rank and a cosine distance are not on comparable
scales; RRF only needs each ranker's ordering.

Retrieval is always hybrid. A keyword-only mode looks precise but silently drops
the question that asks the same thing in different words, which is exactly the
question a learner searching a concept wants.
"""
import hashlib
import json
import logging
import re

from sqlalchemy import text as sa_text
from sqlalchemy.orm import Session

logger = logging.getLogger(__name__)

# Rank-fusion constant. 60 is the value from the original RRF paper; it damps
# the head of each list so one ranker cannot dominate on its top hit alone.
RRF_K = 60
# Per-ranker candidate pool. Wider than the page so fusion has room to reorder.
POOL_MULTIPLIER = 4
MIN_POOL = 60
SEMANTIC_FLOOR = 0.45
# Query embeddings are deterministic per model, so a day is safe and cheap.
QUERY_CACHE_TTL = 24 * 3600


# Enough prose that it is a document, not a query.
LONG_TEXT_CHARS = 300
# Words carried by every clinical vignette, so useless for discriminating.
STOPWORDS = {
    "the", "and", "for", "with", "that", "this", "from", "have", "has", "had",
    "are", "was", "were", "been", "being", "which", "their", "there", "would",
    "could", "should", "about", "after", "before", "during", "into", "over",
    "under", "than", "then", "them", "they", "you", "your", "his", "her",
    "its", "not", "but", "all", "any", "one", "two", "who", "whom", "what",
    "when", "where", "how", "why", "most", "more", "also", "other", "some",
    "such", "only", "own", "same", "each", "both", "will", "can", "may",
    "old", "year", "years", "month", "months", "day", "days", "patient",
    "following", "next", "step", "best", "likely", "presents", "history",
}


def _query_terms(text: str, limit: int = 18) -> str:
    """Reduce a document to its most distinctive terms, OR-joined.

    A whole document handed to `websearch_to_tsquery` becomes one enormous
    conjunction that matches nothing. The salient nouns, joined with OR, are
    what actually discriminate between bank questions.
    """
    from collections import Counter

    words = re.findall(r"[a-zA-Z][a-zA-Z-]{3,}", text.lower())
    counts = Counter(word for word in words if word not in STOPWORDS)
    if not counts:
        return ""
    return " or ".join(word for word, _ in counts.most_common(limit))


def _is_postgres(db: Session) -> bool:
    return db.bind is not None and db.bind.dialect.name == "postgresql"


# Where each searchable kind lives, and which columns a SQLite fallback scans.
CORPORA = {
    "question": ("questions", ("question_text", "CAST(options AS TEXT)")),
    "article": ("articles", ("title", "summary", "content")),
    "flashcard": ("flashcards", ("front", "back")),
    "article_section": ("article_section_index", ("title", "content")),
    "media": ("media_assets", ("title", "caption", "alt_text")),
}


def _lexical_ranked(db: Session, query_text: str, pool: int, kind: str = "question") -> list[int]:
    """Row ids by full-text relevance, best first.

    `websearch_to_tsquery` gives quoted phrases exact-match semantics for free:
    "absence seizure" matches the phrase, bare words match either. That covers
    the one case a keyword-only mode was ever needed for, per query rather than
    as a sticky setting.
    """
    table, columns = CORPORA[kind]
    if len(query_text) > LONG_TEXT_CHARS:
        query_text = _query_terms(query_text) or query_text[:200]
    if not _is_postgres(db):
        # SQLite (tests): substring matching keeps the fusion path exercised.
        # Mirror the OR semantics of the Postgres path on any input length.
        terms = [t for t in re.split(r"[^a-z0-9]+", query_text.strip('"').lower()) if len(t) > 3][:8]
        clauses, params = [], {"pool": pool}
        for index, term in enumerate(terms):
            params[f"t{index}"] = f"%{term}%"
            clauses.append(" OR ".join(f"lower({column}) LIKE :t{index}" for column in columns))
        if not clauses:
            return []
        rows = db.execute(sa_text(
            f"SELECT id FROM {table} WHERE {' OR '.join(clauses)} ORDER BY id LIMIT :pool"
        ), params).fetchall()
        return [row[0] for row in rows]
    rows = db.execute(sa_text(f"""
        SELECT id FROM {table}
        WHERE search_vector @@ websearch_to_tsquery('english', :q)
        ORDER BY ts_rank_cd(search_vector, websearch_to_tsquery('english', :q)) DESC, id
        LIMIT :pool
    """), {"q": query_text, "pool": pool}).fetchall()
    return [row[0] for row in rows]


def _query_embedding(query_text: str) -> list[float] | None:
    """Embed a search query, cached per model so typing is not a round-trip per keystroke.

    Cached under the active model's name, so switching models cannot serve a
    vector from the previous embedding space.
    """
    from app.services.embedding_service import _get_embedding_model, generate_embedding

    model = _get_embedding_model()
    key = f"qsearch:emb:{model}:{hashlib.sha256(query_text.encode()).hexdigest()[:32]}"
    cache = None
    try:
        import redis as redis_lib

        from app.config import settings

        cache = redis_lib.from_url(settings.REDIS_URL, decode_responses=True)
        hit = cache.get(key)
        if hit:
            return json.loads(hit)
    except Exception:
        cache = None  # Redis is a nicety here; the embedder still answers.

    embedding = generate_embedding(query_text)
    if embedding and cache is not None:
        try:
            cache.setex(key, QUERY_CACHE_TTL, json.dumps(embedding))
        except Exception:
            logger.debug("Could not cache query embedding", exc_info=True)
    return embedding


def _semantic_ranked(db: Session, query_text: str, pool: int, kind: str = "question") -> list[int]:
    """Row ids by embedding similarity, nearest first."""
    if not _is_postgres(db):
        return []
    # A quoted phrase asks for an exact lookup, so the fuzzy ranker sits it out.
    if query_text.startswith('"') and query_text.endswith('"') and len(query_text) > 2:
        return []
    embedding = _query_embedding(query_text)
    if not embedding:
        return []
    table, _ = CORPORA[kind]
    literal = "[" + ",".join(str(float(value)) for value in embedding) + "]"
    rows = db.execute(sa_text(f"""
        SELECT id, 1 - (embedding <=> CAST(:vec AS vector)) AS similarity
        FROM {table}
        WHERE embedding IS NOT NULL
        ORDER BY embedding <=> CAST(:vec AS vector)
        LIMIT :pool
    """), {"vec": literal, "pool": pool}).fetchall()
    return [row.id for row in rows if float(row.similarity) >= SEMANTIC_FLOOR]


def hybrid_ids(db: Session, query_text: str, kind: str = "question",
               limit: int = 200) -> tuple[list[int], set[int]]:
    """Return (ids best-first, ids the semantic ranker contributed), for any corpus.

    The result is the *union* of both rankers. An earlier implementation
    intersected them, so a row that matched the meaning but not the literal
    string could never be returned no matter how well it scored.
    """
    query_text = (query_text or "").strip()
    if not query_text or kind not in CORPORA:
        return [], set()
    pool = max(MIN_POOL, limit * POOL_MULTIPLIER)

    try:
        lexical = _lexical_ranked(db, query_text, pool, kind)
    except Exception:
        logger.warning("Lexical search unavailable for %s; semantic only", kind, exc_info=True)
        lexical = []
    try:
        semantic = _semantic_ranked(db, query_text, pool, kind)
    except Exception:
        logger.warning("Semantic search unavailable for %s; lexical only", kind, exc_info=True)
        semantic = []

    scores: dict[int, float] = {}
    for ranked in (lexical, semantic):
        for position, question_id in enumerate(ranked):
            scores[question_id] = scores.get(question_id, 0.0) + 1.0 / (RRF_K + position + 1)

    ordered = sorted(scores, key=lambda qid: (-scores[qid], qid))
    return ordered[:limit], set(semantic)


def hybrid_question_ids(db: Session, query_text: str, limit: int = 200) -> tuple[list[int], set[int]]:
    """Questions, for callers that predate the multi-corpus signature."""
    return hybrid_ids(db, query_text, "question", limit)