0f38417c60
`CitationPolicyCapability` makes the single enforcement decision, whatever mix of evidence capabilities is registered: two of them must not each demand a citation for one answer. It decides in `after_model_request`, when a response carries no tool calls and the question can still be redirected. An explicitly ungrounded answer is a declaration and is left alone; a question that gathered no evidence is left alone too, read from the ledger rather than from a searches dict that a new question clears. When the cite tool is already withdrawn the question is recorded in `CitationPolicyState.violations` instead of pointing the model at a tool that is gone. Registering it is the only switch. `DiscoveredEvidence` and discovery move to `capabilities.evidence` so both optional capabilities share them, and `cite_available` joins `evidence_tool_names` as public for the same reason. Measured on Qwen3.6-35B over two arms of 37 questions, 29 of them unanswerable from the corpus: explicit ungrounded declarations rose from 23 to 26, grounded answers to unanswerable questions fell from 4 to 2, answerable questions stayed at 8 of 8, and one redirect fired in the whole arm. Reading every answer found no invented grounding.
282 lines
9.4 KiB
Python
282 lines
9.4 KiB
Python
from dataclasses import dataclass, field
|
|
from typing import Any, cast
|
|
from unittest.mock import patch
|
|
|
|
import pytest
|
|
from pydantic_ai import Agent
|
|
from pydantic_ai.exceptions import UserError
|
|
from pydantic_ai.messages import ModelResponse, TextPart, ToolCallPart
|
|
from pydantic_ai.models.function import FunctionModel
|
|
|
|
from haiku.rag.capabilities.analysis import create_capability as create_analysis
|
|
from haiku.rag.capabilities.policy import (
|
|
REDIRECT_HINT,
|
|
CitationPolicyState,
|
|
)
|
|
from haiku.rag.capabilities.policy import (
|
|
create_capability as create_policy,
|
|
)
|
|
from haiku.rag.capabilities.rag import RAGCapability
|
|
from haiku.rag.capabilities.rag import create_capability as create_rag
|
|
from haiku.rag.config.models import AppConfig
|
|
from haiku.rag.store.models.chunk import SearchResult
|
|
|
|
|
|
@dataclass
|
|
class Deps:
|
|
state: dict[str, Any] = field(default_factory=dict)
|
|
|
|
|
|
async def stub_search(self, query: str, _limit: int | None) -> str:
|
|
cast(Any, self.state).searches[query] = [
|
|
SearchResult(content="evidence", score=1.0, chunk_id="chunk-1")
|
|
]
|
|
self._note_evidence()
|
|
return "EVIDENCE"
|
|
|
|
|
|
def prompts_of(messages) -> list[str]:
|
|
return [
|
|
str(part.content)
|
|
for message in messages
|
|
for part in message.parts
|
|
if type(part).__name__ == "UserPromptPart"
|
|
]
|
|
|
|
|
|
async def run_with_policy(temp_db_path, responses, *, policy=True, config=None):
|
|
"""Answer one question with the given model responses, policy optional."""
|
|
rag = create_rag(
|
|
db_path=temp_db_path, config=config or AppConfig(), defer_loading=False
|
|
)
|
|
capabilities: list[Any] = [rag]
|
|
if policy:
|
|
capabilities.append(create_policy())
|
|
turns = iter(responses)
|
|
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=capabilities)
|
|
deps = Deps()
|
|
with patch.object(RAGCapability, "_search", stub_search):
|
|
result = await agent.run("what does the supervisor do?", deps=deps)
|
|
return result, deps, sent
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_an_answer_without_a_citation_is_sent_back_once(temp_db_path):
|
|
"""The last response of a question is the last moment to notice."""
|
|
result, deps, sent = await run_with_policy(
|
|
temp_db_path,
|
|
[
|
|
[ToolCallPart("rag_search", {"query": "supervisor"}, "call-1")],
|
|
[TextPart("an answer with no citation")],
|
|
[ToolCallPart("rag_cite", {"chunk_ids": ["chunk-1"]}, "call-2")],
|
|
[TextPart("an answer with no citation")],
|
|
],
|
|
)
|
|
|
|
redirects = [prompt for prompt in prompts_of(sent[-1]) if REDIRECT_HINT in prompt]
|
|
assert len(redirects) == 1
|
|
assert deps.state["rag"]["citations"] == ["chunk-1"]
|
|
assert result.output == "an answer with no citation"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_a_grounded_answer_is_left_alone(temp_db_path):
|
|
_, _, sent = await run_with_policy(
|
|
temp_db_path,
|
|
[
|
|
[ToolCallPart("rag_search", {"query": "supervisor"}, "call-1")],
|
|
[ToolCallPart("rag_cite", {"chunk_ids": ["chunk-1"]}, "call-2")],
|
|
[TextPart("a grounded answer")],
|
|
],
|
|
)
|
|
|
|
assert not [p for p in prompts_of(sent[-1]) if REDIRECT_HINT in p]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_an_explicitly_ungrounded_answer_is_left_alone(temp_db_path):
|
|
"""Citing nothing is a declaration, not an omission."""
|
|
_, deps, sent = await run_with_policy(
|
|
temp_db_path,
|
|
[
|
|
[ToolCallPart("rag_search", {"query": "supervisor"}, "call-1")],
|
|
[ToolCallPart("rag_cite", {"chunk_ids": []}, "call-2")],
|
|
[TextPart("I cannot find this in the knowledge base")],
|
|
],
|
|
)
|
|
|
|
assert not [p for p in prompts_of(sent[-1]) if REDIRECT_HINT in p]
|
|
assert deps.state["citation_policy"]["violations"] == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_a_question_that_gathered_no_evidence_is_left_alone(temp_db_path):
|
|
"""Nothing was retrieved, so there is no grounding to declare."""
|
|
_, _, 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_violation_is_recorded_when_the_cite_tool_is_gone(temp_db_path):
|
|
"""Asking for a withdrawn tool costs the agent's unknown-tool retries."""
|
|
with patch(
|
|
"haiku.rag.capabilities._base.RAGCapabilityBase.cite_available",
|
|
new_callable=lambda: property(lambda self: False),
|
|
):
|
|
_, deps, sent = await run_with_policy(
|
|
temp_db_path,
|
|
[
|
|
[ToolCallPart("rag_search", {"query": "supervisor"}, "call-1")],
|
|
[TextPart("an answer with no citation")],
|
|
],
|
|
)
|
|
|
|
assert not [p for p in prompts_of(sent[-1]) if REDIRECT_HINT in p]
|
|
assert deps.state["citation_policy"]["violations"] == [0]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_without_the_policy_capability_nothing_is_enforced(temp_db_path):
|
|
"""Omission is the switch, so there is no flag to test."""
|
|
_, deps, sent = await run_with_policy(
|
|
temp_db_path,
|
|
[
|
|
[ToolCallPart("rag_search", {"query": "supervisor"}, "call-1")],
|
|
[TextPart("an answer with no citation")],
|
|
],
|
|
policy=False,
|
|
)
|
|
|
|
assert not [p for p in prompts_of(sent[-1]) if REDIRECT_HINT in p]
|
|
assert "citation_policy" not in deps.state
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_one_decision_is_made_with_both_evidence_capabilities(temp_db_path):
|
|
"""Two capabilities must not each demand a citation for one answer."""
|
|
rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False)
|
|
analysis = create_analysis(
|
|
db_path=temp_db_path, config=AppConfig(), defer_loading=False
|
|
)
|
|
turns = iter(
|
|
[
|
|
[ToolCallPart("rag_search", {"query": "supervisor"}, "call-1")],
|
|
[TextPart("an answer with no citation")],
|
|
[TextPart("an answer with no citation")],
|
|
]
|
|
)
|
|
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, analysis, create_policy()],
|
|
)
|
|
with patch.object(RAGCapability, "_search", stub_search):
|
|
await agent.run("what does the supervisor do?", deps=Deps())
|
|
|
|
assert len([p for p in prompts_of(sent[-1]) if REDIRECT_HINT in p]) == 1
|
|
|
|
|
|
def test_two_policy_capabilities_fail_fast(temp_db_path):
|
|
rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False)
|
|
|
|
async def model(_messages, _info): # pragma: no cover - never reached
|
|
return ModelResponse(parts=[TextPart("answer")])
|
|
|
|
with pytest.raises(UserError, match="unique within a run"):
|
|
Agent(
|
|
FunctionModel(model),
|
|
deps_type=Deps,
|
|
capabilities=[rag, create_policy(), create_policy()],
|
|
)
|
|
|
|
|
|
def test_the_policy_state_round_trips():
|
|
state = CitationPolicyState(violations=[4, 12])
|
|
|
|
restored = CitationPolicyState.model_validate(state.model_dump(mode="json"))
|
|
|
|
assert restored.violations == [4, 12]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_a_second_question_can_be_redirected_again(temp_db_path):
|
|
"""The redirect fires once per question, not once per conversation."""
|
|
rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False)
|
|
turns = iter(
|
|
[
|
|
[ToolCallPart("rag_search", {"query": "supervisor"}, "call-1")],
|
|
[TextPart("first, uncited")],
|
|
[TextPart("first, uncited")],
|
|
[ToolCallPart("rag_search", {"query": "again"}, "call-2")],
|
|
[TextPart("second, uncited")],
|
|
[TextPart("second, uncited")],
|
|
]
|
|
)
|
|
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 who supervises them?", deps=deps, message_history=first.all_messages()
|
|
)
|
|
|
|
assert len([p for p in prompts_of(sent[-1]) if REDIRECT_HINT in p]) == 2
|
|
|
|
|
|
@dataclass
|
|
class StatelessDeps:
|
|
"""A host that keeps no capability state, which is allowed."""
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_a_violation_with_nowhere_to_record_it_does_not_fail_the_run(
|
|
temp_db_path,
|
|
):
|
|
rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False)
|
|
turns = iter(
|
|
[
|
|
[ToolCallPart("rag_search", {"query": "supervisor"}, "call-1")],
|
|
[TextPart("an answer with no citation")],
|
|
]
|
|
)
|
|
|
|
async def model(_messages, _info):
|
|
return ModelResponse(parts=next(turns))
|
|
|
|
agent = Agent(
|
|
FunctionModel(model),
|
|
deps_type=StatelessDeps,
|
|
capabilities=[rag, create_policy()],
|
|
)
|
|
|
|
with (
|
|
patch(
|
|
"haiku.rag.capabilities._base.RAGCapabilityBase.cite_available",
|
|
new_callable=lambda: property(lambda self: False),
|
|
),
|
|
patch.object(RAGCapability, "_search", stub_search),
|
|
):
|
|
result = await agent.run("what does the supervisor do?", deps=StatelessDeps())
|
|
|
|
assert result.output == "an answer with no citation"
|