9beafff0bf
The design settled earlier, built as described: retrieval decides what may be cited, and the server enforces it. The model is handed a shortlist of at most fourteen sources from the learner's own library and told to cite them by marker. Afterwards every citation it wrote is checked against that shortlist and anything else is deleted before it is stored or shown. A hallucinated citation is not unlikely here, it is impossible — surviving is not a decision the model gets to make. A URL it invents is not a citation either: only the marker form counts, so a plausible-looking link stays in the prose citing nothing. Retrieval reuses the hybrid search already in place, and each corpus keeps its own visibility rules — the bank predicate and exam scope for questions, the draft rule for articles, deck ownership for cards. A question source carries the stem only: a chat that printed the answer would hand away the practice it exists to prepare you for. Curated links do the job they were built for. A retrieved row an educator tied to another retrieved row is boosted, because two things somebody already linked surfacing for one query is evidence rather than coincidence. Nothing is stored for this; the boost lives only in that ordering, and the answer marks those sources so the reader knows which claim rests on an educator's judgement rather than on a ranking. Citations are stored with the answer as filtered, so reopening a thread shows the links it showed at the time rather than a fresh retrieval that may now rank differently. In the page the markers become numbers and each number opens its source; a section citation deep-links into that section. Two smaller decisions worth naming: a question appears in the thread the moment you send it and is handed back to the input if the answer fails, because typed words are not something to lose on a 502; and someone else's thread returns 404 rather than 403, since whether it exists is not your business either. 182 backend, 206 frontend green — 16 of the backend tests are the citation contract and the retrieval boundary. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01TqXevQJhxFrM7jJg82cgZN
246 lines
10 KiB
Python
246 lines
10 KiB
Python
"""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.
|
||
"""
|
||
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.quiz_builder import bank_query, exam_scope_predicate
|
||
from app.services.search_service import hybrid_ids
|
||
|
||
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
|
||
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[: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[: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,
|
||
"title": 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[: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[: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 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)
|
||
_apply_curated_boost(db, sources)
|
||
sources.sort(key=lambda s: -s["score"])
|
||
return sources[:MAX_SOURCES]
|
||
|
||
|
||
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)
|
||
|
||
|
||
def build_prompt(sources: list[dict]) -> str:
|
||
if not sources:
|
||
return (
|
||
"You are a study assistant for a pediatrics learning platform.\n"
|
||
"Nothing in this learner's library matches their question. Say so plainly "
|
||
"in one or two sentences and suggest what they might search for instead. "
|
||
"Do not answer from your own knowledge, and do not cite anything."
|
||
)
|
||
return (
|
||
"You are a study assistant for a pediatrics learning platform.\n\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 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"
|
||
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
|