3ec72363bc
Resumption was inferred from the transcript, and no shape says it. A missing prompt is how UI adapters ask their first question as much as how pydantic-ai resumes one, so every AG-UI host failed on its first message. Reading the tail instead moved the error rather than fixing it: a settled structured answer ends with an output tool's return, indistinguishable from results delivered to a question still in progress, so following questions inherited the first one's identity. `CapabilityEvidenceRecord.in_progress` is now the authority. `begin_question` sets it, `after_run` clears it unless the run is only pausing for deferred work, and a run that raised never reaches `after_run`, which is what leaves an interrupted question claimable. The history is consulted only to catch a host that dropped the state of a question the model is unmistakably still owed.
227 lines
9.2 KiB
Python
227 lines
9.2 KiB
Python
from collections.abc import Iterable
|
|
from typing import Literal
|
|
|
|
from pydantic import BaseModel, Field
|
|
|
|
CitationStatus = Literal["missing", "grounded", "ungrounded"]
|
|
|
|
|
|
class EvidenceRef(BaseModel):
|
|
"""One piece of evidence, identified by its owner as well as its chunk.
|
|
|
|
A chunk id alone is not an identity: the same id can be reported by more than
|
|
one capability, and ownership is what tells compaction whose output it may
|
|
touch.
|
|
"""
|
|
|
|
capability: str
|
|
chunk_id: str
|
|
|
|
|
|
class EvidenceOccurrence(BaseModel):
|
|
"""Which questions retrieved a piece of evidence, and which cited it."""
|
|
|
|
capability: str
|
|
chunk_id: str
|
|
retrieved_in_questions: list[int] = Field(default_factory=list)
|
|
cited_in_questions: list[int] = Field(default_factory=list)
|
|
|
|
|
|
class CitationDeclaration(BaseModel):
|
|
"""What a question declared as its grounding, and when.
|
|
|
|
Bound to a question *and* an epoch: the epoch outlives a question, so matching
|
|
it alone would let a question that gathered no evidence inherit the previous
|
|
declaration and read as compliant having declared nothing.
|
|
"""
|
|
|
|
question: int
|
|
epoch: int
|
|
refs: list[EvidenceRef] = Field(default_factory=list)
|
|
|
|
|
|
class CapabilityEvidenceRecord(BaseModel):
|
|
"""What one evidence capability wrote, in its own state namespace.
|
|
|
|
Holds no content: ``Citation`` in ``citation_index`` is the canonical record
|
|
and already persists content, document id and picture refs. This is the index
|
|
over it that compaction needs and the transcript cannot provide.
|
|
|
|
Single-writer by construction. A record shared between capabilities would be
|
|
overwritten by whichever of them synced its state last; merging happens in the
|
|
transient views built by ``citation_status`` and the optional capabilities.
|
|
|
|
``in_progress`` is whether the question this record names is still being
|
|
answered. It is the only authority on that: a transcript ending in a tool
|
|
return is a settled structured answer and a question awaiting its model
|
|
equally, so the shape of the history cannot decide it.
|
|
|
|
``question`` is the number of messages that existed when the question arrived,
|
|
and ``epoch`` the number when an outcome occurred. Both are derived from the
|
|
conversation rather than counted locally, so every participant computes the
|
|
same values without sharing a counter. ``question`` is unset until a run
|
|
establishes it, so a record a host merely created is distinguishable from one
|
|
that has been through a question.
|
|
"""
|
|
|
|
occurrences: dict[str, EvidenceOccurrence] = Field(default_factory=dict)
|
|
question: int | None = None
|
|
in_progress: bool = False
|
|
latest_evidence_epoch: int = 0
|
|
declaration: CitationDeclaration | None = None
|
|
|
|
def _reject_regression(self, count: int, what: str) -> None:
|
|
"""Refuse a message count below one already recorded in this question.
|
|
|
|
Identities and epochs are both message counts, and every comparison
|
|
between them assumes the conversation only grows while a question is
|
|
being answered. One capability truncating or reordering the history
|
|
breaks that, and each way of recording it has to refuse the same way: an
|
|
unchecked evidence outcome freezes every later declaration as stale, while
|
|
an unchecked declaration replaces a newer one with an older one and
|
|
revives the answer it grounded.
|
|
"""
|
|
recorded = max(
|
|
self.question or 0,
|
|
self.latest_evidence_epoch,
|
|
self.declaration.epoch if self.declaration else 0,
|
|
)
|
|
if count < recorded:
|
|
raise ValueError(
|
|
f"{what} at message count {count} is behind {recorded}, which is "
|
|
"already recorded: message history must be append-only for "
|
|
"question identities and epochs to hold."
|
|
)
|
|
|
|
def begin_question(self, identity: int) -> None:
|
|
"""Take the identity of a question that has just arrived.
|
|
|
|
The evidence epoch and the declaration describe the question that has
|
|
just ended, and outlive their meaning the moment the next one begins:
|
|
kept, they let the last question's citations ground this one, and hold a
|
|
horizon this question's own declarations cannot pass. Clearing them
|
|
confines those comparisons to one question, so a host whose stored
|
|
history shifted between two of them is answered rather than refused.
|
|
|
|
Identities themselves must still separate one question from the next.
|
|
Occurrences outlive the question that recorded them and carry the
|
|
identities that cited them, which compaction groups and orders the capsule
|
|
by: a reused identity merges two questions into one group, and a lower one
|
|
renders later evidence as though it were cited earlier.
|
|
"""
|
|
if self.question is not None and identity <= self.question:
|
|
raise ValueError(
|
|
f"A question at message count {identity} is not past question "
|
|
f"{self.question}, which is already answered: identities separate "
|
|
"one question from the next and are compared as recency."
|
|
)
|
|
self.question = identity
|
|
self.in_progress = True
|
|
self.latest_evidence_epoch = 0
|
|
self.declaration = None
|
|
|
|
def end_question(self) -> None:
|
|
"""Mark the question answered, so the next run knows it is a new one."""
|
|
self.in_progress = False
|
|
|
|
def note_evidence(self, epoch: int) -> None:
|
|
"""Record that the model has seen an evidence outcome.
|
|
|
|
Called for anything an answer could rest on, including a search that
|
|
returned nothing and a failed execution that still printed output — a
|
|
fruitless search grounds a refusal. Not called for a failure that yields
|
|
no evidence at all, such as an exhausted budget.
|
|
"""
|
|
self._reject_regression(epoch, "Evidence")
|
|
self.latest_evidence_epoch = epoch
|
|
|
|
def declare(
|
|
self,
|
|
refs: list[EvidenceRef],
|
|
*,
|
|
epoch: int,
|
|
retrieved_now: set[str] | None = None,
|
|
) -> None:
|
|
"""Record validated citations for the current question.
|
|
|
|
Citing again cannot narrow what a question already declared: calls merge
|
|
while no evidence outcome has followed the standing declaration, whatever
|
|
epoch they arrive at, so an empty second thought leaves a grounded question
|
|
grounded. Only genuinely newer evidence starts a declaration afresh, since
|
|
what the model saw in between may be what it is now citing.
|
|
"""
|
|
if self.question is None:
|
|
raise ValueError(
|
|
"Citations cannot be declared before a run establishes the "
|
|
"question identity."
|
|
)
|
|
self._reject_regression(epoch, "A declaration")
|
|
current = self.declaration
|
|
if (
|
|
current is not None
|
|
and current.question == self.question
|
|
and self.latest_evidence_epoch <= current.epoch
|
|
):
|
|
known = {(ref.capability, ref.chunk_id) for ref in current.refs}
|
|
current.refs.extend(
|
|
ref for ref in refs if (ref.capability, ref.chunk_id) not in known
|
|
)
|
|
current.epoch = max(current.epoch, epoch)
|
|
else:
|
|
self.declaration = CitationDeclaration(
|
|
question=self.question, epoch=epoch, refs=list(refs)
|
|
)
|
|
|
|
for ref in refs:
|
|
occurrence = self.occurrences.setdefault(
|
|
ref.chunk_id,
|
|
EvidenceOccurrence(capability=ref.capability, chunk_id=ref.chunk_id),
|
|
)
|
|
if self.question not in occurrence.cited_in_questions:
|
|
occurrence.cited_in_questions.append(self.question)
|
|
if (
|
|
retrieved_now
|
|
and ref.chunk_id in retrieved_now
|
|
and self.question not in occurrence.retrieved_in_questions
|
|
):
|
|
occurrence.retrieved_in_questions.append(self.question)
|
|
|
|
|
|
def citation_status(
|
|
records: Iterable[CapabilityEvidenceRecord], *, question: int
|
|
) -> CitationStatus:
|
|
"""Derived, never stored, so refs and status cannot contradict.
|
|
|
|
A declaration is current only for the question it was made in, and only if it
|
|
followed the newest evidence outcome of *every* capability: a question where
|
|
one capability cited and another then searched without citing is not grounded.
|
|
Strictly later, since a citation made in the same request as an outcome cannot
|
|
have read it.
|
|
|
|
A grounding *violation* is not one of these: that is an enforcement outcome
|
|
recorded by the policy capability, not something the model declared.
|
|
"""
|
|
records = list(records)
|
|
horizon = max((record.latest_evidence_epoch for record in records), default=0)
|
|
current = [
|
|
record.declaration
|
|
for record in records
|
|
if record.declaration is not None
|
|
and record.declaration.question == question
|
|
and record.declaration.epoch > horizon
|
|
]
|
|
if not current:
|
|
return "missing"
|
|
return (
|
|
"grounded" if any(declaration.refs for declaration in current) else "ungrounded"
|
|
)
|
|
|
|
|
|
__all__ = [
|
|
"CapabilityEvidenceRecord",
|
|
"CitationDeclaration",
|
|
"CitationStatus",
|
|
"EvidenceOccurrence",
|
|
"EvidenceRef",
|
|
"citation_status",
|
|
]
|