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

"""Retrieval and citation handling for AI Mode.

The safety property, stated plainly: **retrieval decides what may be cited, and
the server enforces it.** The model is handed a numbered shortlist and told to
cite from it by id. Afterwards every citation it wrote is checked against that
shortlist and anything else is deleted. A citation the model invented cannot
survive, because surviving is not a thing the model gets to decide.

This is the same discipline as the article page not printing answers: a property
the system holds, not one the model is trusted to respect.

Two kinds of link exist and are treated differently. *Curated* links are rows an
educator created — a question tied to an article section, a card to a question.
They are assertions, so they are trusted, and a retrieved row that carries one to
another retrieved row is boosted: two things an educator already tied together
answering the same query is evidence, not coincidence. *Retrieved* links are
ranked guesses computed per query and stored nowhere.

The shortlist is chosen twice. Each corpus offers twice as many candidates as it
can have places for, and a cross-encoder reads all of them against the query and
decides which survive and in what order. Unlike everywhere else on the site,
where reranking is a strict permutation of a page of results, here it selects —
because here the shortlist is not a page somebody skims past but the entire
evidence an answer may be built from, and a source that does not make it is one
the model cannot lean on. If no reranker answers, each corpus keeps the six it
elected, in the order it elected them, exactly as before there was one.
"""
import logging
import re

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

from app.models.article import Article, ArticleSectionIndex, QuestionArticleLink
from app.models.flashcard import Flashcard, FlashcardDeck
from app.models.question import Question
from app.models.user import User
from app.services import rerank_service
from app.services.quiz_builder import bank_query, exam_scope_predicate
from app.services.search_service import hybrid_ids, top_similarity

logger = logging.getLogger(__name__)

# How many of each kind retrieval offers the model. Small on purpose: a
# shortlist the model can hold is worth more than a corpus it skims.
PER_KIND = 6
# How many of each kind the cross-encoder gets to choose from. Handing it only
# the six a corpus already elected would leave the membership of the shortlist
# entirely to rankers that never read the query against the document; twelve
# lets a rank-nine section reach the answer, and keeps the union of four corpora
# inside a single rerank call.
CANDIDATES_PER_KIND = PER_KIND * 2
MAX_SOURCES = 14
# An excerpt long enough to answer from, short enough that fourteen of them fit.
EXCERPT_CHARS = 700
# A curated tie between two retrieved rows is evidence an educator left behind.
CURATED_BOOST = 0.5

CITATION_RE = re.compile(r"\[\[(article|section|question|card):([A-Za-z0-9#_-]+)\]\]")


def _clean(value: str | None, limit: int = EXCERPT_CHARS) -> str:
    text = re.sub(r"!\[[^\]]*\]\([^)]*\)", " ", value or "")
    text = re.sub(r"\[([^\]]*)\]\([^)]*\)", r"\1", text)
    text = re.sub(r"<[^>]+>", " ", text)
    text = re.sub(r"\s+", " ", text).strip()
    return text[:limit] + ("…" if len(text) > limit else "")


def _articles(db: Session, user: User, query: str) -> list[dict]:
    ranked, _ = hybrid_ids(db, query, "article", limit=PER_KIND * 3)
    if not ranked:
        return []
    rows = db.query(Article).filter(Article.id.in_(ranked)).all()
    if not user.is_moderator:
        rows = [a for a in rows if a.status == "published" or a.user_id == user.id]
    order = {rid: i for i, rid in enumerate(ranked)}
    rows.sort(key=lambda a: order.get(a.id, len(order)))
    return [{
        "kind": "article", "ref": str(a.id), "id": a.id,
        "title": a.title,
        "text": _clean(a.summary or a.content),
        "score": 1.0 / (1 + order.get(a.id, 0)),
    } for a in rows[:CANDIDATES_PER_KIND]]


def _sections(db: Session, user: User, query: str) -> list[dict]:
    ranked, _ = hybrid_ids(db, query, "article_section", limit=PER_KIND * 3)
    if not ranked:
        return []
    rows = db.query(ArticleSectionIndex).filter(ArticleSectionIndex.id.in_(ranked)).all()
    if not rows:
        return []
    articles = {a.id: a for a in db.query(Article).filter(
        Article.id.in_({r.article_id for r in rows})).all()}
    order = {rid: i for i, rid in enumerate(ranked)}
    rows.sort(key=lambda r: order.get(r.id, len(order)))
    out = []
    for row in rows[:CANDIDATES_PER_KIND]:
        article = articles.get(row.article_id)
        if not article:
            continue
        if article.status != "published" and not user.is_moderator and article.user_id != user.id:
            continue
        out.append({
            "kind": "section", "ref": f"{article.id}#{row.section_id}",
            "id": article.id, "section_id": row.section_id,
            "title": f"{article.title} › {row.title or 'section'}",
            "text": _clean(row.content),
            "score": 1.0 / (1 + order.get(row.id, 0)),
        })
    return out


def _questions(db: Session, user: User, query: str) -> list[dict]:
    ranked, _ = hybrid_ids(db, query, "question", limit=PER_KIND * 3)
    if not ranked:
        return []
    q = bank_query(db, user).filter(Question.id.in_(ranked))
    scope = exam_scope_predicate(db, user)
    if scope is not None:
        q = q.filter(scope)
    order = {rid: i for i, rid in enumerate(ranked)}
    rows = sorted(q.all(), key=lambda r: order.get(r.id, len(order)))
    return [{
        "kind": "question", "ref": str(row.id), "id": row.id,
        # The opening of the stem, not the row's number. "Question #2320" tells
        # a reader nothing about whether it is worth following, and the list of
        # sources under an answer is exactly where that has to be decidable.
        "title": _clean(row.question_text, 90) or f"Question #{row.id}",
        # The stem only. An answer belongs to the quiz runner, and a chat that
        # printed it would hand away the practice it is meant to prepare for.
        "text": _clean(row.question_text, 320),
        "score": 1.0 / (1 + order.get(row.id, 0)),
    } for row in rows[:CANDIDATES_PER_KIND]]


def _cards(db: Session, user: User, query: str) -> list[dict]:
    ranked, _ = hybrid_ids(db, query, "flashcard", limit=PER_KIND * 3)
    if not ranked:
        return []
    own = [d.id for d in db.query(FlashcardDeck.id).filter(
        FlashcardDeck.user_id == user.id, FlashcardDeck.deleted_at.is_(None)).all()]
    if not own:
        return []
    rows = db.query(Flashcard).filter(
        Flashcard.id.in_(ranked), Flashcard.deck_id.in_(own)).all()
    order = {rid: i for i, rid in enumerate(ranked)}
    rows.sort(key=lambda r: order.get(r.id, len(order)))
    return [{
        "kind": "card", "ref": str(row.id), "id": row.id,
        "title": _clean(row.front, 90),
        "text": _clean(row.back, 260),
        "score": 1.0 / (1 + order.get(row.id, 0)),
    } for row in rows[:CANDIDATES_PER_KIND]]


def _apply_curated_boost(db: Session, sources: list[dict]) -> None:
    """Lift a source that an educator tied to another source in this shortlist.

    Two rows an educator already linked, both surfacing for the same question, is
    a stronger signal than either ranking alone — that is what a curated link is
    for. Nothing is stored; the boost lives only in this ordering.
    """
    article_ids = {s["id"] for s in sources if s["kind"] in ("article", "section")}
    question_ids = {s["id"] for s in sources if s["kind"] == "question"}
    if not article_ids or not question_ids:
        return
    pairs = db.query(QuestionArticleLink).filter(
        QuestionArticleLink.article_id.in_(article_ids),
        QuestionArticleLink.question_id.in_(question_ids)).all()
    if not pairs:
        return
    linked_articles = {p.article_id for p in pairs}
    linked_questions = {p.question_id for p in pairs}
    for source in sources:
        if source["kind"] in ("article", "section") and source["id"] in linked_articles:
            source["score"] += CURATED_BOOST
            source["curated"] = True
        elif source["kind"] == "question" and source["id"] in linked_questions:
            source["score"] += CURATED_BOOST
            source["curated"] = True


def _cross_encode(query: str, sources: list[dict]) -> None:
    """Rescore the whole shortlist on one scale, in place, if a reranker answers.

    The scores the finders assign are `1/(1+rank)` *within their own corpus*, so
    the best article, the best section, the best question and the best card all
    score 1.0 and the order between them is whatever the sort happened to do.
    A cross-encoder is the first thing in this pipeline that can compare a
    section against a question, because it reads both against the same query.

    The new score is the position, not the raw relevance: 0.75 from one reranker
    and 0.78 from another mean nothing to `CURATED_BOOST`, whereas `1/(1+rank)`
    keeps a curated pair worth the same few places it has always been worth.

    Silence leaves the scores alone. That is the whole degradation story here:
    without a reranker the shortlist is exactly what it was before there was one.
    """
    if len(sources) < 2:
        return
    scores = rerank_service.rerank(query, [f"{s['title']}. {s['text']}" for s in sources])
    if scores is None:
        return
    ranking = sorted(range(len(sources)), key=lambda i: (-scores[i], i))
    for position, index in enumerate(ranking):
        sources[index]["score"] = 1.0 / (1 + position)


def retrieve(db: Session, user: User, query: str) -> list[dict]:
    """The only things the model will be allowed to cite for this message."""
    query = (query or "").strip()
    if len(query) < 2:
        return []
    sources: list[dict] = []
    for finder in (_sections, _articles, _questions, _cards):
        try:
            sources.extend(finder(db, user, query))
        except Exception:
            # One corpus failing narrows the answer; it does not end it.
            logger.warning("AI Mode retrieval failed for %s", finder.__name__, exc_info=True)
    _cross_encode(query, sources)
    _apply_curated_boost(db, sources)
    sources.sort(key=lambda s: -s["score"])

    # Still no more than `PER_KIND` of anything. A shortlist the cross-encoder
    # filled with fourteen sections of one article would score well and read
    # like a single source quoted fourteen times, and it would leave the chat
    # with no question to send the learner to practise.
    kept: list[dict] = []
    seen: dict[str, int] = {}
    for source in sources:
        if seen.get(source["kind"], 0) >= PER_KIND:
            continue
        seen[source["kind"]] = seen.get(source["kind"], 0) + 1
        kept.append(source)
        if len(kept) >= MAX_SOURCES:
            break
    return kept


def sources_block(sources: list[dict]) -> str:
    """The shortlist, as the model sees it."""
    lines = []
    for source in sources:
        marker = f"[[{source['kind']}:{source['ref']}]]"
        curated = " (an educator linked this to another source here)" if source.get("curated") else ""
        lines.append(f"{marker} {source['title']}{curated}\n{source['text']}")
    return "\n\n".join(lines)


# How close the nearest thing in the library has to be before the answer is
# treated as coming from it. Measured against this corpus with the bodies
# embedded, eight clearly on-topic questions and eight clearly off-topic ones:
#
#   off-topic  0.339 – 0.499   (the French revolution … photosynthesis)
#   on-topic   0.586 – 0.740   (what causes croup … posterior urethral valves)
#
# The gap between them is where these sit. They are not the retrieval floor and
# must not be: `SEMANTIC_FLOOR` decides what is worth showing in a list, where a
# weak hit costs a reader one glance. Here it decides whether an answer claims
# to come from the library, and a wrong claim costs them their trust in every
# other answer.
#
# Worth re-measuring when the corpus changes size or subject; the method and
# the full measurement are in docs/retrieval-thresholds.md. Anything else is
# tuning by feel against numbers nobody wrote down.
#
# Deliberately not the cross-encoder's score, though it is the better judge of a
# pair. The question here is "is there anything in this library about this at
# all", and `top_similarity` answers it by scanning every embedded row in two
# corpora through the vector index. A reranker can only score candidates that
# were shortlisted first, so a reranked closeness cannot tell "the library does
# not cover this" from "retrieval had a bad day", and it would put a network hop
# in the path of a decision about what the answer *claims* — the one place where
# a service being down must not change the output. Order is a preference;
# whether an answer says it came from the library is a promise.
STRONG_MATCH = 0.55
ADJACENT_MATCH = 0.50

ROLE = "You are a study assistant for a pediatrics learning platform.\n"

CITE = (
    "Cite with the exact marker shown, for example [[article:7]] or "
    "[[section:7#abc123]], placed at the end of the sentence it supports. Never "
    "write a URL and never cite a marker that is not listed here.\n\n"
    "Never reveal the answer to a practice question. You may say what a question "
    "is about so the learner can go and attempt it.\n\n"
    "Be brief: a few sentences or a short list.\n\n"
)


def closeness(db: Session, query: str) -> float | None:
    """How close the nearest thing in the library is, or None if unmeasurable."""
    try:
        return top_similarity(db, query)
    except Exception:
        logger.warning("AI Mode could not measure closeness", exc_info=True)
        return None


def answer_mode(similarity: float | None, sources: list[dict]) -> str:
    """Which of the three answers this question gets: sourced, adjacent, or open.

    Decided by a number rather than by asking the model to work out which
    situation it is in. Classification written as prose in a prompt is the part
    that does not work, and it is also the part that makes the prompt long.

    An unmeasurable closeness — no vector database, encoder down — is not a low
    one. Retrieval still found these rows by other means, and discarding them
    because the ruler is missing would silently drop every citation on a
    deployment where semantic search happens to be unavailable.
    """
    if not sources:
        return "open"
    if similarity is None or similarity >= STRONG_MATCH:
        return "sourced"
    return "adjacent" if similarity >= ADJACENT_MATCH else "open"


def build_prompt(sources: list[dict], mode: str = "sourced") -> str:
    if mode == "open" or not sources:
        # Nothing in the library is close, so the honest answer is to say that
        # and then help anyway. Refusing outright was the old behaviour and it
        # reads as a broken assistant rather than a careful one; answering as
        # though the shortlist supported it would be worse still.
        return (
            ROLE +
            "Nothing in this learner's library covers their question.\n\n"
            "Open with one short sentence saying so. Then answer from general "
            "knowledge, briefly and plainly. Do not cite anything: there is "
            "nothing here to cite, and a marker you invent points nowhere."
        )
    if mode == "adjacent":
        # Something related, nothing direct. Naming the gap is the point: a
        # learner told "here is what is closest" can judge the answer, where one
        # handed an adjacent source as though it were the answer cannot.
        return (
            ROLE +
            "Nothing in this learner's library covers their question directly. "
            "The sources below are the closest things in it.\n\n"
            "Open with one short sentence saying that, naming what the closest "
            "material is about. Then answer using those sources where they help, "
            "citing them, and from general knowledge where they do not — saying "
            "which is which.\n\n"
            + CITE +
            f"SOURCES\n\n{sources_block(sources)}"
        )
    return (
        ROLE + "\n"
        "Answer only from the sources below. They are the learner's own library — "
        "if they do not contain the answer, say so rather than filling the gap from "
        "your own knowledge, which the learner cannot check against anything.\n\n"
        + CITE +
        f"SOURCES\n\n{sources_block(sources)}"
    )


def enforce_citations(reply: str, sources: list[dict]) -> tuple[str, list[dict]]:
    """Drop every citation that was not retrieved, and report the ones that stand.

    This is the step that makes a hallucinated citation impossible rather than
    unlikely. The model can write whatever marker it likes; only markers in the
    shortlist survive contact with this function.
    """
    allowed = {f"{s['kind']}:{s['ref']}": s for s in sources}
    used: dict[str, dict] = {}

    def replace(match: re.Match) -> str:
        key = f"{match.group(1)}:{match.group(2)}"
        source = allowed.get(key)
        if source is None:
            return ""   # Invented, or pointing at something this learner may not see.
        used[key] = source
        return match.group(0)

    cleaned = CITATION_RE.sub(replace, reply)
    # Deleting a marker can leave a double space or a space before a full stop.
    cleaned = re.sub(r"[ \t]{2,}", " ", cleaned)
    cleaned = re.sub(r"\s+([.,;:!?])", r"\1", cleaned).strip()

    citations = [{
        "marker": f"[[{s['kind']}:{s['ref']}]]",
        "kind": s["kind"], "id": s["id"],
        "section_id": s.get("section_id"),
        "title": s["title"],
        "curated": bool(s.get("curated")),
    } for s in used.values()]
    return cleaned, citations


#: A session assembled from a conversation. Enough to be worth sitting, few
#: enough that it follows from what was just discussed rather than becoming a
#: general exam on the topic.
PRACTICE_MAX = 20


def practice_ids(db: Session, user: User, citations, question: str) -> list[int]:
    """Which questions a chat turn should send a learner to practise.

    Three sources, in the order they deserve trust. A question the answer
    actually cited is the closest thing to "this is what we were talking
    about". An article it cited stands for a topic, so questions filed under
    that article's category follow. Retrieval on the learner's own words
    catches the rest.

    Everything is put through the bank's own visibility rules on the way out —
    a chat is not a way to reach questions a learner could not otherwise see.
    """
    from app.models.question_category import QuestionCategoryLink

    cited_questions = [int(c["id"]) for c in (citations or []) if c.get("kind") == "question"]
    article_ids = [int(c["id"]) for c in (citations or []) if c.get("kind") == "article"]

    from_articles: list[int] = []
    if article_ids:
        categories = [cid for (cid,) in db.query(Article.category_id).filter(
            Article.id.in_(article_ids), Article.category_id.isnot(None)).all()]
        if categories:
            direct = db.query(Question.id).filter(
                Question.question_category_id.in_(categories)).limit(PRACTICE_MAX * 3).all()
            linked = db.query(QuestionCategoryLink.question_id).filter(
                QuestionCategoryLink.category_id.in_(categories)).limit(PRACTICE_MAX * 3).all()
            from_articles = [qid for (qid,) in [*direct, *linked]]

    retrieved: list[int] = []
    try:
        retrieved, _ = hybrid_ids(db, question or "", "question", limit=PRACTICE_MAX * 2)
    except Exception:
        logger.warning("Practice retrieval failed", exc_info=True)

    ordered: list[int] = []
    for group in (cited_questions, from_articles, retrieved):
        for qid in group:
            if qid not in ordered:
                ordered.append(qid)
    if not ordered:
        return []

    allowed = bank_query(db, user).filter(Question.id.in_(ordered))
    scope = exam_scope_predicate(db, user)
    if scope is not None:
        allowed = allowed.filter(scope)
    visible = {row.id for row in allowed.all()}
    return [qid for qid in ordered if qid in visible][:PRACTICE_MAX]