8eef0b3734
A run carrying no prompt is continuing a question rather than asking one: pydantic-ai resumes that way for deferred tool results, interruptions and suspended responses. len(ctx.messages) then counts the live question's own messages, so its search result was replaced by the earlier-question notice and the model was asked to answer with the evidence removed. Reproduced: resuming with an in-flight history left a notice where the only evidence was. Switch compaction off for the whole run when ctx.prompt is None. The absence of a prompt is the signal rather than the message layout — the resume shapes differ from each other, and deriving the boundary from layout is what broke this to begin with.
464 lines
19 KiB
Python
464 lines
19 KiB
Python
import asyncio
|
|
import os
|
|
from dataclasses import dataclass, field, replace
|
|
from difflib import get_close_matches
|
|
from pathlib import Path
|
|
from typing import Any, cast
|
|
|
|
from pydantic import BaseModel
|
|
from pydantic_ai import ModelRetry, RunContext, ToolFailed
|
|
from pydantic_ai.capabilities import AbstractCapability, WrapModelRequestHandler
|
|
from pydantic_ai.messages import (
|
|
InstructionPart,
|
|
ModelMessage,
|
|
ModelRequest,
|
|
ModelResponse,
|
|
ToolCallPart,
|
|
ToolReturn,
|
|
ToolReturnPart,
|
|
)
|
|
from pydantic_ai.models import ModelRequestContext
|
|
from pydantic_ai.run import AgentRunResult
|
|
from pydantic_ai.tools import ToolDefinition
|
|
from pydantic_ai.toolsets import AgentToolset
|
|
|
|
from haiku.rag.capabilities._tools import CodeExecutionEntry, search_corpus
|
|
from haiku.rag.client import HaikuRAG
|
|
from haiku.rag.config.models import AppConfig
|
|
from haiku.rag.store.models.chunk import SearchResult
|
|
from haiku.rag.store.models.citation import Citation, resolve_citations
|
|
from haiku.rag.tools.search import build_image_content_from_results
|
|
|
|
CITATION_GRACE_REQUESTS = 2
|
|
"""Requests calling this capability's tools that its cite tool outlives the rest by.
|
|
|
|
A loop guard, not a budget: cite consumes no retry budget and raises nothing, so
|
|
left available forever a stuck model calls it until the agent's own request limit
|
|
raises ``UsageLimitExceeded`` and the question returns no answer at all. Only
|
|
engagement can loop, which is why other capabilities' turns do not spend it.
|
|
"""
|
|
|
|
CHUNK_ID_MATCH_CUTOFF = 0.75
|
|
"""Similarity a cited chunk id needs to be treated as a corrupted known id.
|
|
|
|
Calibration knob. Two unrelated UUID4s reach about 0.5, while dropping or
|
|
duplicating a character or a whole group stays above 0.75, so the gap is wide.
|
|
"""
|
|
|
|
|
|
def _nearest_known_id(chunk_id: str, known_ids: list[str]) -> str:
|
|
"""Recover a chunk id the model damaged while transcribing it.
|
|
|
|
Models copying opaque UUIDs drop and duplicate characters and whole
|
|
hyphen-separated groups. Candidates are limited to ids the run actually
|
|
retrieved, so a wrong match needs both a near miss and a same-run neighbour.
|
|
Ids that match nothing are returned unchanged for the caller to report.
|
|
"""
|
|
if not known_ids or chunk_id in known_ids:
|
|
return chunk_id
|
|
match = get_close_matches(chunk_id, known_ids, n=1, cutoff=CHUNK_ID_MATCH_CUTOFF)
|
|
return match[0] if match else chunk_id
|
|
|
|
|
|
def resolve_db_path(db_path: Path | None, config: AppConfig) -> Path:
|
|
if db_path is not None:
|
|
return db_path
|
|
if env_db := os.environ.get("HAIKU_RAG_DB"):
|
|
return Path(env_db).expanduser()
|
|
return config.storage.data_dir / "haiku.rag.lancedb"
|
|
|
|
|
|
def _clear_invocation_state(state: BaseModel) -> None:
|
|
for field_name in ("citations", "searches", "executions"):
|
|
value = getattr(state, field_name, None)
|
|
if hasattr(value, "clear"):
|
|
value.clear()
|
|
|
|
|
|
PRIOR_TURN_NOTICE = (
|
|
"[Evidence retrieved for an earlier question, no longer shown. It does not "
|
|
"count as cited for the current question.]"
|
|
)
|
|
|
|
|
|
def _compact_old_tool_returns(
|
|
messages: list[ModelMessage],
|
|
tool_names: frozenset[str],
|
|
*,
|
|
turn_start: int,
|
|
) -> list[ModelMessage]:
|
|
"""Remove bulky earlier-question evidence while retaining the current one.
|
|
|
|
Tool call and return parts remain paired; only the old return payload is
|
|
replaced.
|
|
|
|
Page images attached to a replaced return are deliberately left in place.
|
|
Dropping them alongside their text is tempting — they are the bulk, and
|
|
they accumulate — but a follow-up about a figure already shown ("what
|
|
colour is that box?") carries no terms that could retrieve it again, so
|
|
removing the image turns an answerable question into a refusal.
|
|
|
|
``turn_start`` is how many messages existed when the current question
|
|
arrived, so everything below it belongs to an earlier one. The run reports
|
|
it rather than this function deriving it from message shape: a
|
|
``UserPromptPart`` mid-question is as likely to be page images on a tool
|
|
return, or a notice a capability injected, and reading either as the next
|
|
question strips evidence the model is still answering from.
|
|
"""
|
|
if turn_start <= 0:
|
|
return messages
|
|
|
|
compacted = list(messages)
|
|
for index, message in enumerate(messages[:turn_start]):
|
|
if not isinstance(message, ModelRequest):
|
|
continue
|
|
parts = [
|
|
replace(part, content=PRIOR_TURN_NOTICE)
|
|
if isinstance(part, ToolReturnPart) and part.tool_name in tool_names
|
|
else part
|
|
for part in message.parts
|
|
]
|
|
if parts != message.parts:
|
|
compacted[index] = replace(message, parts=parts)
|
|
return compacted
|
|
|
|
|
|
def _called_own_tool(messages: list[ModelMessage], tool_names: frozenset[str]) -> bool:
|
|
"""Whether the model's most recent response called one of these tools."""
|
|
for message in reversed(messages):
|
|
if isinstance(message, ModelResponse):
|
|
return any(
|
|
isinstance(part, ToolCallPart) and part.tool_name in tool_names
|
|
for part in message.parts
|
|
)
|
|
return False
|
|
|
|
|
|
@dataclass
|
|
class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
|
|
db_path: Path
|
|
config: AppConfig
|
|
state_type: type[StateT]
|
|
state_namespace: str
|
|
instruction_text: str
|
|
vision: bool
|
|
tool_names: frozenset[str]
|
|
request_limit: int | None = None
|
|
state: StateT | None = field(default=None, repr=False)
|
|
outer_state: dict[str, Any] | None = field(default=None, repr=False)
|
|
rag: HaikuRAG | None = field(default=None, repr=False)
|
|
rag_lock: asyncio.Lock = field(default_factory=asyncio.Lock, repr=False)
|
|
resource_lock: asyncio.Lock = field(default_factory=asyncio.Lock, repr=False)
|
|
search_count: int = field(default=0, repr=False)
|
|
request_count: int = field(default=0, repr=False)
|
|
grace_requests_used: int = field(default=0, repr=False)
|
|
turn_start: int = field(default=0, repr=False)
|
|
|
|
async def for_run(self, ctx: RunContext[Any]) -> "RAGCapabilityBase[StateT]":
|
|
"""Start a run's own copy, and decide what counts as an earlier question.
|
|
|
|
A run carrying no prompt is continuing a question rather than asking one:
|
|
pydantic-ai resumes that way for deferred tool results, interruptions and
|
|
suspended responses. ``len(ctx.messages)`` would then count the live
|
|
question's own messages and hand the model a notice where its search
|
|
result should be, so compaction is switched off for the whole run
|
|
(``turn_start=0``). The absence of a prompt is the signal — the three
|
|
resume shapes differ in message layout, and reading the layout is what
|
|
broke this in the first place.
|
|
"""
|
|
outer = getattr(ctx.deps, "state", None)
|
|
outer_state = outer if isinstance(outer, dict) else None
|
|
raw_state = outer_state.get(self.state_namespace) if outer_state else None
|
|
state = self.state_type.model_validate(raw_state or {})
|
|
_clear_invocation_state(state)
|
|
run_capability = replace(
|
|
self,
|
|
state=state,
|
|
outer_state=outer_state,
|
|
rag=None,
|
|
rag_lock=asyncio.Lock(),
|
|
resource_lock=asyncio.Lock(),
|
|
search_count=0,
|
|
request_count=0,
|
|
grace_requests_used=0,
|
|
turn_start=0 if ctx.prompt is None else len(ctx.messages),
|
|
)
|
|
run_capability._sync_state()
|
|
return run_capability
|
|
|
|
def get_instructions(self) -> str:
|
|
if self.config.prompts.domain_preamble:
|
|
return f"{self.config.prompts.domain_preamble}\n\n{self.instruction_text}"
|
|
return self.instruction_text
|
|
|
|
async def wrap_model_request(
|
|
self,
|
|
ctx: RunContext[Any],
|
|
*,
|
|
request_context: ModelRequestContext,
|
|
handler: WrapModelRequestHandler,
|
|
) -> ModelResponse:
|
|
"""Trim earlier-question evidence off the wire only.
|
|
|
|
Deliberately not ``before_model_request``: that hook's result is
|
|
assigned back onto the run's message history, so trimming there would
|
|
destroy the host's record of what was retrieved.
|
|
"""
|
|
request_context.messages = _compact_old_tool_returns(
|
|
request_context.messages,
|
|
self.tool_names - {self._cite_tool_name},
|
|
turn_start=self.turn_start,
|
|
)
|
|
return await handler(request_context)
|
|
|
|
async def before_model_request(
|
|
self, ctx: RunContext[Any], request_context: ModelRequestContext
|
|
) -> ModelRequestContext:
|
|
if instruction := self._budget_notice():
|
|
current_request = request_context.messages[-1]
|
|
if isinstance(current_request, ModelRequest):
|
|
current_request.instructions = "\n\n".join(
|
|
part for part in (current_request.instructions, instruction) if part
|
|
)
|
|
parameters = request_context.model_request_parameters
|
|
request_context.model_request_parameters = replace(
|
|
parameters,
|
|
instruction_parts=[
|
|
*(parameters.instruction_parts or []),
|
|
InstructionPart(content=instruction, dynamic=True),
|
|
],
|
|
)
|
|
if self._request_limit_reached and _called_own_tool(
|
|
request_context.messages, self.tool_names
|
|
):
|
|
self.grace_requests_used += 1
|
|
self.request_count += 1
|
|
return request_context
|
|
|
|
def _budget_notice(self) -> str | None:
|
|
"""Tell the model which of this capability's budgets just ran out.
|
|
|
|
Never names a tool ``prepare_tools`` has already withdrawn: pointing the
|
|
model at a tool that is gone costs it the agent's unknown-tool retry
|
|
budget and can abort the run.
|
|
"""
|
|
if self._citation_grace_expired:
|
|
return (
|
|
f"The {self.state_namespace} capability's tools are no longer "
|
|
"available. Give the best answer possible using the evidence "
|
|
"already gathered."
|
|
)
|
|
if self._request_limit_reached:
|
|
return (
|
|
f"The {self.state_namespace} capability has reached its request "
|
|
f"limit. Only {self._cite_tool_name} remains among its tools: "
|
|
"register the chunk_ids supporting your answer, then answer from "
|
|
"the evidence already gathered."
|
|
)
|
|
if spent := self._spent_tool_names():
|
|
names = ", ".join(sorted(spent))
|
|
if remaining := sorted(self._evidence_tool_names() - spent):
|
|
return (
|
|
f"The {self.state_namespace} capability has spent its budget "
|
|
f"for {names}; further calls to them fail. Gather any further "
|
|
f"evidence with {', '.join(remaining)}, or call "
|
|
f"{self._cite_tool_name} with the chunk_ids you have and "
|
|
"answer."
|
|
)
|
|
return (
|
|
f"The {self.state_namespace} capability has spent its budget for "
|
|
f"{names}; further calls to them fail. Answer from the evidence "
|
|
f"already gathered and call {self._cite_tool_name} with the "
|
|
"chunk_ids supporting it."
|
|
)
|
|
return None
|
|
|
|
async def prepare_tools(
|
|
self,
|
|
ctx: RunContext[Any],
|
|
tool_defs: list[ToolDefinition],
|
|
) -> list[ToolDefinition]:
|
|
"""Remove this capability's tools past its limit, cite tool last.
|
|
|
|
Tools whose own budget is spent stay declared on purpose. Removing one
|
|
makes a model that calls it anyway hit ``Unknown tool name``, charged
|
|
against the agent's unknown-tool retry budget, which kills the run after
|
|
two attempts. A spent tool that keeps failing only wastes requests.
|
|
"""
|
|
if self._citation_grace_expired:
|
|
return [tool for tool in tool_defs if tool.capability_id != self.id]
|
|
if not self._request_limit_reached:
|
|
return tool_defs
|
|
return [
|
|
tool
|
|
for tool in tool_defs
|
|
if tool.capability_id != self.id or tool.name == self._cite_tool_name
|
|
]
|
|
|
|
def _evidence_tool_names(self) -> set[str]:
|
|
"""Tools that can bring new evidence into the run."""
|
|
return {f"{self.state_namespace}_search"}
|
|
|
|
def _spent_tool_names(self) -> set[str]:
|
|
"""This capability's tools whose own budget is exhausted."""
|
|
if self.search_count >= self._max_searches:
|
|
return {f"{self.state_namespace}_search"}
|
|
return set()
|
|
|
|
@property
|
|
def _cite_tool_name(self) -> str:
|
|
return f"{self.state_namespace}_cite"
|
|
|
|
@property
|
|
def _max_searches(self) -> int:
|
|
return self.config.qa.max_searches
|
|
|
|
@property
|
|
def _request_limit_reached(self) -> bool:
|
|
return (
|
|
self.request_limit is not None and self.request_count >= self.request_limit
|
|
)
|
|
|
|
@property
|
|
def _citation_grace_expired(self) -> bool:
|
|
# No `request_limit is None` guard: the counter only advances under
|
|
# `_request_limit_reached`, which already requires a limit.
|
|
return self.grace_requests_used >= CITATION_GRACE_REQUESTS
|
|
|
|
async def after_run(
|
|
self, ctx: RunContext[Any], *, result: AgentRunResult[Any]
|
|
) -> AgentRunResult[Any]:
|
|
await self._close()
|
|
return result
|
|
|
|
async def on_run_error(
|
|
self, ctx: RunContext[Any], *, error: BaseException
|
|
) -> AgentRunResult[Any]:
|
|
await self._close()
|
|
raise error
|
|
|
|
async def _ensure_rag(self) -> HaikuRAG:
|
|
if self.rag is None:
|
|
async with self.resource_lock:
|
|
if self.rag is None:
|
|
rag = HaikuRAG(self.db_path, config=self.config, read_only=True)
|
|
await rag.__aenter__()
|
|
self.rag = rag
|
|
return self.rag
|
|
|
|
async def _close(self) -> None:
|
|
if self.rag is not None:
|
|
await self.rag.__aexit__(None, None, None)
|
|
self.rag = None
|
|
|
|
def _sync_state(self) -> None:
|
|
if self.outer_state is not None and self.state is not None:
|
|
self.outer_state[self.state_namespace] = self.state.model_dump(mode="json")
|
|
|
|
async def _with_state(self, operation: Any) -> Any:
|
|
"""Execute an operation and copy its state back to the host dependencies.
|
|
|
|
A failing tool still syncs, so evidence it gathered before the failure
|
|
reaches the host.
|
|
"""
|
|
try:
|
|
return await operation
|
|
finally:
|
|
self._sync_state()
|
|
|
|
async def _search(self, query: str, limit: int | None) -> str | ToolReturn:
|
|
assert self.state is not None
|
|
self.search_count += 1
|
|
if self.search_count > self._max_searches:
|
|
raise ToolFailed(
|
|
"Search limit reached. Answer the question using "
|
|
"the results you already have."
|
|
)
|
|
async with self.rag_lock:
|
|
formatted, results = await search_corpus(
|
|
await self._ensure_rag(),
|
|
query,
|
|
limit=limit,
|
|
document_filter=getattr(self.state, "document_filter", None),
|
|
)
|
|
state = cast(Any, self.state)
|
|
state.searches[query] = results
|
|
if self.vision and (parts := build_image_content_from_results(results)):
|
|
return ToolReturn(return_value=formatted, content=parts)
|
|
return formatted
|
|
|
|
async def _cite(self, chunk_ids: list[str]) -> str:
|
|
assert self.state is not None
|
|
if not chunk_ids:
|
|
raise ModelRetry(
|
|
"No citations registered: chunk_ids was empty. Pass the chunk_ids "
|
|
"you want to cite, copied verbatim from search results."
|
|
)
|
|
|
|
all_results: list[SearchResult] = []
|
|
state = cast(Any, self.state)
|
|
for results in state.searches.values():
|
|
all_results.extend(results)
|
|
known_ids = [result.chunk_id for result in all_results if result.chunk_id]
|
|
requested = [_nearest_known_id(cid.strip("[]"), known_ids) for cid in chunk_ids]
|
|
citations = resolve_citations(requested, all_results)
|
|
resolved = {citation.chunk_id for citation in citations}
|
|
missing = [cid for cid in requested if cid not in resolved]
|
|
|
|
if missing:
|
|
async with self.rag_lock:
|
|
rag = await self._ensure_rag()
|
|
synthetic: list[SearchResult] = []
|
|
documents: dict[str, Any] = {}
|
|
for chunk_id in missing:
|
|
chunk = await rag.get_chunk_by_id(chunk_id)
|
|
if chunk is None or not chunk.document_id:
|
|
continue
|
|
document = documents.get(chunk.document_id)
|
|
if chunk.document_id not in documents:
|
|
document = await rag.get_document_by_id(chunk.document_id)
|
|
documents[chunk.document_id] = document
|
|
chunk.document_uri = document.uri if document else None
|
|
chunk.document_title = document.title if document else None
|
|
chunk.document_meta = document.metadata if document else {}
|
|
synthetic.append(SearchResult.from_chunk(chunk, score=1.0))
|
|
citations.extend(resolve_citations(missing, synthetic))
|
|
|
|
if not citations:
|
|
raise ModelRetry(
|
|
f"None of the supplied chunk_ids {list(chunk_ids)} could be resolved. "
|
|
"Copy chunk_ids verbatim from search results."
|
|
)
|
|
self._register_citations(citations)
|
|
resolved = {citation.chunk_id for citation in citations}
|
|
unresolved = [cid for cid in missing if cid not in resolved]
|
|
if unresolved:
|
|
return (
|
|
f"Registered {len(citations)} citation(s); "
|
|
f"ignored {len(unresolved)} unresolvable id(s): "
|
|
f"{unresolved}. Copy chunk_ids verbatim from search "
|
|
"results and cite again."
|
|
)
|
|
return f"Registered {len(citations)} citation(s)."
|
|
|
|
def _register_citations(self, citations: list[Citation]) -> None:
|
|
assert self.state is not None
|
|
state = cast(Any, self.state)
|
|
next_index = len(state.citation_index) + 1
|
|
for citation in citations:
|
|
if citation.chunk_id not in state.citation_index:
|
|
citation.index = next_index
|
|
next_index += 1
|
|
state.citation_index[citation.chunk_id] = citation
|
|
if citation.chunk_id not in state.citations:
|
|
state.citations.append(citation.chunk_id)
|
|
|
|
def get_toolset(self) -> AgentToolset[Any] | None:
|
|
raise NotImplementedError
|
|
|
|
|
|
__all__ = [
|
|
"CodeExecutionEntry",
|
|
"RAGCapabilityBase",
|
|
"resolve_db_path",
|
|
]
|