fix per-question capability limits and tool isolation

CHANGELOG.md

@@ -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
 

docs/capabilities/analysis.md

@@ -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 |

docs/capabilities/rag.md

@@ -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
 

evaluations/evaluations/capability_runner.py

@@ -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 = (
+    agent_result = await agent.run(question, deps=deps)
-        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
-        ),
-    )
     state = capability.state_type.model_validate(deps.state[capability.state_namespace])
     typed = cast(_RagLikeState, state)
 

evaluations/tests/test_capability_runner.py

@@ -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

haiku_rag_slim/haiku/rag/capabilities/_base.py

@@ -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]:

haiku_rag_slim/haiku/rag/capabilities/analysis.py

@@ -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,
     )

haiku_rag_slim/haiku/rag/capabilities/rag.py

@@ -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."

haiku_rag_slim/haiku/rag/chat/app.py

@@ -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(

haiku_rag_slim/haiku/rag/client/agents.py

@@ -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(
+    result = await agent.run(question, deps=deps)
-        question,
-        deps=deps,
-        usage_limits=UsageLimits(request_limit=capability.default_request_limit),
-    )
     state = AnalysisState.model_validate(deps.state["analysis"])
     citations = [
         state.citation_index[cid]

tests/capabilities/test_capabilities.py

@@ -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())

tests/chat/test_chat_app.py

@@ -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."""