diff --git a/CHANGELOG.md b/CHANGELOG.md index 148e1207..79744d48 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/docs/capabilities/analysis.md b/docs/capabilities/analysis.md index 1ab4d410..97430702 100644 --- a/docs/capabilities/analysis.md +++ b/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 | diff --git a/docs/capabilities/rag.md b/docs/capabilities/rag.md index a6d312a4..503bed41 100644 --- a/docs/capabilities/rag.md +++ b/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 diff --git a/evaluations/evaluations/capability_runner.py b/evaluations/evaluations/capability_runner.py index baf044e3..3957d737 100644 --- a/evaluations/evaluations/capability_runner.py +++ b/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 = ( - 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) diff --git a/evaluations/tests/test_capability_runner.py b/evaluations/tests/test_capability_runner.py index c966ccc4..bc310ac4 100644 --- a/evaluations/tests/test_capability_runner.py +++ b/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 diff --git a/haiku_rag_slim/haiku/rag/capabilities/_base.py b/haiku_rag_slim/haiku/rag/capabilities/_base.py index bc65da0a..8f5687ce 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/_base.py +++ b/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]: diff --git a/haiku_rag_slim/haiku/rag/capabilities/analysis.py b/haiku_rag_slim/haiku/rag/capabilities/analysis.py index e645bfca..f25ee0a8 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/analysis.py +++ b/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, ) diff --git a/haiku_rag_slim/haiku/rag/capabilities/rag.py b/haiku_rag_slim/haiku/rag/capabilities/rag.py index 4ef80293..815e95c8 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/rag.py +++ b/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." diff --git a/haiku_rag_slim/haiku/rag/chat/app.py b/haiku_rag_slim/haiku/rag/chat/app.py index a7eaa888..6d5f0ea1 100644 --- a/haiku_rag_slim/haiku/rag/chat/app.py +++ b/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( diff --git a/haiku_rag_slim/haiku/rag/client/agents.py b/haiku_rag_slim/haiku/rag/client/agents.py index a6b22582..948f4819 100644 --- a/haiku_rag_slim/haiku/rag/client/agents.py +++ b/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( - 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] diff --git a/tests/capabilities/test_capabilities.py b/tests/capabilities/test_capabilities.py index 42d9c4a4..16628e51 100644 --- a/tests/capabilities/test_capabilities.py +++ b/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()) diff --git a/tests/chat/test_chat_app.py b/tests/chat/test_chat_app.py index 645d9ad5..eff2e579 100644 --- a/tests/chat/test_chat_app.py +++ b/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."""