ae2755e88f
Requiring an evidence outcome from the current question exempted the case enforcement exists for: a follow-up about evidence already cited needs no new search, since that evidence is still on the wire — in a capsule when a compactor is registered, in full when not. The condition is now that the conversation has something to declare, either an outcome in this question or evidence it has already cited, which is independent of whether anything compacts. A conversation that has neither is still left alone. Citing again cannot narrow a question at any epoch. Declarations merged only within one epoch, so an empty second thought a request later replaced the refs with nothing and reported a grounded question ungrounded. They merge while no evidence outcome has followed the standing declaration, and only genuinely newer evidence starts one afresh. Whether a question has already been asked to declare is read from the message history rather than remembered on the run instance, which a resumption's `for_run` discarded — the same question was asked twice. Reading the history 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. Violations are recorded once per question for the same reason.
379 lines
13 KiB
Python
379 lines
13 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 (
|
|
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
|