a79f348ac7
`_already_asked` matched the redirect's phrasing, so a question that merely contained it read as a redirect we had sent and switched enforcement off for that question. The redirect carries `[haiku.rag/citation-redirect]` and detection matches that, the same way retrieved pictures are identified by a tag rather than by the prose beside them.
414 lines
14 KiB
Python
414 lines
14 KiB
Python
from dataclasses import dataclass, field
|
|
from typing import Any, cast
|
|
from unittest.mock import patch
|
|
|
|
import pytest
|
|
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]
|