cee0824744
The documentation claimed a conversation that has never cited anything is not enforced, which reads as though a first question could answer from fresh evidence without declaring it. Enforcement needs neither condition to hold: no evidence outcome in this question, and nothing cited earlier. The structured-output test accepted a redirect or a recorded violation, so a break in ending detection would still have passed on the backstop alone. The cite tool is available in that scenario, so it asserts the redirect.
497 lines
17 KiB
Python
497 lines
17 KiB
Python
from dataclasses import dataclass, field
|
|
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
|
|
from pydantic_ai.models.function import FunctionModel
|
|
|
|
from haiku.rag.capabilities.analysis import create_capability as create_analysis
|
|
from haiku.rag.capabilities.policy import (
|
|
CITATION_REDIRECT_TAG,
|
|
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"
|
|
|
|
|
|
@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
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_a_user_quoting_the_redirect_does_not_suppress_enforcement(temp_db_path):
|
|
"""Prose is not proof that we asked: a user can write any phrase.
|
|
|
|
Matching the wording let a question that merely mentioned it pass as already
|
|
asked, which silently switches enforcement off.
|
|
"""
|
|
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")],
|
|
]
|
|
)
|
|
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()]
|
|
)
|
|
|
|
with patch.object(RAGCapability, "_search", stub_search):
|
|
await agent.run(
|
|
"Please record what grounded the answer you already gave, in your notes.",
|
|
deps=Deps(),
|
|
)
|
|
|
|
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)
|
|
|
|
# The cite tool is available here, so the redirect is what must happen; the
|
|
# backstop recording a violation would pass an `or` even with detection broken.
|
|
assert [p for p in prompts_of(sent[-1]) if CITATION_REDIRECT_TAG in p]
|
|
|
|
|
|
@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]
|