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


# One vector cannot hold a whole article: the body averages 8k characters of
# section text and runs past 13k, while `_join` and `generate_embedding` clamp
# at 4000. Budget below that so the clamp never has the last word — a silent
# truncation is how an index comes to be confidently wrong.
ARTICLE_EMBED_BUDGET = 3600
# Below this an excerpt is too short to say anything, so a very long article
# spends a little over budget rather than reducing every section to a phrase.
ARTICLE_MIN_EXCERPT = 90


def article_embedding_text(article) -> str:
    """Compose what a whole article contributes to its vector.

    `content` holds the body for only eight hand-seeded articles; every
    generated one puts its prose in the `sections` JSON, so title and summary
    alone described 98% of the library and the article vector could not tell
    two respiratory topics apart.

    Every section gets an even slice rather than the head being kept: head
    truncation of a twelve-section article stops somewhere in the
    pathophysiology, so treatment, management and complications — the half a
    learner actually searches — contribute nothing at all. Section titles go in
    whole, being the densest signal per character available.

    This stays a topical signal by design. Depth belongs to `article_section`
    rows, where the longest section in the corpus still fits under the clamp
    intact, so no sentence of the body goes unembedded anywhere.
    """
    sections = [s for s in (article.sections or []) if isinstance(s, dict)]
    parts = [article.title, article.summary, article.content]
    parts.extend(s.get("title") for s in sections if s.get("title"))
    spent = sum(len(part) + 1 for part in parts if part)

    bodies = [body for body in ((s.get("content") or "").strip() for s in sections) if body]
    if bodies:
        excerpt = max(ARTICLE_MIN_EXCERPT, (ARTICLE_EMBED_BUDGET - spent) // len(bodies))
        parts.extend(body[:excerpt] for body in bodies)
    return _join(*parts)


# 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": article_embedding_text,
    "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]


# One request per row costs a round trip each; an article with fourteen sections
# spent fourteen of them on a single save. The proxy takes a list, so a save is
# one call and a corpus sweep is hundreds rather than thousands.
EMBED_BATCH = 32


def generate_embeddings(texts: list[str]) -> list[list[float] | None]:
    """Embed several texts in one round trip, aligned to the input list.

    A batch is all-or-nothing at the transport level, so any failure falls back
    to embedding the texts one at a time rather than dropping the lot: the
    fallback chain in `generate_embedding` (Bedrock, dimension checks) is the
    only place that logic should live, and a partial batch must not skip it.
    """
    clean = [" ".join(text.split())[:4000] if text else "" for text in texts]
    wanted = [index for index, text in enumerate(clean) if text]
    out: list[list[float] | None] = [None] * len(clean)
    if not wanted:
        return out

    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[i] for i in wanted]}
            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=120,
            )
            resp.raise_for_status()
            # The proxy is not required to preserve order, but it does return
            # the index it was given; trusting position alone would silently
            # attach one document's vector to another.
            data = sorted(resp.json()["data"], key=lambda item: item.get("index", 0))
            if len(data) == len(wanted):
                for slot, item in zip(wanted, data):
                    vector = item["embedding"]
                    if len(vector) == settings.EMBEDDING_DIMENSIONS:
                        out[slot] = vector
                if all(out[i] is not None for i in wanted):
                    return out
            logger.warning("Batch embedding returned %d vectors for %d inputs", len(data), len(wanted))
        except Exception as e:
            logger.warning(f"Batch embedding failed, falling back to one at a time: {e}")

    for index in wanted:
        if out[index] is None:
            out[index] = generate_embedding(clean[index])
    return out


def embed_records(rows: list, kind: str) -> int:
    """Embed a list of rows of one kind in batches; returns how many got a vector.

    Rows that fail keep whatever they had, including nothing — the retry task
    exists for exactly that, and a half-written vector is worse than none.
    """
    build = EMBEDDABLE.get(kind)
    if build is None:
        raise ValueError(f"Unknown embeddable kind: {kind}")
    model, now, done = _get_embedding_model(), datetime.utcnow(), 0
    for start in range(0, len(rows), EMBED_BATCH):
        chunk = rows[start:start + EMBED_BATCH]
        for row, embedding in zip(chunk, generate_embeddings([build(row) for row in chunk])):
            if not embedding:
                continue
            row.embedding = embedding
            row.embedding_model = model
            row.embedded_at = now
            done += 1
    return done


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,
    }