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

"""Embedding generation for semantic search via pgvector.

Priority:
  1. LiteLLM proxy — model from Redis settings (overrides env)
  2. LiteLLM proxy — model from LITELLM_EMBEDDING_MODEL env
  3. AWS Bedrock Titan Embed V2 (direct fallback)
"""
import hashlib
import json
import logging
from datetime import datetime

from app.config import settings

logger = logging.getLogger(__name__)


def _get_embedding_model() -> str:
    """The active embedding model name — also the provenance stamp on each vector."""
    try:
        import redis as redis_lib
        r = redis_lib.from_url(settings.REDIS_URL, decode_responses=True)
        model = r.get("settings:embedding_model")
        if model:
            return model
    except Exception:
        pass
    return settings.LITELLM_EMBEDDING_MODEL


def _text_for_question(question_text: str, options: list[str] | None) -> str:
    """Build the text to embed for a question — stem + options, no explanation."""
    parts = [question_text]
    if options:
        parts.extend(options)
    return " ".join(parts)[:4000]


def generate_embedding(text: str) -> list[float] | None:
    """Return a vector of EMBEDDING_DIMENSIONS floats, or None.

    None means "not embedded yet", never "embedded badly": a wrong-sized vector
    is rejected rather than stored, because the index cannot compare it. The
    `retry_missing_embeddings` task refills whatever a failure leaves behind.

    Order: the LiteLLM proxy (BGE-M3), then Bedrock.
    """
    clean = " ".join(text.split())[:4000]
    if not clean:
        return None

    # ── LiteLLM proxy (direct httpx — avoids LiteLLM param validation) ──
    # Must be the same resolver that stamps provenance, or a vector gets labelled
    # with a model that did not produce it — which is the exact failure the
    # provenance columns exist to catch.
    embedding_model = _get_embedding_model()
    api_base = (settings.LITELLM_API_BASE or "").rstrip("/").removesuffix("/v1")
    if embedding_model and settings.LITELLM_API_KEY and api_base:
        try:
            import httpx, json as _json
            body: dict = {"model": embedding_model, "input": [clean]}
            # Only OpenAI's embedding-3 family accepts a truncation size; sending
            # it to BGE-M3 is rejected outright.
            if "embedding-3" in embedding_model:
                body["dimensions"] = settings.EMBEDDING_DIMENSIONS
            resp = httpx.post(
                f"{api_base}/v1/embeddings",
                headers={"Authorization": f"Bearer {settings.LITELLM_API_KEY}", "Content-Type": "application/json"},
                content=_json.dumps(body),
                timeout=30,
            )
            resp.raise_for_status()
            emb = resp.json()["data"][0]["embedding"]
            if len(emb) == settings.EMBEDDING_DIMENSIONS:
                return emb
            logger.warning(f"Embedding dim mismatch: got {len(emb)}, expected {settings.EMBEDDING_DIMENSIONS}")
        except Exception as e:
            logger.warning(f"LiteLLM embedding failed: {e}")

    # ── 2. AWS Bedrock Titan direct (fallback) ──────────────────
    if settings.AWS_ACCESS_KEY_ID and settings.AWS_SECRET_ACCESS_KEY:
        try:
            import boto3, json
            client = boto3.client(
                "bedrock-runtime",
                aws_access_key_id=settings.AWS_ACCESS_KEY_ID,
                aws_secret_access_key=settings.AWS_SECRET_ACCESS_KEY,
                region_name=settings.AWS_BEDROCK_REGION or "us-east-1",
            )
            body = json.dumps({
                "inputText": clean,
                "dimensions": settings.EMBEDDING_DIMENSIONS,
                "normalize": True,
            })
            resp = client.invoke_model(
                modelId="amazon.titan-embed-text-v2:0",
                body=body,
                contentType="application/json",
                accept="application/json",
            )
            return json.loads(resp["body"].read())["embedding"]
        except Exception as e:
            logger.warning(f"Bedrock embedding failed: {e}")

    return None


# What each embeddable type contributes to its vector. Adding a type here is
# all that is needed for the retry task and the health report to cover it.
EMBEDDABLE = {
    "question": lambda row: _join(row.question_text, *(row.options or [])),
    "article": lambda row: _join(row.title, row.summary, row.content),
    "flashcard": lambda row: _join(row.front, row.back),
    "article_section": lambda row: _join(row.title, row.content),
    # Text today; a vision model can embed the image itself later without
    # changing anything here but this line.
    "media": lambda row: _join(row.title, row.caption, row.alt_text),
}


def _join(*parts) -> str:
    return " ".join(part for part in parts if part)[:4000]


def embed_record(row, kind: str) -> bool:
    """Embed any supported row, stamping the model that produced the vector.

    One code path for questions, articles and cards: they share an embedding
    space, so they must share the model and the provenance rules too.
    """
    build = EMBEDDABLE.get(kind)
    if build is None:
        raise ValueError(f"Unknown embeddable kind: {kind}")
    embedding = generate_embedding(build(row))
    if not embedding:
        return False
    row.embedding = embedding
    row.embedding_model = _get_embedding_model()
    row.embedded_at = datetime.utcnow()
    return True


def embed_question(question) -> bool:
    """Generate and store an embedding, recording which model produced it.

    Stamping the model is what makes a later model change detectable: vectors
    from two different models share no space, so cosine distance between them is
    noise. `stale_embedding_counts` reports the mix; the regenerate task re-embeds
    anything not on the active model.
    """
    return embed_record(question, "question")


def embeddable_models() -> dict:
    """The ORM class behind each embeddable kind."""
    from app.models import media  # noqa — registers media_assets on the mapper.
    from app.models.article import Article, ArticleSectionIndex
    from app.models.flashcard import Flashcard
    from app.models.question import Question

    from app.models.media import MediaAsset

    return {"question": Question, "article": Article, "flashcard": Flashcard,
            "article_section": ArticleSectionIndex, "media": MediaAsset}


def stale_embedding_counts(db) -> dict:
    """How much of each corpus is semantically searchable, and how much is stale.

    A vector from another model sits in a different space, so "stale" is not a
    performance nicety — those rows return meaningless distances until re-embedded.
    """
    from sqlalchemy import func

    active = _get_embedding_model()
    by_kind, totals = {}, {"total": 0, "missing": 0, "stale": 0, "current": 0}
    for kind, model in embeddable_models().items():
        total = db.query(func.count(model.id)).scalar() or 0
        missing = db.query(func.count(model.id)).filter(model.embedding.is_(None)).scalar() or 0
        stale = db.query(func.count(model.id)).filter(
            model.embedding.isnot(None),
            (model.embedding_model.is_(None)) | (model.embedding_model != active),
        ).scalar() or 0
        current = max(0, total - missing - stale)
        by_kind[kind] = {"total": total, "missing": missing, "stale": stale, "current": current}
        for key, value in (("total", total), ("missing", missing), ("stale", stale), ("current", current)):
            totals[key] += value
    return {
        "active_model": active,
        **totals,
        "by_kind": by_kind,
        "needs_regeneration": totals["missing"] + totals["stale"] > 0,
    }