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

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


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


def _lexical_ranked(db: Session, query_text: str, pool: int) -> list[int]:
    """Question ids by full-text relevance, best first."""
    if not _is_postgres(db):
        # SQLite (tests): substring matching keeps the fusion path exercised.
        rows = db.execute(sa_text("""
            SELECT id FROM questions
            WHERE lower(question_text) LIKE :like OR lower(CAST(options AS TEXT)) LIKE :like
            ORDER BY id LIMIT :pool
        """), {"like": f"%{query_text.lower()}%", "pool": pool}).fetchall()
        return [row[0] for row in rows]
    rows = db.execute(sa_text("""
        SELECT id FROM questions
        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) -> list[int]:
    """Question ids by embedding similarity, nearest first."""
    if not _is_postgres(db):
        return []
    embedding = _query_embedding(query_text)
    if not embedding:
        return []
    literal = "[" + ",".join(str(float(value)) for value in embedding) + "]"
    rows = db.execute(sa_text("""
        SELECT id, 1 - (embedding <=> CAST(:vec AS vector)) AS similarity
        FROM questions
        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_question_ids(db: Session, query_text: str, limit: int = 200) -> tuple[list[int], set[int]]:
    """Return (ids best-first, ids the semantic ranker contributed).

    The result is the *union* of both rankers. An earlier implementation
    intersected them, so a question 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:
        return [], set()
    pool = max(MIN_POOL, limit * POOL_MULTIPLIER)

    try:
        lexical = _lexical_ranked(db, query_text, pool)
    except Exception:
        logger.warning("Lexical search unavailable; falling back to semantic only", exc_info=True)
        lexical = []
    try:
        semantic = _semantic_ranked(db, query_text, pool)
    except Exception:
        logger.warning("Semantic search unavailable; falling back to lexical only", 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)