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.

Fusion decides *which* rows are candidates. `rerank_ids` may then reorder the
head of that list with a cross-encoder (see `rerank_service`), which is a strict
permutation: recall belongs to the two rankers here, and a scoring service being
unwell must never subtract a result.
"""
import hashlib
import json
import logging
import re

from sqlalchemy import bindparam as sa_bindparam
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 top_similarity(db: Session, query_text: str,
                   kinds=("article", "article_section")) -> float | None:
    """The best cosine similarity any row in these corpora has to the query.

    The one calibrated number retrieval produces. `hybrid_ids` fuses two
    rankers by reciprocal rank and throws the distances away, so what comes
    back is an order with no sense of scale — and an order is always non-empty
    if either ranker matched anything at all. That is why a question about
    photosynthesis came back with six paediatric sources and an instruction to
    answer only from them.

    `None` means the question could not be asked — no vector database, or the
    encoder is down — which is a different thing from "nothing is close" and
    must not be collapsed into it. Zero is a real measurement of nothing.
    """
    if not _is_postgres(db):
        return None
    embedding = _query_embedding((query_text or "").strip())
    if not embedding:
        return None
    literal = "[" + ",".join(str(float(value)) for value in embedding) + "]"
    best = 0.0
    for kind in kinds:
        if kind not in CORPORA:
            continue
        table, _ = CORPORA[kind]
        row = db.execute(sa_text(f"""
            SELECT 1 - (embedding <=> CAST(:vec AS vector)) AS similarity
            FROM {table}
            WHERE embedding IS NOT NULL
            ORDER BY embedding <=> CAST(:vec AS vector)
            LIMIT 1
        """), {"vec": literal}).fetchone()
        if row is not None:
            best = max(best, float(row.similarity))
    return best


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)


def article_ids_with_sections(db: Session, query_text: str,
                              limit: int = 200) -> tuple[list[int], dict[int, list]]:
    """Article ids for a query, plus the section rows that matched, grouped by article.

    An article's body lives in its sections, and the article's own row carries
    only a topical vector and a weighted summary of that body. So a term that
    appears in one section and nowhere else — a drug, a procedure, an eponym —
    is found by searching the section corpus, not the article corpus.

    Section hits are reported under their article rather than beside it: ten
    matching sections of one article are one result with ten places to start
    reading, not ten results that bury everything else.
    """
    from app.models.article import ArticleSectionIndex

    ranked, _ = hybrid_ids(db, query_text, "article", limit=limit)
    section_ranked, _ = hybrid_ids(db, query_text, "article_section", limit=limit)
    # Only the sections are reranked, and it is the article ordering that this
    # buys. An article row carries a title, a summary and a topical vector; the
    # prose a query is actually about is in its sections, so handing a
    # cross-encoder the article row means handing it a stub and asking about a
    # document it cannot see. Its best section stands for it instead — which is
    # already how the fusion below decides an article's place.
    section_ranked = rerank_ids(db, query_text, "article_section", section_ranked)

    rows = {row.id: row for row in db.query(ArticleSectionIndex).filter(
        ArticleSectionIndex.id.in_(section_ranked)).all()} if section_ranked else {}
    by_article: dict[int, list] = {}
    best: dict[int, int] = {}
    for index, row_id in enumerate(section_ranked):
        row = rows.get(row_id)
        if row is None:
            continue
        by_article.setdefault(row.article_id, []).append(row)
        # Only an article's best-placed section scores. Counting them all would
        # rank a fourteen-section article above a better two-section one on
        # length alone.
        best.setdefault(row.article_id, index)

    # Fused, not concatenated. Appending the section hits behind the article
    # hits put the strongest evidence there is — a section that matched at rank
    # one — behind every weak whole-article match, so it never reached the page.
    scores: dict[int, float] = {}
    for index, article_id in enumerate(ranked):
        scores[article_id] = scores.get(article_id, 0.0) + 1.0 / (RRF_K + index + 1)
    for article_id, index in best.items():
        scores[article_id] = scores.get(article_id, 0.0) + 1.0 / (RRF_K + index + 1)

    ordered = sorted(scores, key=lambda article_id: (-scores[article_id], article_id))[:limit]
    return ordered, by_article


# What a cross-encoder is shown of each row. Not the same as the columns the
# lexical fallback scans: `CAST(options AS TEXT)` helps a substring match find a
# drug name buried in an option, but as prose it is a JSON array, and paying a
# cross-encoder to read punctuation makes its judgement worse rather than
# better. Title first everywhere, because a truncated document keeps its head.
RERANK_TEXT = {
    "question": "question_text",
    "article": "coalesce(title, '') || '. ' || coalesce(summary, '')",
    "article_section": "coalesce(title, '') || '. ' || coalesce(content, '')",
    "flashcard": "coalesce(front, '') || ' — ' || coalesce(back, '')",
    "media": "coalesce(title, '') || '. ' || coalesce(caption, '') || ' ' || coalesce(alt_text, '')",
}


def rerank_ids(db: Session, query_text: str, kind: str, ranked: list[int]) -> list[int]:
    """`ranked` reordered by a cross-encoder, or exactly `ranked` if there isn't one.

    The contract is narrow on purpose: this is a permutation. Fusion decides
    which rows exist and every visibility rule downstream still runs on the same
    set, so a reranker that is off, down or wrong costs a worse order and
    nothing else.

    Only the head is fetched and scored. The rows past it are ones neither cheap
    ranker put near the front, and reading a 200-row pool out of the database to
    ask a model about it would cost more than the answer is worth.
    """
    from app.services import rerank_service

    query_text = (query_text or "").strip()
    if len(ranked) < 2 or kind not in RERANK_TEXT or not query_text:
        return ranked
    if not rerank_service.is_configured():
        return ranked

    head = ranked[:rerank_service.MAX_CANDIDATES]
    table, _ = CORPORA[kind]
    try:
        statement = sa_text(
            f"SELECT id, {RERANK_TEXT[kind]} AS body FROM {table} WHERE id IN :ids"
        ).bindparams(sa_bindparam("ids", expanding=True))
        bodies = {row.id: row.body or "" for row in db.execute(statement, {"ids": head}).fetchall()}
    except Exception:
        logger.warning("Could not read %s text for reranking; keeping fused order",
                       kind, exc_info=True)
        return ranked
    if not bodies:
        return ranked
    return rerank_service.reorder(query_text, ranked, lambda row_id: bodies.get(row_id, ""))