diff --git a/haiku_rag_slim/haiku/rag/capabilities/ledger.py b/haiku_rag_slim/haiku/rag/capabilities/ledger.py index 2750bd6c..778ae7ec 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/ledger.py +++ b/haiku_rag_slim/haiku/rag/capabilities/ledger.py @@ -133,10 +133,11 @@ class CapabilityEvidenceRecord(BaseModel): ) -> None: """Record validated citations for the current question. - Repeated calls at the same epoch merge, so citing again cannot narrow what - was already declared: an empty call after a grounded one leaves it - grounded. A call at a later epoch declares afresh, because evidence the - model saw in between may be what it is now citing. + 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( @@ -145,14 +146,16 @@ class CapabilityEvidenceRecord(BaseModel): ) self._reject_regression(epoch, "A declaration") current = self.declaration - if current is not None and (current.question, current.epoch) == ( - self.question, - epoch, + 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) diff --git a/haiku_rag_slim/haiku/rag/capabilities/policy.py b/haiku_rag_slim/haiku/rag/capabilities/policy.py index 662518f3..b497eb21 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/policy.py +++ b/haiku_rag_slim/haiku/rag/capabilities/policy.py @@ -1,10 +1,15 @@ -from dataclasses import dataclass, field, replace +from dataclasses import dataclass from typing import Any from pydantic import BaseModel, Field from pydantic_ai import RunContext from pydantic_ai.capabilities import AbstractCapability -from pydantic_ai.messages import ModelResponse, ToolCallPart +from pydantic_ai.messages import ( + ModelMessage, + ModelResponse, + ToolCallPart, + UserPromptPart, +) from pydantic_ai.models import ModelRequestContext from haiku.rag.capabilities.evidence import ( @@ -58,12 +63,6 @@ class CitationPolicyCapability(AbstractCapability[Any]): would share this capability's id. """ - redirected: set[int] = field(default_factory=set, repr=False) - - async def for_run(self, ctx: RunContext[Any]) -> "CitationPolicyCapability": - """Give the run its own record of what it has already asked for.""" - return replace(self, redirected=set()) - async def after_model_request( self, ctx: RunContext[Any], @@ -82,13 +81,14 @@ class CitationPolicyCapability(AbstractCapability[Any]): return response evidence = discover_evidence(ctx) question = question_in_progress(evidence) - if question in self.redirected or not _gathered_evidence(evidence): + if not _has_evidence_to_declare(evidence): return response records = [found.record for found in evidence] if citation_status(records, question=question) != "missing": return response + if _already_asked(ctx.messages, question): + return response - self.redirected.add(question) if any(found.cite_available for found in evidence): ctx.enqueue(REDIRECT, priority="when_idle") else: @@ -96,12 +96,17 @@ class CitationPolicyCapability(AbstractCapability[Any]): return response def _record_violation(self, ctx: RunContext[Any], question: int) -> None: - """Note a question that could not be asked to cite, the tool being gone.""" + """Note a question that could not be asked to cite, the tool being gone. + + Recorded once per question: a resumption of the same question decides + again, and one question is one outcome. + """ outer = getattr(ctx.deps, "state", None) if not isinstance(outer, dict): return state = CitationPolicyState.model_validate(outer.get(STATE_NAMESPACE) or {}) - state.violations.append(question) + if question not in state.violations: + state.violations.append(question) outer[STATE_NAMESPACE] = state.model_dump(mode="json") async def before_run(self, ctx: RunContext[Any]) -> None: @@ -113,16 +118,42 @@ class CitationPolicyCapability(AbstractCapability[Any]): ) -def _gathered_evidence(evidence: list[DiscoveredEvidence]) -> bool: - """Whether this question produced anything an answer could be grounded on. +def _already_asked(messages: list[ModelMessage], question: int) -> bool: + """Whether this question has already been asked to declare its grounding. - A question with no evidence outcome has nothing to declare — a greeting, or a - conversational aside. Read from the ledger rather than from ``state.searches``, - which a new question clears, so an answer grounded on code execution or on a - document read counts as well. + Read from the history rather than remembered on the instance, which a + resumption's ``for_run`` would forget — the same question would then be asked + twice. It also makes the right call when a redirect was enqueued but the run + ended before it reached the model: nothing is in the history, so it is asked + again, which is what the model needs. + """ + return any( + isinstance(part, UserPromptPart) + and isinstance(part.content, str) + and REDIRECT_HINT in part.content + for message in messages[question:] + for part in message.parts + ) + + +def _has_evidence_to_declare(evidence: list[DiscoveredEvidence]) -> bool: + """Whether anything exists that this answer could have been grounded on. + + Either this question produced an evidence outcome, or the conversation has + already cited something — which stays available to a later answer, in a capsule + if a compactor is registered and in full if not. Requiring a fresh outcome + exempted exactly the follow-up that reuses earlier evidence, which is the case + enforcement exists for. + + A conversation that has neither has nothing to declare: a greeting, an aside. + Read from the ledger rather than from ``state.searches``, which a new question + clears, so an answer grounded on code execution or a document read counts too. """ question = question_in_progress(evidence) - return any(found.record.latest_evidence_epoch > question for found in evidence) + return any( + found.record.latest_evidence_epoch > question or found.record.occurrences + for found in evidence + ) def create_capability() -> CitationPolicyCapability: diff --git a/tests/capabilities/test_citation_policy.py b/tests/capabilities/test_citation_policy.py index 0c43afa6..ef916420 100644 --- a/tests/capabilities/test_citation_policy.py +++ b/tests/capabilities/test_citation_policy.py @@ -3,7 +3,7 @@ from typing import Any, cast from unittest.mock import patch import pytest -from pydantic_ai import Agent +from pydantic_ai import Agent, DeferredToolResults from pydantic_ai.exceptions import UserError from pydantic_ai.messages import ModelResponse, TextPart, ToolCallPart from pydantic_ai.models.function import FunctionModel @@ -280,3 +280,100 @@ async def test_a_violation_with_nowhere_to_record_it_does_not_fail_the_run( result = await agent.run("what does the supervisor do?", deps=StatelessDeps()) assert result.output == "an answer with no citation" + + +@pytest.mark.asyncio +async def test_a_follow_up_answered_from_retained_evidence_is_enforced(temp_db_path): + """The multi-turn case is the one enforcement exists for. + + A follow-up about something already cited needs no new search — the evidence is + still on the wire, whether in a capsule or in full — so requiring a fresh + evidence outcome let exactly those answers through undeclared. + """ + rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False) + turns = iter( + [ + [ToolCallPart("rag_search", {"query": "supervisor"}, "call-1")], + [ToolCallPart("rag_cite", {"chunk_ids": ["chunk-1"]}, "call-2")], + [TextPart("first answer")], + [TextPart("a follow-up answered from what is already here")], + [TextPart("a follow-up answered from what is already here")], + ] + ) + 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, capabilities=[rag, create_policy()] + ) + deps = Deps() + + with patch.object(RAGCapability, "_search", stub_search): + first = await agent.run("what does the supervisor do?", deps=deps) + await agent.run( + "and what colour is the box in it?", + deps=deps, + message_history=first.all_messages(), + ) + + assert [p for p in prompts_of(sent[-1]) if REDIRECT_HINT in p] + + +@pytest.mark.asyncio +async def test_a_conversation_that_never_cited_anything_is_still_left_alone( + temp_db_path, +): + """A greeting has nothing to declare, and no evidence exists to declare from.""" + _, _, sent = await run_with_policy(temp_db_path, [[TextPart("hello back")]]) + + assert not [p for p in prompts_of(sent[-1]) if REDIRECT_HINT in p] + + +@pytest.mark.asyncio +async def test_a_resumed_question_is_not_redirected_twice(temp_db_path): + """Once per question has to mean once, across every run of that question. + + Tracking it on the run instance forgot it at the next `for_run`, so resuming an + interrupted question asked for the citation again. + """ + 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("uncited again")], + [TextPart("uncited a third time")], + [TextPart("uncited a fourth time")], + [TextPart("uncited a fifth time")], + ] + ) + 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, capabilities=[rag, create_policy()] + ) + deps = Deps() + + with patch.object(RAGCapability, "_search", stub_search): + first = await agent.run("what does the supervisor do?", deps=deps) + # The same question again, continued rather than asked anew. + await agent.run( + deps=deps, + message_history=[ + *first.all_messages(), + ModelResponse(parts=[ToolCallPart("external_tool", {}, "call-9")]), + ], + deferred_tool_results=DeferredToolResults( + calls={"call-9": "external result"} + ), + ) + + redirects = [p for p in prompts_of(sent[-1]) if REDIRECT_HINT in p] + assert len(redirects) == 1 diff --git a/tests/capabilities/test_evidence_ledger.py b/tests/capabilities/test_evidence_ledger.py index c9f94353..a495a7a7 100644 --- a/tests/capabilities/test_evidence_ledger.py +++ b/tests/capabilities/test_evidence_ledger.py @@ -240,3 +240,38 @@ def test_a_question_starts_clear_of_the_one_before_it(): assert record.latest_evidence_epoch == 0 assert record.declaration is None assert citation_status([record], question=9) == "missing" + + +def test_an_empty_citation_after_a_grounded_one_cannot_narrow_it(): + """Citing again must not weaken a declaration, at any epoch. + + Merging only within one epoch meant a second thought a request later replaced + the refs with nothing and reported the question ungrounded. + """ + record = CapabilityEvidenceRecord(question=0) + record.declare([rag_ref()], epoch=3) + record.declare([], epoch=5) + + assert record.declaration is not None + assert [ref.chunk_id for ref in record.declaration.refs] == ["c1"] + assert citation_status([record], question=0) == "grounded" + + +def test_a_declaration_after_newer_evidence_starts_afresh(): + """Evidence the model has since seen may be what it is now citing.""" + record = CapabilityEvidenceRecord(question=0) + record.declare([rag_ref("first")], epoch=3) + record.note_evidence(4) + record.declare([rag_ref("second")], epoch=5) + + assert record.declaration is not None + assert [ref.chunk_id for ref in record.declaration.refs] == ["second"] + + +def test_an_empty_citation_after_newer_evidence_is_ungrounded(): + record = CapabilityEvidenceRecord(question=0) + record.declare([rag_ref()], epoch=3) + record.note_evidence(4) + record.declare([], epoch=5) + + assert citation_status([record], question=0) == "ungrounded"