danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

fix per-question capability limits and tool isolation

Browse source
This commit is contained in:
Yiorgis Gozadinos 2026-07-19 13:15:17 +03:00
parent 43c17a6777
commit 77585ef26a
No known key found for this signature in database
12 changed files with 176 additions and 53 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
CHANGELOG.md
View file

@ -30,6 +30,7 @@
- `hotpotqa` evaluation dataset.
- Native deferred Pydantic AI `RAGCapability` and `AnalysisCapability` implementations under `haiku.rag.capabilities`, with namespaced host state and lazy per-run database and sandbox resources.
- Prior-turn RAG and analysis tool results are compacted before model requests while current-turn evidence remains intact.
- Per-question capability request limits force a final answer from gathered evidence by removing only the exhausted capability's tools; unrelated agent and capability tools remain available.
### Changed

2
docs/capabilities/analysis.md
View file

@ -4,6 +4,8 @@
It is deferred by default, keeping its substantial instructions and tool schemas out of context until the model chooses to load it.
The default request limit is 30 model requests per question. Override it with `create_capability(request_limit=...)`, or set `request_limit=None` to disable it. At the limit, only analysis tools are removed and the model gets one more turn to answer from gathered evidence. Other agent and capability tools remain available, and the budget resets for every agent run.
## Tools
| Tool | Purpose |

4
docs/capabilities/rag.md
View file

@ -24,7 +24,9 @@ result = await agent.run("What safety equipment does the manual require?")
print(result.output)
```
`create_capability` accepts `db_path`, `config`, and `defer_loading`. Set `defer_loading=False` for a dedicated RAG agent where routing is unnecessary.
`create_capability` accepts `db_path`, `config`, `defer_loading`, and `request_limit`. Set `defer_loading=False` for a dedicated RAG agent where routing is unnecessary. The default request limit is 20 model requests per question; set `request_limit=None` to disable it.
When the limit is reached, only the RAG capability's tools are removed. The model gets one more turn to answer from evidence already gathered, while unrelated agent and capability tools remain available. A new agent run starts a fresh limit, so multi-turn chat does not consume one shared budget.
## State

16
evaluations/evaluations/capability_runner.py
View file

@ -5,7 +5,6 @@ from typing import Any, Protocol, cast
from pydantic_ai import Agent
from pydantic_ai.models import Model
from pydantic_ai.usage import UsageLimits
from haiku.rag.capabilities import RAGCapabilityBase
from haiku.rag.config.models import AppConfig
@ -61,6 +60,8 @@ async def run_capability_question(
config=config,
defer_loading=False,
)
if request_limit is not None:
capability.request_limit = request_limit
state = capability.state_type()
typed = cast(_RagLikeState, state)
if document_filter is not None:
@ -72,18 +73,7 @@ async def run_capability_question(
deps_type=_EvalDeps,
capabilities=[capability],
)
effective_request_limit = (
request_limit if request_limit is not None else capability.default_request_limit
)
agent_result = await agent.run(
question,
deps=deps,
usage_limits=(
UsageLimits(request_limit=effective_request_limit)
if effective_request_limit is not None
else None
),
)
agent_result = await agent.run(question, deps=deps)
state = capability.state_type.model_validate(deps.state[capability.state_namespace])
typed = cast(_RagLikeState, state)

10
evaluations/tests/test_capability_runner.py
View file

@ -41,13 +41,18 @@ async def test_runs_analysis_capability_without_legacy_capability_layer(tmp_path
@pytest.mark.parametrize(("override", "expected"), [(None, 30), (5, 5)])
async def test_analysis_capability_applies_request_limit(tmp_path, override, expected):
capability = create_analysis(
db_path=tmp_path / "rag.lancedb",
config=AppConfig(),
defer_loading=False,
)
with patch(
"evaluations.capability_runner.Agent.run", new_callable=AsyncMock
) as run:
run.return_value = SimpleNamespace(output="done")
await run_capability_question(
create_analysis,
lambda **_kwargs: capability,
tmp_path / "rag.lancedb",
AppConfig(),
"hello",
@ -55,4 +60,5 @@ async def test_analysis_capability_applies_request_limit(tmp_path, override, exp
request_limit=override,
)
assert run.call_args.kwargs["usage_limits"].request_limit == expected
assert capability.request_limit == expected
assert "usage_limits" not in run.call_args.kwargs

43
haiku_rag_slim/haiku/rag/capabilities/_base.py
View file

@ -8,6 +8,7 @@ from pydantic import BaseModel
from pydantic_ai import ModelRetry, RunContext
from pydantic_ai.capabilities import AbstractCapability
from pydantic_ai.messages import (
InstructionPart,
ModelMessage,
ModelRequest,
ToolReturn,
@ -16,6 +17,7 @@ from pydantic_ai.messages import (
)
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
@ -87,13 +89,14 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
instruction_text: str
model: ModelConfig
tool_names: frozenset[str]
default_request_limit: int | None = None
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)
async def for_run(self, ctx: RunContext[Any]) -> "RAGCapabilityBase[StateT]":
outer = getattr(ctx.deps, "state", None)
@ -109,6 +112,7 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
rag_lock=asyncio.Lock(),
resource_lock=asyncio.Lock(),
search_count=0,
request_count=0,
)
run_capability._sync_state()
return run_capability
@ -124,8 +128,45 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
request_context.messages = _compact_old_tool_returns(
request_context.messages, self.tool_names
)
if self._request_limit_reached:
current_request = request_context.messages[-1]
if isinstance(current_request, ModelRequest):
instruction = (
f"The {self.state_namespace} capability has reached its request "
"limit. Its tools are no longer available. Give the best answer "
"possible using the evidence already gathered."
)
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),
],
)
else:
self.request_count += 1
return request_context
async def prepare_tools(
self,
ctx: RunContext[Any],
tool_defs: list[ToolDefinition],
) -> list[ToolDefinition]:
"""Remove only this capability's tools after its per-question limit."""
if not self._request_limit_reached:
return tool_defs
return [tool for tool in tool_defs if tool.capability_id != self.id]
@property
def _request_limit_reached(self) -> bool:
return (
self.request_limit is not None and self.request_count >= self.request_limit
)
async def after_run(
self, ctx: RunContext[Any], *, result: AgentRunResult[Any]
) -> AgentRunResult[Any]:

7
haiku_rag_slim/haiku/rag/capabilities/analysis.py
View file

@ -128,6 +128,7 @@ def create_capability(
config: AppConfig | None = None,
*,
defer_loading: bool = True,
request_limit: int | None = 30,
) -> AnalysisCapability:
"""Create a native Pydantic AI analysis capability."""
if config is None:
@ -142,10 +143,12 @@ def create_capability(
instruction_text=instructions(),
model=config.analysis.model or config.qa.model,
tool_names=_TOOL_NAMES,
default_request_limit=30,
request_limit=request_limit,
id=_CAPABILITY_ID,
description=(
"Analyze the haiku.rag corpus with search and sandboxed Python code."
"Analyze the haiku.rag corpus with search and sandboxed Python code. "
"Use for counting, aggregation, statistics, data traversal, comparison "
"across documents, and other tasks best solved by writing Python code."
),
defer_loading=defer_loading,
)

2
haiku_rag_slim/haiku/rag/capabilities/rag.py
View file

@ -70,6 +70,7 @@ def create_capability(
config: AppConfig | None = None,
*,
defer_loading: bool = True,
request_limit: int | None = 20,
) -> RAGCapability:
"""Create a native Pydantic AI RAG capability."""
if config is None:
@ -84,6 +85,7 @@ def create_capability(
instruction_text=instructions(),
model=config.qa.model,
tool_names=_TOOL_NAMES,
request_limit=request_limit,
id=_CAPABILITY_ID,
description=(
"Search the haiku.rag knowledge base and cite evidence for grounded answers."

10
haiku_rag_slim/haiku/rag/chat/app.py
View file

@ -17,7 +17,6 @@ from pydantic_ai.messages import (
TextPartDelta,
)
from pydantic_ai.run import AgentRunResultEvent
from pydantic_ai.usage import UsageLimits
from textual.app import App, SystemCommand
from textual.binding import Binding
from textual.widgets import Footer, Header, Input
@ -86,14 +85,6 @@ class ChatApp(App):
super().__init__()
self.db_path = db_path
self._capabilities = capabilities
request_limits = [
capability.default_request_limit
for capability in capabilities
if capability.default_request_limit is not None
]
self._usage_limits = (
UsageLimits(request_limit=min(request_limits)) if request_limits else None
)
self.read_only = read_only
self._model = model
self.client: HaikuRAG | None = None
@ -200,7 +191,6 @@ class ChatApp(App):
message_history=self._messages,
conversation_id=self._conversation_id,
deps=deps,
usage_limits=self._usage_limits,
) as stream:
async for event in stream:
if isinstance(event, PartStartEvent) and isinstance(

7
haiku_rag_slim/haiku/rag/client/agents.py
View file

@ -2,7 +2,6 @@ from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any
from pydantic_ai import Agent
from pydantic_ai.usage import UsageLimits
if TYPE_CHECKING:
from haiku.rag.client import HaikuRAG
@ -103,11 +102,7 @@ async def analyze(
deps_type=_AgentDeps,
capabilities=[capability],
)
result = await agent.run(
question,
deps=deps,
usage_limits=UsageLimits(request_limit=capability.default_request_limit),
)
result = await agent.run(question, deps=deps)
state = AnalysisState.model_validate(deps.state["analysis"])
citations = [
state.citation_index[cid]

109
tests/capabilities/test_capabilities.py
View file

@ -49,6 +49,7 @@ def test_rag_capability_api(temp_db_path):
assert toolset.sequential is True
assert capability.state_type is RAGState
assert capability.state_namespace == "rag"
assert capability.request_limit == 20
def test_analysis_capability_api(temp_db_path):
@ -66,6 +67,7 @@ def test_analysis_capability_api(temp_db_path):
assert toolset.max_retries == 3
assert toolset.sequential is True
assert capability.state_type is AnalysisState
assert capability.request_limit == 30
def test_domain_preamble_is_added_to_capability_instructions(temp_db_path):
@ -117,6 +119,113 @@ async def test_capability_instructions_are_injected_once(
assert seen_instructions[0].count(heading) == 1
@pytest.mark.asyncio
async def test_request_limit_removes_only_exhausted_capability_tools_per_run(
temp_db_path,
):
calls = 0
seen_tools = []
seen_instructions = []
def model_function(_messages, info):
nonlocal calls
calls += 1
seen_tools.append({tool.name for tool in info.function_tools})
seen_instructions.append(info.instructions or "")
if calls % 2 == 1:
return ModelResponse(parts=[ToolCallPart("host_tool", {})])
return ModelResponse(parts=[TextPart("best available answer")])
def host_tool(_ctx: RunContext[Deps]) -> str:
"""Return host-owned context."""
return "host context"
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,
request_limit=1,
)
agent = Agent(
FunctionModel(model_function),
deps_type=Deps,
tools=[host_tool],
capabilities=[rag, analysis],
)
first = await agent.run("Analyze this", deps=Deps())
second = await agent.run("Analyze another question", deps=Deps())
assert first.output == "best available answer"
assert second.output == "best available answer"
analysis_tools = {
"analysis_search",
"analysis_execute_code",
"analysis_cite",
}
for initial, exhausted in ((0, 1), (2, 3)):
assert analysis_tools <= seen_tools[initial]
assert analysis_tools.isdisjoint(seen_tools[exhausted])
assert {"host_tool", "rag_search", "rag_cite"} <= seen_tools[exhausted]
assert (
"analysis capability has reached its request limit"
in (seen_instructions[exhausted])
)
@pytest.mark.asyncio
async def test_deferred_request_limit_starts_after_capability_load(temp_db_path):
seen_tools = []
seen_instructions = []
def model_function(_messages, info):
seen_tools.append({tool.name for tool in info.function_tools})
seen_instructions.append(info.instructions or "")
if len(seen_tools) == 1:
return ModelResponse(
parts=[
ToolCallPart(
"load_capability",
{"id": "haiku-rag-analysis"},
)
]
)
if len(seen_tools) == 2:
return ModelResponse(parts=[ToolCallPart("host_tool", {})])
return ModelResponse(parts=[TextPart("best available answer")])
def host_tool(_ctx: RunContext[Deps]) -> str:
"""Return host-owned context."""
return "host context"
agent = Agent(
FunctionModel(model_function),
deps_type=Deps,
tools=[host_tool],
capabilities=[
create_analysis(
db_path=temp_db_path,
config=AppConfig(),
request_limit=1,
)
],
)
result = await agent.run("Analyze this", deps=Deps())
assert result.output == "best available answer"
assert "load_capability" in seen_tools[0]
assert "analysis_search" in seen_tools[1]
assert "analysis_search" not in seen_tools[2]
assert "host_tool" in seen_tools[2]
assert "analysis capability has reached its request limit" in seen_instructions[2]
@pytest.mark.asyncio
async def test_capability_isolated_per_run_and_round_trips_state(temp_db_path):
capability = create_rag(db_path=temp_db_path, config=AppConfig())

18
tests/chat/test_chat_app.py
View file

@ -4,7 +4,6 @@ from unittest.mock import AsyncMock, patch
import pytest
from typer.testing import CliRunner
from haiku.rag.capabilities.analysis import create_capability as create_analysis
from haiku.rag.capabilities.rag import RAGState, create_capability
from haiku.rag.cli import _cli as cli
@ -83,23 +82,6 @@ def _make_app_with_state(db_path: Path, mock_client: AsyncMock | None = None):
), mock_client
def test_chat_uses_capability_request_limit(temp_db_path: Path):
"""Test chat applies the strictest request guard from its capabilities."""
from haiku.rag.chat.app import ChatApp
app = ChatApp(
db_path=temp_db_path,
capabilities=[
create_capability(db_path=temp_db_path),
create_analysis(db_path=temp_db_path),
],
read_only=True,
)
assert app._usage_limits is not None
assert app._usage_limits.request_limit == 30
@pytest.mark.asyncio
async def test_chat_app_has_required_widgets(temp_db_path: Path):
"""Test that ChatApp has the required widgets: ChatHistory, Input."""