diff --git a/docs/configuration/qa.md b/docs/configuration/qa.md index f8f76f7f..ec527a1e 100644 --- a/docs/configuration/qa.md +++ b/docs/configuration/qa.md @@ -52,10 +52,12 @@ analysis: temperature: 0.0 # Default: 0.0 (deterministic for code generation) code_timeout: 60.0 # Max seconds for code execution max_output_chars: 50000 # Truncate output after this many chars + max_executions: 15 # Max execute_code calls per question ``` - **model**: LLM configuration (see [Providers](providers.md#model-settings)). When unset, falls back to `qa.model`. - **code_timeout**: Maximum seconds for each code execution (default: 60) - **max_output_chars**: Truncate code output after this many characters (default: 50000) +- **max_executions**: Maximum `execute_code` calls per question before the skill is told to answer from what it has (default: 15) See [Analysis skill](../skills/analysis.md) for usage details. diff --git a/docs/skills/analysis.md b/docs/skills/analysis.md index 2012207a..740678ac 100644 --- a/docs/skills/analysis.md +++ b/docs/skills/analysis.md @@ -37,6 +37,7 @@ The interpreter is [pydantic-monty](https://github.com/pydantic/monty), isolated - **Limited imports.** Only `json`, `re`, `math`, `pathlib`. - **Execution timeout** (default 60s, configurable via `analysis.code_timeout`). - **Output truncation** (default 50000 chars, configurable via `analysis.max_output_chars`). +- **Execution budget** (default 15 calls, configurable via `analysis.max_executions`). Past the budget, `execute_code` returns a notice telling the skill to answer from what it has instead of running more code. Variables persist between `execute_code` calls within one invocation, so the agent can build state step by step. A fresh sandbox is built per `client.analyze` call. @@ -190,6 +191,7 @@ analysis: name: claude-sonnet-4-20250514 code_timeout: 60.0 # Max seconds per code execution max_output_chars: 50000 # Truncate output after this many chars + max_executions: 15 # Max execute_code calls per question ``` When `analysis.model` is unset, the skill falls back to `qa.model`. diff --git a/haiku_rag_slim/haiku/rag/config/models.py b/haiku_rag_slim/haiku/rag/config/models.py index a30c4524..3f7a513b 100644 --- a/haiku_rag_slim/haiku/rag/config/models.py +++ b/haiku_rag_slim/haiku/rag/config/models.py @@ -113,6 +113,7 @@ class AnalysisConfig(BaseModel): model: ModelConfig | None = None code_timeout: float = 60.0 max_output_chars: int = 50_000 + max_executions: int = 15 class PictureDescriptionConfig(BaseModel): diff --git a/haiku_rag_slim/haiku/rag/skills/_deps.py b/haiku_rag_slim/haiku/rag/skills/_deps.py index 530d47b4..f113b6ab 100644 --- a/haiku_rag_slim/haiku/rag/skills/_deps.py +++ b/haiku_rag_slim/haiku/rag/skills/_deps.py @@ -27,6 +27,7 @@ class RAGRunDeps(SkillRunDeps): @dataclass class AnalysisRunDeps(RAGRunDeps): sandbox: "Sandbox | None" = None + execute_count: int = 0 def _reset_invocation_state(state: Any) -> None: @@ -73,6 +74,7 @@ def make_analysis_lifespan(db_path: Path, config: AppConfig): async with HaikuRAG(db_path, config=config, read_only=True) as rag: deps.rag = rag deps.search_count = 0 + deps.execute_count = 0 sandbox = Sandbox( db_path=db_path, config=config, diff --git a/haiku_rag_slim/haiku/rag/skills/_tools.py b/haiku_rag_slim/haiku/rag/skills/_tools.py index 7366d794..71132b54 100644 --- a/haiku_rag_slim/haiku/rag/skills/_tools.py +++ b/haiku_rag_slim/haiku/rag/skills/_tools.py @@ -265,6 +265,7 @@ def create_skill_tools( tools["get_document"] = get_document if "execute_code" in tool_names: + max_executions = config.analysis.max_executions async def execute_code(ctx: RunContext[AnalysisRunDeps], code: str) -> str: """Execute Python code in a sandboxed interpreter. @@ -280,6 +281,13 @@ def create_skill_tools( Args: code: Python code to execute. """ + ctx.deps.execute_count += 1 + if ctx.deps.execute_count > max_executions: + return ( + "Code-execution limit reached. Give your final answer now " + "from what you already have; do not call execute_code again." + ) + assert ctx.deps is not None and ctx.deps.sandbox is not None, ( "AnalysisRunDeps.sandbox is not set — skill lifespan must run before execute_code." ) diff --git a/haiku_rag_slim/haiku/rag/skills/analysis.py b/haiku_rag_slim/haiku/rag/skills/analysis.py index bcdc674d..39aa127c 100644 --- a/haiku_rag_slim/haiku/rag/skills/analysis.py +++ b/haiku_rag_slim/haiku/rag/skills/analysis.py @@ -97,4 +97,5 @@ def create_skill( state_namespace=STATE_NAMESPACE, deps_type=AnalysisRunDeps, lifespan=make_analysis_lifespan(db_path, config), + request_limit=30, ) diff --git a/tests/skills/test_analysis.py b/tests/skills/test_analysis.py index 1c28c959..99dac76d 100644 --- a/tests/skills/test_analysis.py +++ b/tests/skills/test_analysis.py @@ -55,6 +55,14 @@ class TestAnalysisSkillCreation: assert skill.metadata.description assert skill.instructions + def test_create_skill_sets_request_limit_backstop( + self, test_app_config, temp_db_path + ): + from haiku.rag.skills.analysis import create_skill + + skill = create_skill(config=test_app_config, db_path=temp_db_path) + assert skill.request_limit == 30 + def test_create_skill_has_expected_tools(self, test_app_config, temp_db_path): from haiku.rag.skills.analysis import create_skill @@ -148,6 +156,23 @@ class TestExecuteCodeTool: assert "ZeroDivisionError" in result assert state.executions[0].success is False + async def test_execute_code_rate_limited(self, rag_db, sandbox_factory): + from haiku.rag.skills.analysis import create_skill + + config = AppConfig() + config.analysis.max_executions = 2 + skill = create_skill(db_path=rag_db, config=config) + execute_code = _get_tool(skill, "execute_code") + state = AnalysisState() + ctx = _make_ctx(state, sandbox=sandbox_factory()) + + await execute_code(ctx, code="print('first')") + await execute_code(ctx, code="print('second')") + result = await execute_code(ctx, code="print('third')") + assert "limit reached" in result.lower() + assert ctx.deps.execute_count == 3 + assert len(state.executions) == 2 + async def test_execute_code_applies_document_filter(self, rag_db, sandbox_factory): from haiku.rag.skills.analysis import create_skill @@ -295,6 +320,7 @@ class TestAnalysisLifespan: assert deps.rag is not None assert deps.rag.is_read_only assert deps.search_count == 0 + assert deps.execute_count == 0 assert isinstance(deps.sandbox, Sandbox) docs = await deps.rag.list_documents() assert len(docs) == 2 @@ -311,6 +337,16 @@ class TestAnalysisLifespan: assert deps.sandbox is not None assert deps.sandbox._context.filter == "title = 'AI Overview'" + async def test_lifespan_resets_counts_per_invocation(self, rag_db): + from haiku.rag.skills._deps import AnalysisRunDeps, make_analysis_lifespan + + config = AppConfig() + lifespan = make_analysis_lifespan(rag_db, config) + deps = AnalysisRunDeps(search_count=7, execute_count=42) + async with lifespan(deps): + assert deps.search_count == 0 + assert deps.execute_count == 0 + async def test_skill_has_lifespan_and_deps_type( self, test_app_config, temp_db_path ):