From eb9934a6e59a566d55b99045068979574a0f1b4b Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Wed, 12 Aug 2026 11:46:27 +0300 Subject: [PATCH] Notice the endings a structured answer arrives in MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit An output tool call is a `ToolCallPart` like any other, and treating every tool call as intermediate meant a model could search, skip citing, emit its structured answer and finish with neither a redirect nor a record. A response ends the question when it carries no tool calls, or when one of its calls names an output tool. Some endings are not visible from a single response — a host running `end_strategy="early"` can finish on text beside a function call — so `after_run` is the backstop: it cannot ask the model for anything by then, but it records a question that reached the end of its run undeclared. That also covers a question that was asked once, ignored, and finished anyway, which previously returned early on the redirect marker and went unrecorded. The capability documentation and the changelog said a question that gathered no evidence is left alone. That describes neither the code nor the intent: enforcement applies wherever there is something to declare, which includes a follow-up that reuses evidence cited earlier without searching again. --- CHANGELOG.md | 2 +- docs/capabilities/index.md | 16 +++- .../haiku/rag/capabilities/policy.py | 53 ++++++++++-- tests/capabilities/test_citation_policy.py | 83 +++++++++++++++++++ 4 files changed, 144 insertions(+), 10 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 6491603a..3ac453f5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,7 +3,7 @@ ### Added -- `CitationPolicyCapability` (`haiku.rag.capabilities.policy.create_capability`): registering it requires every answer to declare its grounding. A question that ends undeclared is sent back once to record what grounded the answer already given; when the cite tool is no longer available the question is recorded in `CitationPolicyState.violations` instead. A question that gathered no evidence is left alone. +- `CitationPolicyCapability` (`haiku.rag.capabilities.policy.create_capability`): registering it requires every answer to declare its grounding, in any conversation that has something to declare — this question retrieved evidence, or something was cited earlier. A question that ends undeclared is sent back once to record what grounded the answer already given, and is recorded in `CitationPolicyState.violations` if it finishes undeclared regardless. A conversation that has never cited anything is not enforced. - `haiku.rag.capabilities.evidence.discover_evidence()` and `DiscoveredEvidence`, moved out of `compaction` so both optional capabilities share them. `RAGCapabilityBase.cite_available`. - `EvidenceCompactionCapability` (`haiku.rag.capabilities.compaction.create_capability`): registering it replaces earlier questions' evidence on the model request with the evidence that was cited, grouped by the question that cited it, cited page images re-attached, other earlier evidence returns reduced to a receipt. Requests only; `all_messages()` is untouched. No configuration. - `RAGState.evidence` / `AnalysisState.evidence` (`CapabilityEvidenceRecord`): which evidence a capability retrieved and cited, per question, keyed by message-count question identities and epochs. `haiku.rag.capabilities.ledger.citation_status(records, question=...)` derives `missing` / `grounded` / `ungrounded` across capabilities. diff --git a/docs/capabilities/index.md b/docs/capabilities/index.md index 1c6f9b7f..817eeadb 100644 --- a/docs/capabilities/index.md +++ b/docs/capabilities/index.md @@ -90,10 +90,18 @@ a declaration possible without forcing the model to invent grounding. When a question ends undeclared, the model is asked once to record what grounded the answer it already gave. It is not asked to change the answer. If the cite tool is no -longer available by then, the question is recorded as a violation in -`CitationPolicyState` under `"citation_policy"` instead, since pointing a model at a -tool that is gone costs it retries. A question that gathered no evidence at all — a -greeting, an aside — is left alone. +longer available by then, or the question finishes undeclared anyway, it is recorded as +a violation in `CitationPolicyState` under `"citation_policy"`, since pointing a model +at a tool that is gone costs it retries. + +What gets enforced is every answer in a conversation that has something to declare: +either this question retrieved evidence, or the conversation has already cited +something, which stays available to later answers. So a follow-up about evidence cited +earlier is enforced even though it searched nothing — that case is the reason the +capability exists. It also means that once anything has been cited, later turns are +enforced too, a greeting included; the model satisfies the policy by citing an empty +list, at the cost of one extra request. A conversation that has never cited anything is +not enforced at all. Exactly one policy capability makes the decision, however many evidence capabilities are registered, so two of them cannot each demand a citation for one answer. diff --git a/haiku_rag_slim/haiku/rag/capabilities/policy.py b/haiku_rag_slim/haiku/rag/capabilities/policy.py index 95ca0214..ce559c18 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/policy.py +++ b/haiku_rag_slim/haiku/rag/capabilities/policy.py @@ -11,6 +11,7 @@ from pydantic_ai.messages import ( UserPromptPart, ) from pydantic_ai.models import ModelRequestContext +from pydantic_ai.run import AgentRunResult from haiku.rag.capabilities.evidence import ( DiscoveredEvidence, @@ -80,12 +81,17 @@ class CitationPolicyCapability(AbstractCapability[Any]): ) -> ModelResponse: """Decide once, at the last moment a question can still be redirected. - A response carrying no tool calls ends the question, so there is no later - opportunity. Citing is unconditional, so an undeclared answer is a protocol - breach whether the model answered or refused, and this never has to guess - which it was. + A response that ends the question is the last opportunity: one carrying no + tool calls, or one whose call is an output tool, which is how a structured + answer arrives and which finishes the run just the same. Citing is + unconditional, so an undeclared answer is a protocol breach whether the model + answered or refused, and this never has to guess which it was. + + Endings this cannot see — a host running ``end_strategy="early"`` can finish + on text beside a function call — are caught by ``after_run``, which can still + record the outcome even though it can no longer ask for a citation. """ - if any(isinstance(part, ToolCallPart) for part in response.parts): + if not _ends_the_question(response, request_context): return response evidence = discover_evidence(ctx) question = question_in_progress(evidence) @@ -117,6 +123,25 @@ class CitationPolicyCapability(AbstractCapability[Any]): state.violations.append(question) outer[STATE_NAMESPACE] = state.model_dump(mode="json") + async def after_run( + self, ctx: RunContext[Any], *, result: AgentRunResult[Any] + ) -> AgentRunResult[Any]: + """Record a question that finished undeclared, whatever ended it. + + The backstop for an ending ``after_model_request`` cannot recognise. Nothing + can be asked of the model now, so this only records: a question that reached + the end of its run without a declaration is a violation, and one that was + asked but never answered is the same. + """ + evidence = discover_evidence(ctx) + question = question_in_progress(evidence) + if not _has_evidence_to_declare(evidence): + return result + records = [found.record for found in evidence] + if citation_status(records, question=question) == "missing": + self._record_violation(ctx, question) + return result + async def before_run(self, ctx: RunContext[Any]) -> None: """Publish an empty outcome, so a host can tell "none" from "not running".""" outer = getattr(ctx.deps, "state", None) @@ -126,6 +151,24 @@ class CitationPolicyCapability(AbstractCapability[Any]): ) +def _ends_the_question( + response: ModelResponse, request_context: ModelRequestContext +) -> bool: + """Whether this response finishes the question rather than continuing it. + + An output tool call is a ``ToolCallPart`` like any other, but it carries the + final answer and ends the run, so treating every tool call as intermediate let a + structured answer finish undeclared. + """ + calls = [part for part in response.parts if isinstance(part, ToolCallPart)] + if not calls: + return True + output_tools = { + tool.name for tool in request_context.model_request_parameters.output_tools + } + return any(call.tool_name in output_tools for call in calls) + + def _already_asked(messages: list[ModelMessage], question: int) -> bool: """Whether this question has already been asked to declare its grounding. diff --git a/tests/capabilities/test_citation_policy.py b/tests/capabilities/test_citation_policy.py index d4ff57ee..bcd1a1f9 100644 --- a/tests/capabilities/test_citation_policy.py +++ b/tests/capabilities/test_citation_policy.py @@ -3,6 +3,7 @@ from typing import Any, cast from unittest.mock import patch import pytest +from pydantic import BaseModel from pydantic_ai import Agent, DeferredToolResults from pydantic_ai.exceptions import UserError from pydantic_ai.messages import ModelResponse, TextPart, ToolCallPart @@ -412,3 +413,85 @@ async def test_a_user_quoting_the_redirect_does_not_suppress_enforcement(temp_db ) assert [p for p in prompts_of(sent[-1]) if CITATION_REDIRECT_TAG in p] + + +class Answer(BaseModel): + """A structured output, which the model returns through an output tool.""" + + text: str + + +@pytest.mark.asyncio +async def test_a_structured_output_answer_does_not_escape_enforcement(temp_db_path): + """An output tool call is a `ToolCallPart` too, and it ends the run. + + Treating every tool call as intermediate let a model search, skip citing, emit + its structured answer and finish with neither a redirect nor a violation. + """ + rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False) + turns = iter( + [ + [ToolCallPart("rag_search", {"query": "supervisor"}, "call-1")], + [ + ToolCallPart( + "final_result", {"text": "uncited structured answer"}, "out" + ) + ], + [ + ToolCallPart( + "final_result", {"text": "uncited structured answer"}, "out" + ) + ], + ] + ) + sent: list[list[Any]] = [] + + async def model(messages, _info): + sent.append(list(messages)) + return ModelResponse(parts=next(turns)) + + agent = Agent( + FunctionModel(model), + deps_type=Deps, + output_type=Answer, + capabilities=[rag, create_policy()], + ) + deps = Deps() + + with patch.object(RAGCapability, "_search", stub_search): + await agent.run("what does the supervisor do?", deps=deps) + + redirected = [p for p in prompts_of(sent[-1]) if CITATION_REDIRECT_TAG in p] + violations = deps.state.get("citation_policy", {}).get("violations", []) + assert redirected or violations + + +@pytest.mark.asyncio +async def test_a_question_asked_once_and_still_undeclared_is_recorded(temp_db_path): + """Being asked is not an outcome; the question still ended undeclared. + + Returning early on the redirect marker meant a question that was asked, ignored, + and then finished — with the cite tool possibly gone by that point — was neither + redirected again nor recorded anywhere. + """ + rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False) + turns = iter( + [ + [ToolCallPart("rag_search", {"query": "supervisor"}, "call-1")], + [TextPart("uncited")], + [TextPart("still uncited after being asked")], + ] + ) + + async def model(_messages, _info): + return ModelResponse(parts=next(turns)) + + agent = Agent( + FunctionModel(model), deps_type=Deps, capabilities=[rag, create_policy()] + ) + deps = Deps() + + with patch.object(RAGCapability, "_search", stub_search): + await agent.run("what does the supervisor do?", deps=deps) + + assert deps.state["citation_policy"]["violations"] == [0]