diff --git a/CHANGELOG.md b/CHANGELOG.md index ec08da1..e25a0f6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,16 @@ and this project adheres to [Clever Semantic Versioning](https://www.w3.org/subm ### Added +- **LLM Agent Token-Budget Awareness and Tool Output Pruning (issue #61)** (`llm.py`): Two complementary mechanisms to prevent context-window exhaustion in the multi-turn tool-call loop. + + **Token-budget awareness** (`token_budget_percent` config, default off): tracks estimated token consumption before each LLM invocation. Emits a warning at 75% budget consumption. When the budget ceiling is exceeded, injects a synthesis prompt, permits one final tool-call round, then forces a text-only response. Token estimation uses `chat_model.get_num_tokens()` when available, falling back to a `len(content)//4` heuristic. + + **Tool output pruning** (`allow_tool_output_pruning` config, default `false`): an implicit LLM extraction pass that runs between tool execution and ToolMessage construction, scoped initially to `file_read` calls. The main model can opt-out per call via the `output_prune` meta-argument, or provide a focused search directive via `output_prune_context`. The pruning model responds with `[PRUNE_INFO_START/END]` and `[PRUNE_OUTPUT_START/END]` markers (ยง4.4.9) to separate pruning metadata from extracted content. An optional `pruning_model` config allows using a lighter model for extraction. On failure, the raw tool output is used as-is. + + **ADR:** `docs/adr/ADR-2031-tool-loop-token-budget-and-pruning.md` documents the specification extensions (three new ยง4.4 config fields, two runtime behaviors, one normative marker format). + + **Module:** `src/cleveractors/agents/llm.py`. 1 source file changed, BDD: 21 new scenarios in `features/llm_agent_tool_calling.feature`, Robot: 5 new integration tests in `robot/llm_token_budget.robot`, ASV: 16 benchmarks in `benchmarks/token_budget_benchmark.py`. + - **Tool Schema Descriptions (`llm_tools.py`) (issue #59)** (`cleveractors.agents.llm_tools`): New module providing `normalize_tool_entry()` and `_BUILTIN_TOOL_SCHEMAS` โ€” a registry of OpenAI function-calling parameter schemas for all built-in tools (`echo`, `math`, `json_parse`, `http_request`, `file_read`, `file_write`, `progress_bar`, `shell`, `python_exec`). Each schema includes targeted LLM guidance: `file_read` requires `max_chars` for files >10KB to prevent context overflow; `file_write` clarifies that response text is NOT persisted to disk; `shell` directs the LLM toward `file_read` for reading files; `python_exec` lists available sandbox builtins and instructs the LLM to use `file_read`/`file_write` instead of `open()`. String tool names in agent config (e.g. `"file_read"`) are normalized to full parameter schemas so the LLM knows about parameters like `max_chars`. Config dicts with a `"name"` key are also normalized; already-formatted OpenAI tool dicts pass through unchanged. - **Directory Listing and `max_chars` Truncation in `file_read` Tool (issue #59)** (`tool.py`): When given a directory path, `file_read` now returns a formatted listing with file sizes and type indicators (`๐Ÿ“` for directories, `๐Ÿ“„` for files). New optional `max_chars` integer parameter truncates file content to the specified limit, preventing models from reading multi-hundred-KB files and overflowing the 262K token context window. Output header includes `TRUNCATED to N chars` when truncation applies. Shell command detection rejects paths containing pipe/redirect operators or starting with known shell commands (`grep`, `ls`, `find`, etc.), returning an explicit error redirecting to the shell tool. File-not-found errors now include the parent directory path and instruct the LLM to list that directory for recovery. diff --git a/benchmarks/token_budget_benchmark.py b/benchmarks/token_budget_benchmark.py new file mode 100644 index 0000000..e864a78 --- /dev/null +++ b/benchmarks/token_budget_benchmark.py @@ -0,0 +1,249 @@ +"""ASV benchmarks for token-budget and pruning features. + +Measures performance of token estimation, pruning marker parsing, +schema augmentation, and the full pruning extraction pass under +varying input sizes. +""" + +from __future__ import annotations + +import asyncio +from typing import Any +from unittest.mock import AsyncMock, MagicMock, Mock, patch + +from langchain_core.messages import AIMessage, HumanMessage, SystemMessage + +from cleveractors.agents.llm import LLMAgent +from cleveractors.templates.renderer import TemplateRenderer + + +def _make_agent(**extra_config: Any) -> LLMAgent: + """Create a minimal LLMAgent with mocked dependencies.""" + renderer = Mock(spec=TemplateRenderer) + renderer.render_string.return_value = "mocked system prompt" + config: dict[str, Any] = { + "name": "bench_agent", + "provider": "openai", + "model": "gpt-3.5-turbo", + "api_key": "bench_key", + **extra_config, + } + agent = LLMAgent(name="bench_agent", config=config, template_renderer=renderer) + mock_model = MagicMock() + mock_model.temperature = 0.7 + mock_model.astream = AsyncMock() + agent.chat_model = mock_model + return agent + + +def _messages_with_size(chars: int) -> list[Any]: + """Build a message list with approximately *chars* characters of content.""" + content = "x" * max(1, chars - 70) + return [ + SystemMessage(content="You are a helpful assistant."), + HumanMessage(content=content), + ] + + +# โ”€โ”€ Token Budget Benchmarks โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + + +class TokenBudgetBenchmark: + """Benchmarks for token-budget awareness (ยง4.4.7).""" + + def setup(self) -> None: + self.agent = _make_agent() + + def teardown(self) -> None: + self.agent = None + + def time_estimate_small_messages(self) -> None: + """Estimate tokens for a small (100 char) conversation.""" + msgs = _messages_with_size(100) + self.agent._estimate_token_count(msgs) + + def time_estimate_medium_messages(self) -> None: + """Estimate tokens for a 10K char conversation.""" + msgs = _messages_with_size(10_000) + self.agent._estimate_token_count(msgs) + + def time_estimate_large_messages(self) -> None: + """Estimate tokens for a 50K char conversation.""" + msgs = _messages_with_size(50_000) + self.agent._estimate_token_count(msgs) + + def time_estimate_very_large_messages(self) -> None: + """Estimate tokens for a 200K char conversation near context limit.""" + msgs = _messages_with_size(200_000) + self.agent._estimate_token_count(msgs) + + def time_context_window_lookup(self) -> None: + """Look up the context window for the configured model.""" + self.agent._get_model_context_window() + + +# โ”€โ”€ Pruning Marker Parsing Benchmarks โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + + +class PruningMarkerBenchmark: + """Benchmarks for pruning response marker parsing (ยง4.4.9).""" + + def setup(self) -> None: + self.agent = _make_agent() + content = "x" * 10_000 + + self.response_full = ( + f"{self.agent._PRUNE_INFO_START}\n" + f"Removed 48K chars of non-relevant data\n" + f"{self.agent._PRUNE_INFO_END}\n" + f"{self.agent._PRUNE_OUTPUT_START}\n" + f"{content}\n" + f"{self.agent._PRUNE_OUTPUT_END}" + ) + self.response_output_only = ( + f"{self.agent._PRUNE_OUTPUT_START}\n{content}\n" + f"{self.agent._PRUNE_OUTPUT_END}" + ) + self.response_no_markers = content + + huge = "y" * 50_000 + self.response_huge = ( + f"{self.agent._PRUNE_INFO_START}\nRemoved data\n" + f"{self.agent._PRUNE_INFO_END}\n" + f"{self.agent._PRUNE_OUTPUT_START}\n" + f"{huge}\n" + f"{self.agent._PRUNE_OUTPUT_END}" + ) + + def teardown(self) -> None: + self.agent = None + + def time_parse_full_markers(self) -> None: + """Parse a response with both info and output markers (10K content).""" + self.agent._parse_pruning_response(self.response_full) + + def time_parse_output_only(self) -> None: + """Parse a response with only output markers (10K content).""" + self.agent._parse_pruning_response(self.response_output_only) + + def time_parse_no_markers(self) -> None: + """Parse a response with no markers (falls back to raw text).""" + self.agent._parse_pruning_response(self.response_no_markers) + + def time_parse_huge_output(self) -> None: + """Parse a response with 50K chars of extracted content.""" + self.agent._parse_pruning_response(self.response_huge) + + +# โ”€โ”€ Schema Augmentation Benchmarks โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + + +class SchemaAugmentationBenchmark: + """Benchmarks for tool-schema augmentation (ยง4.4.8 step 1).""" + + def setup(self) -> None: + self.agent = _make_agent() + self.single_tool = [ + { + "type": "function", + "function": { + "name": "echo", + "parameters": { + "type": "object", + "properties": { + "text": {"type": "string"}, + }, + }, + }, + }, + ] + self.many_tools = [ + { + "type": "function", + "function": { + "name": f"tool_{i}", + "parameters": { + "type": "object", + "properties": { + f"param_{j}": {"type": "string"} for j in range(5) + }, + }, + }, + } + for i in range(20) + ] + + def teardown(self) -> None: + self.agent = None + + def time_augment_single_tool(self) -> None: + """Augment a single tool schema.""" + self.agent._augment_tool_schemas_for_pruning(self.single_tool) + + def time_augment_many_tools(self) -> None: + """Augment 20 tool schemas.""" + self.agent._augment_tool_schemas_for_pruning(self.many_tools) + + +# โ”€โ”€ Full Pruning Pass Benchmarks โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + + +class PruningPassBenchmark: + """Benchmarks for the full pruning extraction pass (ยง4.4.8).""" + + def setup(self) -> None: + self.agent = _make_agent(allow_tool_output_pruning=True) + + mock_prune_model = MagicMock() + mock_prune_model.temperature = 0.0 + + self.extracted = "Extracted: key line 42, key line 43" + + async def _mock_prune_ainvoke(messages, **kw): + resp = MagicMock(spec=AIMessage) + resp.content = ( + f"[PRUNE_INFO_START]\nRemoved non-relevant data\n" + f"[PRUNE_INFO_END]\n" + f"[PRUNE_OUTPUT_START]\n{self.extracted}\n" + f"[PRUNE_OUTPUT_END]" + ) + return resp + + mock_prune_model.ainvoke = _mock_prune_ainvoke + self._prune_patch = patch( + "cleveractors.agents.llm.build_chat_model", + return_value=mock_prune_model, + ) + self._prune_patch.start() + + self.small_output = "line 1\nline 2\n" + self.medium_output = "line " * 500 + self.large_output = "log data\n" * 5_000 # ~50K chars + + self.task_ctx = [ + SystemMessage(content="You are a helpful assistant."), + HumanMessage(content="Find error lines in the output."), + ] + + def teardown(self) -> None: + self._prune_patch.stop() + self._prune_patch = None + self.agent = None + + def time_prune_small_output(self) -> None: + """Run the full pruning pass on a small (~16 char) tool output.""" + asyncio.run( + self.agent._run_pruning_pass("file_read", self.small_output, self.task_ctx) + ) + + def time_prune_medium_output(self) -> None: + """Run the full pruning pass on a medium (~2500 char) tool output.""" + asyncio.run( + self.agent._run_pruning_pass("file_read", self.medium_output, self.task_ctx) + ) + + def time_prune_large_output(self) -> None: + """Run the full pruning pass on a large (~50K char) tool output.""" + asyncio.run( + self.agent._run_pruning_pass("file_read", self.large_output, self.task_ctx) + ) diff --git a/docs/adr/ADR-2031-tool-loop-token-budget-and-pruning.md b/docs/adr/ADR-2031-tool-loop-token-budget-and-pruning.md new file mode 100644 index 0000000..3a92180 --- /dev/null +++ b/docs/adr/ADR-2031-tool-loop-token-budget-and-pruning.md @@ -0,0 +1,393 @@ +# ADR-2031: LLM Agent Token Budget and Tool Output Pruning โ€” Specification Extensions + +**Status:** accepted + +**Date:** 2026-06-23 + +**Author:** Luis Mendes (CoreRasurae) + +**Issue:** #61 โ€” Add token-budget awareness and tool output pruning to the multi-turn tool-call loop + +--- + +## Context + +We propose two complementary mechanisms to address context-window +exhaustion in the multi-turn tool-call loop (implemented in #59/#60): + +1. **Token-budget awareness**: track estimated token consumption, warn + when approaching context-window limits, and gracefully terminate + when the budget is exhausted. +2. **Tool output pruning**: an implicit LLM extraction pass that filters + tool output to only task-relevant content before the output enters + the main conversation context, preserving AIMessage reasoning + (reasoning, tool decisions) intact. + +The current specification (ยง4.4) defines eleven LLM agent configuration +fields. The proposed mechanisms introduce four new configuration fields, +two new runtime behaviors, and one new normative marker format (ยง4.4.9) +that are not currently modelled in the normative specification. + +The Actor Configuration Standard ยง1.3 permits compliant implementations +to "provide additional agent types, node types, operators, conditions, +or routing match types beyond those defined here, provided they do not +conflict with the names defined in this standard." ยง4.4.2 defines +default models; new fields must not redefine any existing `MUST`/`SHALL` +requirement. + +This ADR documents the specification extensions required to support +token-budget awareness and tool output pruning in a standards-compliant +manner. + +--- + +## Decision + +### D-1: New LLM Agent Configuration Field (`token_budget_percent`) + +**What:** One new optional configuration field added to the LLM agent +schema in ยง4.4. + +| Field | Type | Required | Default | Description | +|-------|------|----------|---------|-------------| +| `token_budget_percent` | number | No | `0.85` | Fraction of the model's context window reserved for the conversation. When the estimated token count of accumulated messages exceeds this fraction of the model's advertised context window, a warning is emitted at 75% budget consumption, and the budget-exhausted flow triggers at the defined `token_budget_percent` threshold to reserve the remaining context space for the final LLM reply. | + +**Spec relationship:** ยง4.4 currently defines eleven configuration +fields. This new field extends the field set without removing or +altering any existing mandated parameter. Per ยง1.3, compliant +implementations MAY accept additional agent configuration fields. + +Token estimation MUST use the LangChain `get_num_tokens()` function +when available on the concrete chat model, falling back to a +`len(content) // 4` heuristic when the chat model is a mock or does +not expose a token-counting interface. + +### D-2: New LLM Agent Configuration Fields (`allow_tool_output_pruning`, `pruning_model`, `pruning_threshold`, `pruning_tool_filter`) + +**What:** Four new optional configuration fields controlling the tool +output pruning mechanism. + +| Field | Type | Required | Default | Description | +|-------|------|----------|---------|-------------| +| `allow_tool_output_pruning` | boolean | No | `false` | When true, each raw tool output is passed through an implicit LLM extraction pass before becoming a ToolMessage in the conversation context. The extraction pass asks the LLM to extract only task-relevant content from the tool output. The extraction pass must have access to the conversation context that led to the tool usage (the user's query, the system prompt, and the preceding conversation history up to the AIMessage containing the tool call) so that it can accurately determine what content is relevant. | +| `pruning_model` | string | No | (same as main `model`) | Model identifier used for the extraction LLM call. When absent, the agent's primary model is used. | +| `pruning_threshold` | integer | No | `512` | Positive integer threshold (in characters) for tool output length. When the raw tool output length is at or below this value, no pruning occurs even when `allow_tool_output_pruning` is true and the tool matches the filter. | +| `pruning_tool_filter` | list of strings | No | `["file_read", "shell"]` | List of tool names whose output is eligible for pruning. Pruning only applies when the invoked tool name appears in this list. When the list is empty, no tool output is pruned. | + +**Spec relationship:** Same as D-1 โ€” extensions to ยง4.4 permitted by +ยง1.3. `allow_tool_output_pruning` defaults to `false`, ensuring +backward compatibility: when no new config keys are provided, the +tool-call loop behaves identically to the current specification. + +### D-3: Tool Output Pruning Pass (Runtime Behavior) + +**What:** A runtime behavior that operates between tool execution and +ToolMessage construction. Pruning proceeds only when **all** of the +following conditions are met: + +1. The agent configuration has `allow_tool_output_pruning: true`. +2. The invoked tool name appears in the `pruning_tool_filter` list. +3. The raw tool output length (in characters) exceeds + `pruning_threshold`. +4. The main model did not explicitly disable pruning for this call + via the `output_prune` meta-argument (see step 2 below). + +When any condition is false, the raw tool output passes through +unmodified. + +The pruning pass: + +1. **Schema augmentation.** When `allow_tool_output_pruning: true`, each + tool whose name appears in `pruning_tool_filter` has its function + schema sent to the LLM augmented with two additional optional + parameters: + - `output_prune` (boolean, default `true`): setting it to `false` skips + pruning and delivers the full raw tool output. + - `output_prune_context` (string, default `None`): a focused description + from the main model of what exactly it is looking for in the tool + output โ€” what information it needs and why. This string replaces the + broader task context when constructing the extraction prompt (see + step 3 below). + + Both parameters are **meta-arguments**: they MUST be stripped from the + arguments before the actual tool is executed. + +2. **Opt-out detection.** If the model invokes a tool with + `output_prune: false`, the pruning pass is void โ€” the raw tool output + is used as the ToolMessage content with no modifications. + +3. **Extraction prompt construction.** If the model did not specify + `output_prune`, or set it to `true` (or the provider does not expose + the parameter in the tool call response), the pruning pass executes. + An extraction prompt is built comprising: + + - A **system message** instructing the pruning model to extract only + task-relevant information from the tool output, to indicate clearly + when non-relevant content was removed, and to preserve all data that + might be needed for subsequent decisions. + + - A **user message** whose content depends on whether the main model + provided an `output_prune_context`: + + *When `output_prune_context` is absent or empty:* the user message + contains the full task context (original user query, system prompt, + and conversation history up to the AIMessage containing the tool + calls), the tool name, and the raw tool output. + + *When `output_prune_context` is a non-empty string:* the task + context is **replaced** by `output_prune_context` โ€” the user message + contains only that string, the tool name, and the raw tool output. + The main model is responsible for encoding everything the pruning + model needs into this directive. + +4. **Pruning invocation.** Invoke the pruning model (or the primary model + when `pruning_model` is unset) in a stateless, single-turn call. + The pruning model's response must be structured with distinctive + markers (see ยง4.4.9) separating informational pruning metadata from + the extracted content. + +5. **ToolMessage construction.** Use the marked-up response as the + ToolMessage content. + +6. **Failure fallback.** On pruning call failure (network error, model + error, or malformed response), the raw tool output is used as-is and + a warning is logged. + +**What the pruning pass DOES NOT touch:** + +- System messages (including the agent's `system_prompt`). +- User messages and conversation history. +- AIMessage content (the LLM's reasoning and `tool_calls`). +- Any message already present in the `messages` list. + +The pruning pass operates exclusively on raw tool output before it +enters the conversation. Once a ToolMessage is constructed (with +pruned or raw content), it is never retroactively modified. + +**Spec relationship:** The specification defines the tool invocation +protocol (ยง4.5.3) and agent invocation in stream routes (ยง5.3.4). +Neither section currently addresses post-execution processing of +tool output before it enters the conversation context. This behavior +extends the runtime pipeline (an area that ยง1.3 allows +implementations to augment) without conflicting with any existing +`MUST`/`SHALL` requirement. The extraction prompt is an +implementation detail โ€” it is not user-configurable and does not +appear in the configuration document. + +### D-4: Budget-Exhausted Synthesis Flow (Runtime Behavior) + +**What:** When the accumulated token estimate exceeds the budget +ceiling (`token_budget_percent ร— model_context_window`): + +1. A synthesis prompt is injected as a user-style message (the most + recent message in the conversation) instructing the model to + produce a final answer with the information available. +2. One final tool-call round is permitted (the model may need one + more tool call to complete its reasoning). +3. After that final round, tool definitions are stripped from the + next `ainvoke()` call, forcing a text-only response. +4. The loop terminates with whatever content the model produces. + +If the model produces no tool calls in the final round, the +synthesis prompt alone forces termination without a second LLM call. + +**Spec relationship:** The loop-detection and termination logic is +closely related to the existing loop guards in ยง6.8. The +specification does not normatively define the internal tool-call +loop (it defines agent invocation at a higher abstraction level), +so this runtime behavior is an implementation extension that does +not alter any externally observable behavior defined by the +specification. + +### D-5: Token-Budget Logging Events + +**What:** Three new logging event categories, defined as +implementation-quality extensions: + +| Event | When | Level | +|-------|------|-------| +| `budget_warning` | Remaining budget falls below 25% | WARNING | +| `pruning_info` | Pruning pass completes | INFO (includes bytes saved) | +| `budget_exhausted` | Budget ceiling exceeded, synthesis flow triggered | ERROR | + +**Spec relationship:** The specification does not normatively define +logging or observability surfaces. These events fall within +implementation-quality territory as defined by ADR-2030 (D-4, D-6), +where error and information messages may be enhanced without +altering normative behavior. + +### D-6: Actual Token Usage Tracking + +**What:** The `usage_metadata` (specifically `total_tokens`) returned +by each `ainvoke()` call's response is stored in +`_last_token_usage`. Future token estimates use this actual +consumption data to improve estimation accuracy, supplementing the +`get_num_tokens()` / heuristic estimates. + +**Spec relationship:** Token usage metadata is provider-specific and +not part of the normative specification. Tracking it for estimation +purposes is purely an implementation detail. + +### D-7: Pruning Response Marker Format (ยง4.4.9) + +**What:** A normative marker format for pruning model responses, +consisting of two marker pairs: `[PRUNE_INFO_START]`/`[PRUNE_INFO_END]` +for informational metadata about what was pruned, and +`[PRUNE_OUTPUT_START]`/`[PRUNE_OUTPUT_END]` for the extracted relevant +content. Either block may be omitted. + +**Spec relationship:** This is an entirely new specification section +(ยง4.4.9) defining a concrete wire format that both the pruning model +and the implementation must agree on. Unlike D-5 (logging events, +which are implementation-quality), the marker format is normative +because the implementation must correctly parse the pruning model's +response to construct the ToolMessage. Without a defined format, the +implementation cannot distinguish pruning metadata from pruned +content, and the main model cannot know whether cutting occurred. + +The marker names (`[PRUNE_INFO_START]`, `[PRUNE_INFO_END]`, +`[PRUNE_OUTPUT_START]`, `[PRUNE_OUTPUT_END]`) are reserved strings +under ยง15 (Reserved Names); user configurations MUST NOT use them. + +--- + +## Consequences + +### Positive + +- **Context window safety:** Token-budget tracking prevents the most + common cause of silent LLM degradation (context overflow). The + model is given explicit awareness of its budget and can adapt its + tool-calling strategy accordingly. +- **No data loss from automatic compaction:** By rejecting + lossy/automatic summarization in favor of LLM-directed extraction, + the pruning pass preserves the model's complete reasoning chain. + Only raw tool output is filtered โ€” conversation reasoning is + sacrosanct. +- **LLM self-directed extraction:** The model decides what + information from each tool result is relevant. A 50KB log file + that only contains 3 relevant lines produces a ~200 character + ToolMessage, freeing space for higher-value data from subsequent + tool calls. +- **Backward compatibility:** All new configuration fields are + optional with defaults that preserve existing behavior. Existing + YAML configurations continue to work unchanged. +- **No spec violations:** Every extension is justified by ยง1.3 + (extensibility clause), ยง4.4 (new config fields where none are + mandated), or falls within implementation-quality territory not + normatively constrained by the specification. +- **Graceful degradation:** Pruning call failure falls back to raw + output. Budget exhaustion follows the existing stuck-model + recovery pattern. No failure mode causes data loss or silent + corruption. +- **Active cycle prevention:** By freeing space for additional tool + calls after pruning, the mechanism prevents premature loop + termination and allows the LLM to gather all needed data before + synthesizing its final response. + +### Negative / Risks + +- **Pruning latency:** Each tool call incurs an additional LLM + round-trip for extraction when `allow_tool_output_pruning` is + enabled. For workflows with many small tool calls, this can + meaningfully increase total execution time. Mitigation: pruning is + opt-in and configurable with a potentially lighter `pruning_model`. +- **Pruning quality dependency:** The extraction pass depends on the + LLM's ability to faithfully extract relevant content. An + over-aggressive extraction could discard data the main model later + needs. The extraction prompt includes instructions to "be concise + but complete โ€” include all data that might be needed for + subsequent decisions" to mitigate this. +- **Token estimation accuracy:** The `len(content) // 4` fallback + heuristic is crude for non-English text, code, and structured + data. Actual consumption tracked via `usage_metadata` partially + mitigates this but lags by one round. +- **Specification boundary increase:** Four new configuration fields, + two new runtime behaviors, and one new normative marker format + (ยง4.4.9) extend the specification's surface area. If these + mechanisms later prove insufficient, they may need to be deprecated + or replaced. +- **Config drift with ADR-2030:** ADR-2030's `_BUILTIN_TOOL_SCHEMAS` + and `llm_tools.py` are distinct from this ADR's pruning mechanism. + The pruning pass does not interact with tool schemas, so there is + no direct coupling, but developers must be aware that tool output + filtering operates on post-execution output, not on schema-level + descriptions. + +### Follow-up Required + +- **Specification update:** After battle-testing, if token-budget + awareness and tool output pruning prove stable, graduate the new + configuration fields into ยง4.4 of the Actor Configuration Standard + via a subsequent ADR. +- **ADR-2031 compliance validation:** Add BDD scenarios verifying + each extension's spec compliance boundary (e.g. verify that AIMessage + content is never passed through the pruning pass, verify that + backward-compatible defaults preserve existing behavior). +- **Token estimation accuracy test:** Add a test comparing + `get_num_tokens()` output against the heuristic fallback across a + representative corpus of messages (plain English, code blocks, + JSON, multilingual text). +- **Pruning extraction quality evaluation:** Run a benchmark + comparing pruned-vs-raw tool output across diverse file types + (logs, source code, documentation, JSON payloads) to measure + extraction quality and identify failure modes. + +--- + +## Alternatives Considered + +### A-1: Automatic lossy compaction (rejected) + +**Rejected because:** A compressor cannot know what the LLM still +needs. The model builds a reasoning chain: "I concluded X โ†’ I ran +tool Y โ†’ I got result Z โ†’ therefore...". Breaking that chain with +automatic summarization degrades decision quality. The LLM alone +knows what information is relevant from each tool result. + +### A-2: Implement pruning as a meta-tool (rejected) + +**Rejected because:** A `context_prune` meta-tool would require a +new tool schema, `tool_call_id` mapping, AIMessage-vs-ToolMessage +validation, and a handler โ€” all for a single operation. It would +also require retroactive surgery on the `messages` list (finding +and modifying prior ToolMessages), which is error-prone and +architecturally complex. An inline extraction pass in the tool-call +loop is simpler and has zero risk of missing or corrupting prior +entries. + +### A-3: Hard truncation of tool output at N characters (rejected) + +**Rejected because:** A fixed character limit cannot distinguish +between crucial and irrelevant content. A 10KB log file might +contain the single critical error line at byte 9000, which hard +truncation at 4000 characters would discard. LLM-directed extraction +is the only mechanism that preserves task-relevant information while +removing noise. + +### A-4: Prune AIMessages alongside ToolMessages (rejected) + +**Rejected because:** AIMessages contain the LLM's reasoning +(`content` field) and its tool-calling decisions (`tool_calls` +field). Pruning reasoning breaks the chain that the model needs for +coherent multi-step execution. The model must see its prior +conclusions to build on them. Tool output is the only safe target +for pruning. + +### A-5: Always prune, no opt-in flag (rejected) + +**Rejected because:** Not all workflows benefit from pruning. +Single-tool-call scenarios, workflows where every tool result is +small, and latency-sensitive real-time applications would incur the +pruning overhead with no benefit. Opt-in (`allow_tool_output_pruning: +false` default) ensures only configurations that need pruning pay +its cost, preserving backward compatibility. + +### A-6: Shut down the tool loop without synthesis when budget exhausted (rejected) + +**Rejected because:** This would force the model to produce a final +answer without any signal that its context was exhausted, leading +to the same degraded-output problem that budget tracking is meant to +solve. The synthesis prompt gives the model explicit instructions to +conclude with available information, which produces measurably +better final answers than silent termination. \ No newline at end of file diff --git a/features/environment.py b/features/environment.py index f7babb6..754b9fa 100644 --- a/features/environment.py +++ b/features/environment.py @@ -182,6 +182,14 @@ def after_scenario(context, scenario): except Exception: pass + # Stop any _prune_patch from pruning model tests. + if hasattr(context, "_prune_patch") and context._prune_patch is not None: + try: + context._prune_patch.stop() + except Exception: + pass + context._prune_patch = None + # Remove log-capture handler installed by cleanup warning-log tests. if hasattr(context, "_cleanup_log_handler"): from cleveractors.agents.llm import logger as llm_logger diff --git a/features/llm_agent_tool_calling.feature b/features/llm_agent_tool_calling.feature index a31e90b..6fe8889 100644 --- a/features/llm_agent_tool_calling.feature +++ b/features/llm_agent_tool_calling.feature @@ -80,3 +80,268 @@ Feature: LLMAgent Tool Calling And I set a mock chat model that returns tool_calls with an empty tool name And I tool_calling process a message "Call me" Then the messages list should contain a ToolMessage with content "Tool name is empty" + + # โ”€โ”€ Token Budget Awareness (ยง4.4.7) โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + + Scenario: LLMAgent validates token_budget_percent range + Given I have an LLM agent configuration with token_budget_percent 1.5 + When I attempt to create the LLM agent + Then ConfigurationError should be raised containing "token_budget_percent must be in" + + Scenario: LLMAgent stores token_budget_percent when valid + Given I have an LLM agent configuration with token_budget_percent 0.50 + When I create the LLM agent + Then the agent _token_budget_percent should be 0.5 + + Scenario: LLMAgent _token_budget_percent is None by default + Given I have a basic LLM agent configuration (no tools) + When I create the LLM agent + Then the agent _token_budget_percent should be None + + # โ”€โ”€ Tool Output Pruning (ยง4.4.8) โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + + Scenario: LLMAgent stores allow_tool_output_pruning when true + Given I have an LLM agent config with allow_tool_output_pruning true + When I create the LLM agent + Then the agent _allow_tool_output_pruning should be true + + Scenario: LLMAgent allow_tool_output_pruning defaults to false + Given I have a basic LLM agent configuration (no tools) + When I create the LLM agent + Then the agent _allow_tool_output_pruning should be false + + Scenario: LLMAgent stores pruning_model when configured + Given I have an LLM agent config with pruning_model "gpt-3.5-turbo" + When I create the LLM agent + Then the agent _pruning_model should be "gpt-3.5-turbo" + + Scenario: LLMAgent pruning_model defaults to None + Given I have a basic LLM agent configuration (no tools) + When I create the LLM agent + Then the agent _pruning_model should be None + + # โ”€โ”€ pruning_threshold (ยง4.4.8 D-2) โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + + Scenario: LLMAgent stores pruning_threshold when configured + Given I have an LLM agent config with pruning_threshold 1024 + When I create the LLM agent + Then the agent _pruning_threshold should be 1024 + + Scenario: LLMAgent pruning_threshold defaults to 512 + Given I have a basic LLM agent configuration (no tools) + When I create the LLM agent + Then the agent _pruning_threshold should be 512 + + Scenario: LLMAgent validates pruning_threshold must be positive + Given I have an LLM agent config with pruning_threshold -10 + When I attempt to create the LLM agent + Then ConfigurationError should be raised containing "pruning_threshold must be a positive integer" + + Scenario: LLMAgent validates pruning_threshold must be integer + Given I have an LLM agent config with pruning_threshold "not_an_int" + When I attempt to create the LLM agent + Then ConfigurationError should be raised containing "pruning_threshold must be a positive integer" + + # โ”€โ”€ pruning_tool_filter (ยง4.4.8 D-2) โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + + Scenario: LLMAgent stores pruning_tool_filter when configured + Given I have an LLM agent config with pruning_tool_filter ["file_read", "shell", "http_request"] + When I create the LLM agent + Then the agent _pruning_tool_filter should be ["file_read", "shell", "http_request"] + + Scenario: LLMAgent pruning_tool_filter defaults to file_read and shell + Given I have a basic LLM agent configuration (no tools) + When I create the LLM agent + Then the agent _pruning_tool_filter should be ["file_read", "shell"] + + Scenario: LLMAgent validates pruning_tool_filter must be a list + Given I have an LLM agent config with pruning_tool_filter "not_a_list" + When I attempt to create the LLM agent + Then ConfigurationError should be raised containing "pruning_tool_filter must be a list" + + Scenario: LLMAgent validates pruning_tool_filter items must be strings + Given I have an LLM agent config with pruning_tool_filter ["file_read", 123] + When I attempt to create the LLM agent + Then ConfigurationError should be raised containing "pruning_tool_filter items must be strings" + + Scenario: LLMAgent pruning_tool_filter empty list is preserved (disables pruning) + Given I have an LLM agent config with pruning_tool_filter [] + When I create the LLM agent + Then the agent _pruning_tool_filter should be [] + + # โ”€โ”€ Schema augmentation with pruning_tool_filter (ยง4.4.8 D-3) โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + + Scenario: Tool schema augmented only for tools in pruning_tool_filter + Given I have an LLM agent config with allow_tool_output_pruning true + And I have pruning_tool_filter ["file_read"] + And I have tools [{"type": "function", "function": {"name": "file_read", "parameters": {"type": "object", "properties": {"path": {"type": "string"}}}}},{"type": "function", "function": {"name": "shell", "parameters": {"type": "object", "properties": {"cmd": {"type": "string"}}}}}] + When I create the LLM agent + And I augment the tool schemas for pruning + Then the file_read schema should contain output_prune parameter + And the shell schema should NOT contain output_prune parameter + + # โ”€โ”€ Schema Augmentation (ยง4.4.8 step 1) โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + + Scenario: Tool schemas augmented with output_prune when pruning enabled + Given I have an LLM agent config with allow_tool_output_pruning true + And I have tools [{"type": "function", "function": {"name": "file_read", "parameters": {"type": "object", "properties": {"path": {"type": "string"}}}}}] + When I create the LLM agent + And I augment the tool schemas for pruning + Then the augmented schemas should contain output_prune parameter + + Scenario: Tool schema augmentation does not require pruning_tool_filter to be set + Given I have a basic LLM agent configuration (no tools) + And I have tools [{"type": "function", "function": {"name": "file_read", "parameters": {"type": "object", "properties": {"path": {"type": "string"}}}}}] + When I create the LLM agent + And I augment the tool schemas for pruning + Then the augmented schemas should contain output_prune parameter + + Scenario: process_message passes augmented tools to ainvoke when pruning enabled + Given I have an LLM agent config with allow_tool_output_pruning true + And I have tools [{"type": "function", "function": {"name": "file_read", "parameters": {"type": "object", "properties": {"path": {"type": "string"}}}}}] + When I create the LLM agent + And I set a mock chat model that returns a plain text response + And I tool_calling process a message "Hello" + Then the ainvoke tools kwarg should contain the output_prune parameter + + # โ”€โ”€ Pruning Marker Parsing (ยง4.4.9) โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + + Scenario: Parse pruning response with info and output blocks + Given I have a basic LLM agent configuration (no tools) + When I create the LLM agent + And I parse a pruning response with both info and output blocks + Then the parsed result should contain PRUNE_INFO_START + And the parsed result should contain PRUNE_OUTPUT_START + + Scenario: Parse pruning response with only output block + Given I have a basic LLM agent configuration (no tools) + When I create the LLM agent + And I parse a pruning response with only an output block + Then the parsed result should contain PRUNE_OUTPUT_START + And the parsed result should NOT contain PRUNE_INFO_START + + Scenario: Parse pruning response with no markers falls back to full text + Given I have a basic LLM agent configuration (no tools) + When I create the LLM agent + And I parse a pruning response with no markers containing "some extracted content" + Then the parsed result should be "some extracted content" + + # โ”€โ”€ Model Context Window (ยง4.4.7) โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + + Scenario: _get_model_context_window uses max_tokens when set above 10K + Given I have an LLM agent config with max_tokens 200000 + When I create the LLM agent + Then _get_model_context_window should return 200000 + + Scenario: _get_model_context_window defaults to 128000 when max_tokens is small + Given I have an LLM agent config with max_tokens 1000 + When I create the LLM agent + Then _get_model_context_window should return 128000 + + Scenario: _get_model_context_window defaults to 128000 when max_tokens not set + Given I have a basic LLM agent configuration (no tools) + When I create the LLM agent + Then _get_model_context_window should return 128000 + + # โ”€โ”€ Token Estimation (ยง4.4.7) โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + + Scenario: _estimate_token_count uses heuristic fallback + Given I have a basic LLM agent configuration (no tools) + When I create the LLM agent + And I set a mock chat model that returns a plain text response + And I estimate token count for messages with content "hello" + Then the estimated token count should be greater than 0 + + # โ”€โ”€ _run_pruning_pass with mocked pruning model (ยง4.4.8) โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + + Scenario: _run_pruning_pass returns marked-up response + Given I have an LLM agent config with allow_tool_output_pruning true + When I create the LLM agent + And I set a mock pruning model that returns a marked-up response + And I run the pruning pass on tool "file_read" with output "50KB of logs" + Then the pruned result should contain PRUNE_INFO_START + And the pruned result should contain PRUNE_OUTPUT_START + And the pruned result should contain "extracted content" + + Scenario: _run_pruning_pass falls back to raw output on error + Given I have an LLM agent config with allow_tool_output_pruning true + When I create the LLM agent + And I set a mock pruning model that raises an exception + And I run the pruning pass on tool "echo" with output "raw data" + Then the pruned result should equal "raw data" + + Scenario: _run_pruning_pass returns empty raw output unchanged + Given I have an LLM agent config with allow_tool_output_pruning true + When I create the LLM agent + And I run the pruning pass on tool "echo" with empty output + Then the pruned result should be empty + + # โ”€โ”€ Token budget exhaustion in tool loop (ยง4.4.7) โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + + Scenario: Token budget exhaustion triggers synthesis flow + Given I have an LLM agent configuration with tools [{"name": "echo"}] + And I set token_budget_percent to 0.000001 + And I set tool_max_rounds to 5 + When I create the LLM agent + And I set a mock chat model that returns tool_calls then text for budget exhaustion + And I tool_calling process a message "Test budget" + Then the tool_calling result should contain "Budget exhausted" + + # โ”€โ”€ Pruning pass in tool loop (ยง4.4.8) โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + + Scenario: Pruning pass NOT called for non-file_read tools + Given I have an LLM agent configuration with tools [{"name": "echo"}] + And I set allow_tool_output_pruning to true + And I set tool_max_rounds to 2 + When I create the LLM agent + And I set a mock chat model that returns one tool call then text + And I set a mock pruning model that returns a marked-up response + And I tool_calling process a message "Test pruning scope" + Then the tool_calling result should contain "Final answer" + And the pruning pass should NOT have been called + + # โ”€โ”€ output_prune stripping (ยง4.4.8 step 1) โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + + Scenario: output_prune false skips pruning in tool loop + Given I have an LLM agent configuration with tools [{"name": "file_read"}] + And I set allow_tool_output_pruning to true + And I set tool_max_rounds to 2 + When I create the LLM agent + And I set a mock chat model that returns tool call with output_prune false then text + And I set a mock pruning model that returns a marked-up response + And I tool_calling process a message "Test opt out" + Then the tool_calling result should contain "Final answer after tools" + And the pruning pass should NOT have been called + + # โ”€โ”€ Pruning with prune_context (ยง4.4.8 step 3) โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + + Scenario: _run_pruning_pass with prune_context uses search directive + Given I have an LLM agent config with allow_tool_output_pruning true + When I create the LLM agent + And I set a mock pruning model that returns a marked-up response + And I run the pruning pass on "file_read" with output "100 lines" and prune_context "find errors" + Then the pruned result should contain PRUNE_OUTPUT_START + + # โ”€โ”€ Budget exhaustion with pruning enabled โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + + Scenario: Budget exhaustion with pruning enabled triggers pruning in synthesis round + Given I have an LLM agent configuration with tools [{"name": "file_read"}] + And I set token_budget_percent to 0.000001 + And I set allow_tool_output_pruning to true + And I set tool_max_rounds to 5 + When I create the LLM agent + And I set a mock chat model that triggers budget exhaustion with a file_read tool call + And I set a mock pruning model that returns a marked-up response + And I tool_calling process a message "Test budget with prune" + Then the tool_calling result should contain "Budget exhausted" + And the pruning pass should have been called + + # โ”€โ”€ Stuck-model synthesis with tool calls โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + + Scenario: Stuck-model synthesis executes tool calls in synthesis round + Given I have an LLM agent configuration with tools [{"name": "echo"}] + And I set tool_max_rounds to 2 + When I create the LLM agent + And I set a mock chat model that gets stuck then returns tool_calls in synthesis + And I tool_calling process a message "Get stuck" + Then the tool_calling result should contain "Synthesized answer" diff --git a/features/steps/llm_agent_tool_calling_steps.py b/features/steps/llm_agent_tool_calling_steps.py index 8a550d9..108d736 100644 --- a/features/steps/llm_agent_tool_calling_steps.py +++ b/features/steps/llm_agent_tool_calling_steps.py @@ -184,6 +184,32 @@ def step_ainvoke_no_tools_kwarg(context): assert "tools" not in kwargs +@then("the ainvoke tools kwarg should contain the output_prune parameter") +def step_ainvoke_tools_contains_output_prune(context): + """Verify the tools kwarg in ainvoke contains the output_prune parameter.""" + call = context.llm_agent.chat_model.ainvoke.call_args + if call is None: + raise AssertionError("ainvoke was never called") + kwargs = call.kwargs if hasattr(call, "kwargs") else call[1] + assert "tools" in kwargs, ( + f"ainvoke was called but without a ``tools`` kwarg. " + f"call_kwargs={list(kwargs.keys())}" + ) + tools = kwargs["tools"] + found = False + for tool in tools: + if tool.get("type") == "function": + params = tool.get("function", {}).get("parameters", {}) + if params.get("type") == "object": + properties = params.get("properties", {}) + if "output_prune" in properties: + found = True + break + assert found, ( + f"output_prune parameter not found in ainvoke tools kwarg. tools={tools}" + ) + + @then("the tool_calling result should contain {expected}") def step_result_contains(context, expected): """Check that the result string contains an expected substring.""" @@ -361,3 +387,619 @@ def step_messages_has_exact_toolmessage(context, text): f"No ToolMessage found with exact content {text!r}. " f"ToolMessages: {[str(m.content) for m in tool_msgs]}" ) + + +# โ”€โ”€โ”€ Given steps (token-budget / pruning config) โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + + +@given("I have an LLM agent configuration with token_budget_percent {value:g}") +def step_config_with_budget(context, value): + context.llm_config = { + "name": "budget_agent", + "provider": "openai", + "api_key": "test_key", + "token_budget_percent": value, + } + + +@given("I have an LLM agent config with allow_tool_output_pruning true") +def step_config_with_pruning_true(context): + context.llm_config = { + "name": "prune_agent", + "provider": "openai", + "api_key": "test_key", + "allow_tool_output_pruning": True, + } + + +@given('I have an LLM agent config with pruning_model "{model}"') +def step_config_with_pruning_model(context, model): + context.llm_config = { + "name": "pm_agent", + "provider": "openai", + "api_key": "test_key", + "pruning_model": model, + } + + +@given("I have an LLM agent config with pruning_threshold {n:d}") +def step_config_with_pruning_threshold(context, n): + context.llm_config = { + "name": "pt_agent", + "provider": "openai", + "api_key": "test_key", + "pruning_threshold": n, + } + + +@given("I have an LLM agent config with pruning_threshold {bad_value}") +def step_config_with_pruning_threshold_bad(context, bad_value): + context.llm_config = { + "name": "pt_agent", + "provider": "openai", + "api_key": "test_key", + "pruning_threshold": bad_value, + } + + +@given("I have an LLM agent config with pruning_tool_filter {json}") +def step_config_with_pruning_filter(context, json): + parsed = _json.loads(json) + context.llm_config = { + "name": "pf_agent", + "provider": "openai", + "api_key": "test_key", + "pruning_tool_filter": parsed, + } + + +@given("I have pruning_tool_filter {json}") +def step_set_pruning_filter(context, json): + parsed = _json.loads(json) + context.llm_config["pruning_tool_filter"] = parsed + + +@given("I have an LLM agent config with max_tokens {n:d}") +def step_config_with_max_tokens(context, n): + context.llm_config = { + "name": "max_tokens_agent", + "provider": "openai", + "api_key": "test_key", + "model": "gpt-4o", + "max_tokens": n, + } + + +@given('I have an LLM agent config with model "{model}"') +def step_config_with_model(context, model): + context.llm_config = { + "name": "model_agent", + "provider": "openai", + "api_key": "test_key", + "model": model, + } + + +@given("I have tools {tools_json}") +def step_set_tools_raw(context, tools_json): + parsed = _json.loads(tools_json) + context.llm_config["tools"] = parsed + + +@given("I set token_budget_percent to {value:g}") +def step_set_budget_percent(context, value): + context.llm_config["token_budget_percent"] = value + + +@given("I set allow_tool_output_pruning to true") +def step_set_pruning_true(context): + context.llm_config["allow_tool_output_pruning"] = True + + +# โ”€โ”€โ”€ When steps (token-budget / pruning) โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + + +@when("I attempt to create the LLM agent") +def step_attempt_create_agent(context): + try: + context.llm_agent = _make_llm_agent_with_config(context.llm_config) + context._create_error = None + except ConfigurationError as e: + context._create_error = e + + +@when("I augment the tool schemas for pruning") +def step_augment_schemas(context): + agent = context.llm_agent + lc_tools = agent._lc_tools + context._augmented_schemas = agent._augment_tool_schemas_for_pruning( + list(lc_tools) if lc_tools else [] + ) + + +@when("I parse a pruning response with both info and output blocks") +def step_parse_full_pruning(context): + agent = context.llm_agent + response = ( + agent._PRUNE_INFO_START + + "\nRemoved 200 chars\n" + + agent._PRUNE_INFO_END + + "\n" + + agent._PRUNE_OUTPUT_START + + "\nExtracted content\n" + + agent._PRUNE_OUTPUT_END + ) + context._parsed_pruning = agent._parse_pruning_response(response) + + +@when("I parse a pruning response with only an output block") +def step_parse_output_only(context): + agent = context.llm_agent + response = ( + agent._PRUNE_OUTPUT_START + "\nJust the facts\n" + agent._PRUNE_OUTPUT_END + ) + context._parsed_pruning = agent._parse_pruning_response(response) + + +@when('I parse a pruning response with no markers containing "{text}"') +def step_parse_no_markers(context, text): + context._parsed_pruning = context.llm_agent._parse_pruning_response(text) + + +@when('I estimate token count for messages with content "{text}"') +def step_estimate_tokens(context, text): + from langchain_core.messages import HumanMessage as HM + + msgs = [HM(content=text)] + context._estimated_tokens = context.llm_agent._estimate_token_count(msgs) + + +def _cleanup_prune_patch(context): + p = getattr(context, "_prune_patch", None) + if p is not None: + p.stop() + context._prune_patch = None + + +@when("I set a mock pruning model that returns a marked-up response") +def step_mock_pruning_model(context): + mock_prune_model = Mock(spec=BaseChatModel) + + async def _mock_prune_ainvoke(messages, **kw): + context._pruning_was_called = True + response = Mock(spec=AIMessage) + response.content = ( + "[PRUNE_INFO_START]\nRemoved 48K chars of logs\n[PRUNE_INFO_END]\n" + "[PRUNE_OUTPUT_START]\nextracted content\n[PRUNE_OUTPUT_END]" + ) + return response + + mock_prune_model.ainvoke = _mock_prune_ainvoke + context._prune_patch = patch( + "cleveractors.agents.llm.build_chat_model", + return_value=mock_prune_model, + ) + context._prune_patch.start() + + +@when("I set a mock pruning model that raises an exception") +def step_mock_pruning_model_error(context): + mock_prune_model = Mock(spec=BaseChatModel) + + async def _mock_prune_error(messages, **kw): + raise RuntimeError("pruning model down") + + mock_prune_model.ainvoke = _mock_prune_error + context._prune_patch = patch( + "cleveractors.agents.llm.build_chat_model", + return_value=mock_prune_model, + ) + context._prune_patch.start() + + +@when('I run the pruning pass on tool "{tool_name}" with output "{output}"') +def step_run_pruning_pass(context, tool_name, output): + context._prune_result = asyncio.run( + context.llm_agent._run_pruning_pass(tool_name, output, []) + ) + + +@when( + 'I run the pruning pass on "{tool_name}" with output ' + '"{output}" and prune_context "{prune_context}"' +) +def step_run_pruning_pass_with_context(context, tool_name, output, prune_context): + context._prune_result = asyncio.run( + context.llm_agent._run_pruning_pass( + tool_name, + output, + [], + prune_context=prune_context, + ) + ) + + +@when('I run the pruning pass on tool "{tool_name}" with empty output') +def step_run_pruning_pass_empty(context, tool_name): + context._prune_result = asyncio.run( + context.llm_agent._run_pruning_pass(tool_name, "", []) + ) + + +@when("I set a mock chat model that returns tool_calls then text for budget exhaustion") +def step_mock_budget_exhaustion(context): + call_counter = [0] + + async def _mock_ainvoke(messages, **invoke_kwargs): + call_counter[0] += 1 + if call_counter[0] == 1: + return _create_mock_response( + "", + tool_calls=[ + { + "id": "call_budget_1", + "type": "function", + "function": { + "name": "echo", + "arguments": '{"text": "test"}', + }, + } + ], + ) + return _create_mock_response("Budget exhausted") + + mock_model = Mock(spec=BaseChatModel) + mock_model.ainvoke = _mock_ainvoke + context.llm_agent.chat_model = mock_model + + +@when("I set a mock chat model that returns one tool call then text") +def step_mock_one_tool_call(context): + call_counter = [0] + context._pruning_was_called = False + + async def _mock_ainvoke(messages, **invoke_kwargs): + call_counter[0] += 1 + if call_counter[0] == 1: + return _create_mock_response( + "Thinking...", + tool_calls=[ + { + "id": "call_prune_1", + "type": "function", + "function": { + "name": "echo", + "arguments": '{"text": "hello"}', + }, + } + ], + ) + return _create_mock_response("Final answer") + + mock_model = Mock(spec=BaseChatModel) + mock_model.ainvoke = _mock_ainvoke + context.llm_agent.chat_model = mock_model + + +@when( + "I set a mock chat model that returns tool call with output_prune false then text" +) +def step_mock_output_prune_false(context): + import os + + cwd = os.getcwd() + optout_path = os.path.join(cwd, "test_optout_file.txt") + with open(optout_path, "w") as f: + f.write("line 1\nextra data line 2\nline 3\n") + context._optout_file_path = optout_path + if not hasattr(context, "_cleanup_files"): + context._cleanup_files = [] + context._cleanup_files.append(optout_path) + + call_counter = [0] + context._pruning_was_called = False + + async def _mock_ainvoke(messages, **invoke_kwargs): + call_counter[0] += 1 + if call_counter[0] == 1: + return _create_mock_response( + "", + tool_calls=[ + { + "id": "call_optout_1", + "type": "function", + "function": { + "name": "file_read", + "arguments": _json.dumps( + { + "file": context._optout_file_path, + "output_prune": False, + } + ), + }, + } + ], + ) + return _create_mock_response("Final answer after tools") + + mock_model = Mock(spec=BaseChatModel) + mock_model.ainvoke = _mock_ainvoke + context.llm_agent.chat_model = mock_model + + +# โ”€โ”€โ”€ Then steps (token-budget / pruning) โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + + +@then('ConfigurationError should be raised containing "{text}"') +def step_config_error_raised(context, text): + err = context._create_error + assert err is not None, "Expected ConfigurationError but none was raised" + assert text in str(err), ( + f"ConfigurationError message did not contain {text!r}: {err}" + ) + + +@then("the agent _token_budget_percent should be {value:g}") +def step_budget_stored(context, value): + actual = context.llm_agent._token_budget_percent + assert actual == value, f"_token_budget_percent = {actual!r} != {value!r}" + + +@then("the agent _token_budget_percent should be None") +def step_budget_none(context): + assert context.llm_agent._token_budget_percent is None + + +@then("the agent _allow_tool_output_pruning should be true") +def step_pruning_true(context): + assert context.llm_agent._allow_tool_output_pruning is True + + +@then("the agent _allow_tool_output_pruning should be false") +def step_pruning_false(context): + assert context.llm_agent._allow_tool_output_pruning is False + + +@then('the agent _pruning_model should be "{model}"') +def step_pruning_model_stored(context, model): + assert context.llm_agent._pruning_model == model + + +@then("the agent _pruning_model should be None") +def step_pruning_model_none(context): + assert context.llm_agent._pruning_model is None + + +@then("the agent _pruning_threshold should be {n:d}") +def step_pruning_threshold_stored(context, n): + assert context.llm_agent._pruning_threshold == n, ( + f"_pruning_threshold = {context.llm_agent._pruning_threshold!r} != {n}" + ) + + +@then("the agent _pruning_tool_filter should be {json}") +def step_pruning_filter_stored(context, json): + parsed = _json.loads(json) + assert context.llm_agent._pruning_tool_filter == parsed, ( + f"_pruning_tool_filter = {context.llm_agent._pruning_tool_filter!r} != {parsed!r}" + ) + + +@then("the file_read schema should contain output_prune parameter") +def step_augmented_file_read_has_prune(context): + augmented = context._augmented_schemas + assert augmented is not None + for schema in augmented: + if schema.get("function", {}).get("name") == "file_read": + props = ( + schema.get("function", {}).get("parameters", {}).get("properties", {}) + ) + assert "output_prune" in props, ( + f"file_read schema missing output_prune: {list(props.keys())}" + ) + return + raise AssertionError("file_read schema not found in augmented tools") + + +@then("the shell schema should NOT contain output_prune parameter") +def step_augmented_shell_no_prune(context): + augmented = context._augmented_schemas + assert augmented is not None + for schema in augmented: + if schema.get("function", {}).get("name") == "shell": + props = ( + schema.get("function", {}).get("parameters", {}).get("properties", {}) + ) + assert "output_prune" not in props, ( + f"shell schema should NOT have output_prune but does: {list(props.keys())}" + ) + return + raise AssertionError("shell schema not found in augmented tools") + + +@then("the augmented schemas should contain output_prune parameter") +def step_augmented_has_output_prune(context): + augmented = context._augmented_schemas + assert augmented is not None + assert len(augmented) > 0 + for schema in augmented: + props = schema.get("function", {}).get("parameters", {}).get("properties", {}) + assert "output_prune" in props, f"output_prune not found: {list(props.keys())}" + assert props["output_prune"]["type"] == "boolean" + + +@then("the parsed result should contain PRUNE_INFO_START") +def step_parsed_has_info(context): + assert context.llm_agent._PRUNE_INFO_START in context._parsed_pruning + + +@then("the parsed result should contain PRUNE_OUTPUT_START") +def step_parsed_has_output(context): + assert context.llm_agent._PRUNE_OUTPUT_START in context._parsed_pruning + + +@then("the parsed result should NOT contain PRUNE_INFO_START") +def step_parsed_no_info(context): + assert context.llm_agent._PRUNE_INFO_START not in context._parsed_pruning + + +@then('the parsed result should be "{text}"') +def step_parsed_equals(context, text): + assert context._parsed_pruning == text, ( + f"Expected {text!r}, got {context._parsed_pruning!r}" + ) + + +@then("_get_model_context_window should return {n:d}") +def step_context_window(context, n): + actual = context.llm_agent._get_model_context_window() + assert actual == n, f"_get_model_context_window = {actual} != {n}" + + +@then("the estimated token count should be greater than 0") +def step_estimate_positive(context): + assert context._estimated_tokens > 0, ( + f"Expected >0 tokens, got {context._estimated_tokens}" + ) + + +@then("the pruned result should contain PRUNE_INFO_START") +def step_pruned_has_info(context): + assert context.llm_agent._PRUNE_INFO_START in context._prune_result + _cleanup_prune_patch(context) + + +@then("the pruned result should contain PRUNE_OUTPUT_START") +def step_pruned_has_output(context): + assert context.llm_agent._PRUNE_OUTPUT_START in context._prune_result + + +@then('the pruned result should contain "{text}"') +def step_pruned_contains(context, text): + assert text in context._prune_result, ( + f"Expected {text!r} in pruned result: {context._prune_result!r}" + ) + + +@then('the pruned result should equal "{text}"') +def step_pruned_equals(context, text): + assert context._prune_result == text, ( + f"Expected {text!r}, got {context._prune_result!r}" + ) + + +@then("the pruned result should be empty") +def step_pruned_empty(context): + assert context._prune_result == "", f"Expected empty, got {context._prune_result!r}" + + +@then("the pruning pass should have been called") +def step_pruning_was_called(context): + assert context._pruning_was_called, "Pruning pass was not called" + _cleanup_prune_patch(context) + + +@then("the pruning pass should NOT have been called") +def step_pruning_not_called(context): + assert not context._pruning_was_called, ( + "Pruning pass was called but should not have been" + ) + _cleanup_prune_patch(context) + + +def _cleanup_prune_patch(context): + p = getattr(context, "_prune_patch", None) + if p is not None: + p.stop() + context._prune_patch = None + + +# โ”€โ”€โ”€ Budget exhaustion with file_read tool โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + + +@when( + "I set a mock chat model that triggers budget exhaustion with a file_read tool call" +) +def step_mock_budget_exhaustion_file_read(context): + import os + + cwd = os.getcwd() + test_file_path = os.path.join(cwd, "test_budget_file.txt") + large_content = "x" * 600 + with open(test_file_path, "w") as f: + f.write(large_content) + context._budget_file_path = "test_budget_file.txt" + if not hasattr(context, "_cleanup_files"): + context._cleanup_files = [] + context._cleanup_files.append(test_file_path) + context._pruning_was_called = False + + call_counter = [0] + + async def _mock_ainvoke(messages, **invoke_kwargs): + call_counter[0] += 1 + if call_counter[0] == 1: + return _create_mock_response( + "", + tool_calls=[ + { + "id": "call_budget_fr", + "type": "function", + "function": { + "name": "file_read", + "arguments": _json.dumps( + {"file": context._budget_file_path} + ), + }, + } + ], + ) + return _create_mock_response("Budget exhausted") + + mock_model = Mock(spec=BaseChatModel) + mock_model.ainvoke = _mock_ainvoke + context.llm_agent.chat_model = mock_model + + +# โ”€โ”€โ”€ Stuck-model synthesis with tool calls โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + + +@when("I set a mock chat model that gets stuck then returns tool_calls in synthesis") +def step_mock_stuck_with_synth_tools(context): + call_counter = [0] + + async def _mock_ainvoke(messages, **invoke_kwargs): + call_counter[0] += 1 + if call_counter[0] <= 2: + return _create_mock_response( + "", + tool_calls=[ + { + "id": f"call_stuck_{call_counter[0]}", + "type": "function", + "function": { + "name": "echo", + "arguments": '{"text": "test"}', + }, + } + ], + ) + if call_counter[0] == 3: + return _create_mock_response( + "Here is the synthesis", + tool_calls=[ + { + "id": "call_synth_echo", + "type": "function", + "function": { + "name": "echo", + "arguments": '{"text": "synthesis output"}', + }, + } + ], + ) + return _create_mock_response("Synthesized answer") + + mock_model = Mock(spec=BaseChatModel) + mock_model.ainvoke = _mock_ainvoke + context.llm_agent.chat_model = mock_model diff --git a/robot/TokenBudgetTestLib.py b/robot/TokenBudgetTestLib.py new file mode 100644 index 0000000..3a856d4 --- /dev/null +++ b/robot/TokenBudgetTestLib.py @@ -0,0 +1,381 @@ +"""Library for token-budget and pruning integration tests. + +Provides keywords that create an Executor with a mock chat model that +exercises token-budget awareness and tool-output pruning without +requiring a real LLM API key. +""" + +from __future__ import annotations + +import asyncio +import json as _json +import tempfile +from typing import Any +from unittest.mock import MagicMock, patch + +from langchain_core.messages import AIMessage + +from cleveractors.result import ActorResult +from cleveractors.runtime import Executor, create_executor + + +class TokenBudgetTestLib: + """Keywords for integration tests of token-budget and pruning.""" + + def __init__(self) -> None: + self._executor: Executor | None = None + self._last_result: ActorResult | None = None + self._mock_ainvoke_calls: list[dict[str, Any]] = [] + self._patches: list[Any] = [] + self._pruning_was_called: bool = False + _tf = tempfile.NamedTemporaryFile(delete=False, mode="w") + _tf.write("") + _tf.close() + self._temp_file = _tf.name + + def _teardown_patches(self) -> None: + for p in self._patches: + p.stop() + self._patches.clear() + if hasattr(self, "_temp_file"): + import os + + try: + os.unlink(self._temp_file) + except OSError: + pass + + # โ”€โ”€ Common helpers โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + + def _make_basic_executor(self, extra_config: dict[str, Any] | None = None) -> None: + """Create an Executor with a tool-configured LLM agent.""" + self._teardown_patches() + self._mock_ainvoke_calls = [] + self._pruning_was_called = False + + config: dict[str, Any] = { + "type": "llm", + "name": "budget_test_agent", + "provider": "openai", + "model": "gpt-3.5-turbo", + "config": { + "tools": [{"name": "echo"}], + }, + } + if extra_config: + config["config"].update(extra_config) + + def _build_mock_model(*args, **kwargs): + mock_model = MagicMock() + mock_model.temperature = 0.7 + + async def _mock_ainvoke(messages, **invoke_kwargs): + self._mock_ainvoke_calls.append({"kwargs": invoke_kwargs}) + return AIMessage( + content="Response text", + usage_metadata={ + "input_tokens": 15, + "output_tokens": 10, + "total_tokens": 10, + }, + ) + + mock_model.ainvoke = _mock_ainvoke + mock_model.astream = _mock_ainvoke + return mock_model + + patcher = patch( + "cleveractors.agents.llm.build_chat_model", + side_effect=_build_mock_model, + ) + patcher.start() + self._patches.append(patcher) + + self._executor = create_executor( + config_dict=config, + credentials={"openai": {"api_key": "mock-key"}}, + limits={}, + pricing={}, + ) + + async def _execute_async(self, message: str) -> ActorResult: + return await self._executor.execute(message) + + def _execute(self, message: str) -> None: + try: + self._last_result = asyncio.run(self._execute_async(message)) + finally: + self._teardown_patches() + + # โ”€โ”€ Keywords: Token-Budget Awareness โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + + def create_executor_with_budget(self, budget_percent: str) -> None: + """Create an Executor with token_budget_percent set.""" + pct = float(budget_percent) + self._make_basic_executor({"token_budget_percent": pct}) + + def create_executor_with_low_budget_for_exhaustion(self) -> None: + """Create an Executor with an extremely low budget that triggers + exhaustion on the first message.""" + self._teardown_patches() + self._mock_ainvoke_calls = [] + + config: dict[str, Any] = { + "type": "llm", + "name": "budget_exhaust_agent", + "provider": "openai", + "model": "gpt-3.5-turbo", + "config": { + "tools": [{"name": "echo"}], + "token_budget_percent": 0.000001, + }, + } + + def _build_budget_model(*args, **kwargs): + mock_model = MagicMock() + mock_model.temperature = 0.7 + + async def _mock_ainvoke(messages, **invoke_kwargs): + self._mock_ainvoke_calls.append({"kwargs": invoke_kwargs}) + call_count = len(self._mock_ainvoke_calls) + if call_count == 1 and invoke_kwargs.get("tools"): + return AIMessage( + content="", + usage_metadata={ + "input_tokens": 10, + "output_tokens": 5, + "total_tokens": 5, + }, + tool_calls=[ + { + "id": "call_budget_001", + "name": "echo", + "args": {"text": "hello"}, + }, + ], + ) + return AIMessage( + content="Synthesis completed with available information", + usage_metadata={ + "input_tokens": 15, + "output_tokens": 10, + "total_tokens": 10, + }, + ) + + mock_model.ainvoke = _mock_ainvoke + mock_model.astream = _mock_ainvoke + return mock_model + + patcher = patch( + "cleveractors.agents.llm.build_chat_model", + side_effect=_build_budget_model, + ) + patcher.start() + self._patches.append(patcher) + + self._executor = create_executor( + config_dict=config, + credentials={"openai": {"api_key": "mock-key"}}, + limits={}, + pricing={}, + ) + + def create_executor_with_pruning_enabled(self) -> None: + """Create an Executor with pruning enabled and a file_read tool.""" + self._teardown_patches() + self._mock_ainvoke_calls = [] + self._pruning_was_called = False + + config: dict[str, Any] = { + "type": "llm", + "name": "prune_test_agent", + "provider": "openai", + "model": "gpt-3.5-turbo", + "config": { + "tools": [{"name": "file_read"}], + "allow_tool_output_pruning": True, + }, + } + + def _build_pruning_aware_model(*args, **kwargs): + mock_model = MagicMock() + mock_model.temperature = 0.7 + + async def _mock_ainvoke(messages, **invoke_kwargs): + self._mock_ainvoke_calls.append({"kwargs": invoke_kwargs}) + call_count = len(self._mock_ainvoke_calls) + if call_count == 1 and invoke_kwargs.get("tools"): + return AIMessage( + content="", + usage_metadata={ + "input_tokens": 30, + "output_tokens": 15, + "total_tokens": 15, + }, + tool_calls=[ + { + "id": "call_prune_001", + "name": "file_read", + "args": {"file": self._temp_file}, + }, + ], + ) + return AIMessage( + content="Final response after pruning", + usage_metadata={ + "input_tokens": 20, + "output_tokens": 12, + "total_tokens": 12, + }, + ) + + mock_model.ainvoke = _mock_ainvoke + mock_model.astream = _mock_ainvoke + return mock_model + + patcher = patch( + "cleveractors.agents.llm.build_chat_model", + side_effect=_build_pruning_aware_model, + ) + patcher.start() + self._patches.append(patcher) + + self._executor = create_executor( + config_dict=config, + credentials={"openai": {"api_key": "mock-key"}}, + limits={}, + pricing={}, + ) + + def create_executor_with_pruning_and_opt_out(self) -> None: + """Create an Executor with pruning enabled; the mock returns + output_prune=false in the tool call arguments.""" + self._teardown_patches() + self._mock_ainvoke_calls = [] + self._pruning_was_called = False + + config: dict[str, Any] = { + "type": "llm", + "name": "optout_test_agent", + "provider": "openai", + "model": "gpt-3.5-turbo", + "config": { + "tools": [{"name": "file_read"}], + "allow_tool_output_pruning": True, + }, + } + + def _build_optout_model(*args, **kwargs): + mock_model = MagicMock() + mock_model.temperature = 0.7 + + async def _mock_ainvoke(messages, **invoke_kwargs): + self._mock_ainvoke_calls.append({"kwargs": invoke_kwargs}) + call_count = len(self._mock_ainvoke_calls) + if call_count == 1 and invoke_kwargs.get("tools"): + return AIMessage( + content="", + usage_metadata={ + "input_tokens": 30, + "output_tokens": 15, + "total_tokens": 15, + }, + tool_calls=[ + { + "id": "call_optout_001", + "name": "file_read", + "args": { + "file": self._temp_file, + "output_prune": False, + }, + }, + ], + ) + return AIMessage( + content="Final answer without pruning", + usage_metadata={ + "input_tokens": 20, + "output_tokens": 12, + "total_tokens": 12, + }, + ) + + mock_model.ainvoke = _mock_ainvoke + mock_model.astream = _mock_ainvoke + return mock_model + + patcher = patch( + "cleveractors.agents.llm.build_chat_model", + side_effect=_build_optout_model, + ) + patcher.start() + self._patches.append(patcher) + + self._executor = create_executor( + config_dict=config, + credentials={"openai": {"api_key": "mock-key"}}, + limits={}, + pricing={}, + ) + + def create_executor_without_budget_or_pruning(self) -> None: + """Create an Executor without any budget/pruning config (backward compat).""" + self._make_basic_executor() + + # โ”€โ”€ Keywords: Execution โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + + def execute_with_message(self, message: str) -> None: + self._execute(message) + + def execute_stream_with_message(self, message: str) -> None: + async def _stream(): + chunks = [] + async for token in self._executor.execute_stream(message): + chunks.append(token) + return "".join(chunks) + + try: + asyncio.run(_stream()) + self._last_result = self._executor.last_result + finally: + self._teardown_patches() + + # โ”€โ”€ Keywords: Assertions โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + + def result_is_valid_actor_result(self) -> None: + assert isinstance(self._last_result, ActorResult), ( + f"Expected ActorResult, got {type(self._last_result)}" + ) + + def result_response_contains(self, text: str) -> None: + assert self._last_result is not None + assert text in self._last_result.response, ( + f"Expected '{text}' in response, got: {self._last_result.response[:200]}" + ) + + def result_has_nodes(self) -> None: + assert self._last_result is not None + assert len(self._last_result.nodes) > 0 + + def result_prompt_tokens_greater_than_zero(self) -> None: + assert self._last_result is not None + assert self._last_result.prompt_tokens > 0 + + def agent_initialized_successfully(self) -> None: + assert self._executor is not None + assert self._executor._agent is not None + + def chat_model_was_invoked(self) -> None: + assert len(self._mock_ainvoke_calls) > 0, "chat_model.ainvoke was never called" + + def chat_model_was_invoked_with_tools(self) -> None: + for call in self._mock_ainvoke_calls: + if call.get("kwargs", {}).get("tools"): + return + raise AssertionError("chat_model.ainvoke was never called with tools kwarg") + + def multiple_rounds_occurred(self) -> None: + assert len(self._mock_ainvoke_calls) >= 2, ( + f"Expected >=2 ainvoke calls, got {len(self._mock_ainvoke_calls)}" + ) diff --git a/robot/llm_token_budget.robot b/robot/llm_token_budget.robot new file mode 100644 index 0000000..c1eda4c --- /dev/null +++ b/robot/llm_token_budget.robot @@ -0,0 +1,53 @@ +*** Settings *** +Documentation Token-budget awareness and tool-output pruning integration tests. +... Exercises the full Executor โ†’ LLMAgent pipeline with mock chat +... models, covering budget exhaustion, pruning pass, opt-out, and +... backward compatibility โ€” no real LLM API key required. +Library TokenBudgetTestLib.py + +*** Test Cases *** +Token Budget Exhaustion Triggers Synthesis Flow + [Documentation] When token_budget_percent is set to near-zero, the first + ... message should trigger the budget-exhausted synthesis flow. + Create Executor With Low Budget For Exhaustion + Execute With Message Tell me about the project + Result Is Valid Actor Result + Chat Model Was Invoked + Multiple Rounds Occurred + Result Response Contains information + +Pruning-Enabled Agent With File Read Tool Initializes Correctly + [Documentation] An agent with allow_tool_output_pruning=true and a + ... file_read tool should initialize and execute without errors. + Create Executor With Pruning Enabled + Execute With Message Read a file for me + Result Is Valid Actor Result + Chat Model Was Invoked With Tools + Result Response Contains Final response after pruning + +output_prune False Skips Pruning + [Documentation] When the mock model returns output_prune=false in the + ... tool call arguments, the raw tool output should be used as-is. + Create Executor With Pruning And Opt Out + Execute With Message Read without pruning + Result Is Valid Actor Result + Chat Model Was Invoked With Tools + Result Response Contains Final answer without pruning + +Token Budget Percent Config Is Accepted + [Documentation] Setting token_budget_percent to a valid value should + ... initialise the agent without errors. + Create Executor With Budget 0.85 + Execute With Message Hello + Result Is Valid Actor Result + Chat Model Was Invoked + Result Prompt Tokens Greater Than Zero + +Backward Compatible When No Budget Or Pruning Config + [Documentation] An agent without token_budget_percent or + ... allow_tool_output_pruning should behave identically to before. + Create Executor Without Budget Or Pruning + Execute With Message Hello + Result Is Valid Actor Result + Chat Model Was Invoked + Result Prompt Tokens Greater Than Zero diff --git a/src/cleveractors/agents/llm.py b/src/cleveractors/agents/llm.py index 7067a93..47c2276 100644 --- a/src/cleveractors/agents/llm.py +++ b/src/cleveractors/agents/llm.py @@ -199,6 +199,59 @@ class LLMAgent(AgentWithMemory): self.max_tokens: int = config.get("max_tokens", DEFAULT_MAX_TOKENS) self.system_message: str = config.get("system_prompt", DEFAULT_SYSTEM_MESSAGE) + # Token-budget awareness (ยง4.4.7). Default None (disabled). + _raw_budget = config.get("token_budget_percent") + if _raw_budget is not None: + if not isinstance(_raw_budget, (int, float)): + raise ConfigurationError( + f"token_budget_percent must be a number, " + f"got {type(_raw_budget).__name__}" + ) + if not 0.0 < _raw_budget <= 1.0: + raise ConfigurationError( + f"token_budget_percent must be in (0.0, 1.0], got {_raw_budget}" + ) + self._token_budget_percent: float | None = ( + float(_raw_budget) if _raw_budget is not None else None + ) + + # Tool output pruning (ยง4.4.8). Default false (opt-in). + self._allow_tool_output_pruning: bool = bool( + config.get("allow_tool_output_pruning", False) + ) + self._pruning_model: str | None = config.get("pruning_model") or None + + # pruning_threshold: character length threshold for pruning (ยง4.4.8 D-2). + # When raw tool output length is at or below this value, no pruning + # occurs even when allow_tool_output_pruning is true and tool matches. + _raw_threshold = config.get("pruning_threshold") + if _raw_threshold is not None: + if not isinstance(_raw_threshold, int) or _raw_threshold <= 0: + raise ConfigurationError( + f"pruning_threshold must be a positive integer, " + f"got {_raw_threshold!r}" + ) + self._pruning_threshold: int = _raw_threshold if _raw_threshold else 512 + + # pruning_tool_filter: list of tool names eligible for pruning (ยง4.4.8 D-2). + # Pruning only applies when the invoked tool name appears in this list. + _raw_filter = config.get("pruning_tool_filter") + if _raw_filter is not None: + if not isinstance(_raw_filter, list): + raise ConfigurationError( + f"pruning_tool_filter must be a list of strings, " + f"got {type(_raw_filter).__name__}" + ) + for item in _raw_filter: + if not isinstance(item, str): + raise ConfigurationError( + f"pruning_tool_filter items must be strings, " + f"got {type(item).__name__}" + ) + self._pruning_tool_filter: list[str] = ( + _raw_filter if _raw_filter is not None else ["file_read", "shell"] + ) + # Tool definitions extracted from agent config. Stored as a list of # dicts in OpenAI-compatible function-calling format (``{"type": # "function", "function": {...}}``). Populated during __init__ only @@ -304,6 +357,184 @@ class LLMAgent(AgentWithMemory): } return default_models.get(self.provider.lower(), DEFAULT_MODEL) + def _get_model_context_window(self) -> int: + """Return the model's advertised context window size in tokens. + + Uses max_tokens from config if set to a value large enough to be a + context window (>10K tokens). Otherwise falls back to a safe default + of 128K tokens. + """ + if self.max_tokens > 10_000: + return self.max_tokens + return 128_000 + + def _estimate_token_count(self, messages: list[Any]) -> int: + """Estimate total token count using get_num_tokens or heuristic fallback. + + Note: the heuristic fallback (chars//4) does not count AIMessage.tool_calls + JSON. Per ADR D-6 this is a known limitation โ€” tool-call payloads can be + hundreds of characters and are excluded from the estimate.""" + try: + _estimator = getattr(self.chat_model, "get_num_tokens", None) + if _estimator is not None: + return int(_estimator(messages)) + except Exception: # nosec B110 โ€” intentional fallback to heuristic + pass + total_chars = 0 + for msg in messages: + total_chars += len(str(getattr(msg, "content", ""))) + return max(1, total_chars // 4) + + def _augment_tool_schemas_for_pruning( + self, tools: list[dict[str, Any]] + ) -> list[dict[str, Any]]: + """Return a deep copy of *tools* with ``output_prune`` and + ``output_prune_context`` injected for tools in pruning_tool_filter.""" + import copy + + augmented: list[dict[str, Any]] = [] + for tool in tools: + t = copy.deepcopy(tool) + fn = t.get("function", {}) + tool_name = fn.get("name", "") if isinstance(fn, dict) else "" + if tool_name in self._pruning_tool_filter: + params = fn.get("parameters", {}) + if isinstance(params, dict): + props = params.setdefault("properties", {}) + if not isinstance(props, dict): + props = {} + params["properties"] = props + props["output_prune"] = { + "type": "boolean", + "description": ( + "Set to false to receive the full raw tool output " + "without any content pruning or extraction." + ), + "default": True, + } + props["output_prune_context"] = { + "type": "string", + "description": ( + "A focused description of what you are looking for " + "in the tool output: what information you need and " + "why. This will be given to the content extraction " + "assistant so it knows what to preserve." + ), + } + fn["parameters"] = params + t["function"] = fn + augmented.append(t) + return augmented + + _PRUNE_INFO_START: str = "[PRUNE_INFO_START]" + _PRUNE_INFO_END: str = "[PRUNE_INFO_END]" + _PRUNE_OUTPUT_START: str = "[PRUNE_OUTPUT_START]" + _PRUNE_OUTPUT_END: str = "[PRUNE_OUTPUT_END]" + + def _parse_pruning_response(self, response_text: str) -> str: + """Parse a pruning model response into ToolMessage content (ยง4.4.9).""" + text = response_text + + def _extract(start_tag: str, end_tag: str, source: str) -> tuple[str, str]: + s = source.find(start_tag) + if s == -1: + return ("", source) + body_start = s + len(start_tag) + e = source.find(end_tag, body_start) + if e == -1: + return ("", source) + body = source[body_start:e].strip() + before = source[:s] + after = source[e + len(end_tag) :] + return (body, before + after) + + info_block, text = _extract(self._PRUNE_INFO_START, self._PRUNE_INFO_END, text) + output_block, text = _extract( + self._PRUNE_OUTPUT_START, self._PRUNE_OUTPUT_END, text + ) + if not output_block and not info_block: + return response_text.strip() + parts: list[str] = [] + if info_block: + parts.append( + f"{self._PRUNE_INFO_START}\n{info_block}\n{self._PRUNE_INFO_END}" + ) + if output_block: + parts.append( + f"{self._PRUNE_OUTPUT_START}\n{output_block}\n{self._PRUNE_OUTPUT_END}" + ) + return "\n".join(parts) + + async def _run_pruning_pass( + self, + tool_name: str, + raw_output: str, + task_context_msgs: list[Any], + prune_context: str | None = None, + ) -> str: + """Execute the tool-output pruning extraction pass (ยง4.4.8). + + When *prune_context* is a non-empty string, it replaces the + broader task context in the extraction prompt. + """ + if not raw_output: + return raw_output + try: + prune_model_name = self._pruning_model or self.model + prune_model = build_chat_model( + provider=self.provider, + config=self.config, + credentials=self._credentials, + model=prune_model_name, + temperature=0.0, + max_tokens=self.max_tokens, + chat_model_globals=globals(), + ) + system_content = ( + "You are a content extraction assistant. Extract only information " + "relevant to the current task from the tool output below. " + "Structure your response with these markers:\n" + f"{self._PRUNE_INFO_START}\n\n" + f"{self._PRUNE_INFO_END}\n" + f"{self._PRUNE_OUTPUT_START}\n\n" + f"{self._PRUNE_OUTPUT_END}" + ) + from langchain_core.messages import SystemMessage as _SM + + if prune_context and prune_context.strip(): + prune_messages: list[Any] = [ + _SM(content=system_content), + HumanMessage( + content=( + f"Search directive: {prune_context.strip()}\n\n" + f"Tool: {tool_name}\n\n" + f"Tool output:\n{raw_output}" + ) + ), + ] + else: + import copy + + user_msgs = copy.deepcopy(task_context_msgs) + if not isinstance(user_msgs, list): + user_msgs = [user_msgs] + user_msgs.append( + HumanMessage( + content=(f"Tool: {tool_name}\n\nTool output:\n{raw_output}") + ) + ) + prune_messages = [_SM(content=system_content)] + user_msgs + prune_response = await prune_model.ainvoke(prune_messages) + return self._parse_pruning_response(str(prune_response.content)) + except Exception: + logger.warning( + "Agent %s: pruning pass failed for tool %r; using raw output", + self.name, + tool_name, + exc_info=True, + ) + return raw_output + # ------------------------------------------------------------------ # Public credentials property # ------------------------------------------------------------------ @@ -504,12 +735,187 @@ class LLMAgent(AgentWithMemory): ) from _mre _has_tools = self._lc_tools is not None and LANGCHAIN_AVAILABLE if _has_tools: - invoke_kwargs["tools"] = self._lc_tools + # When pruning is enabled, augment each tool's schema + # with the ``output_prune`` meta-parameter (ยง4.4.8). + if self._allow_tool_output_pruning: + _all_functions = all( + isinstance(t, dict) and "function" in t for t in self._lc_tools + ) + if _all_functions: + invoke_kwargs["tools"] = self._augment_tool_schemas_for_pruning( + self._lc_tools + ) + else: + invoke_kwargs["tools"] = self._lc_tools + else: + invoke_kwargs["tools"] = self._lc_tools + + _budget_exhausted: bool = False for _tool_round in range(_TOOL_MAX_ROUNDS): if _tool_round > 0 and not _has_tools: break + # โ”€โ”€ Token-budget check (ยง4.4.7) โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + if self._token_budget_percent is not None and _has_tools: + _budget_ceiling = max( + 1, + int( + self._token_budget_percent + * self._get_model_context_window() + ), + ) + _est = self._estimate_token_count(messages) + _remaining = max(0, _budget_ceiling - _est) + if _est > int(0.75 * _budget_ceiling): + logger.warning( + "Agent %s: token budget at %.0f%% " + "(%d / %d tokens used, %d remaining)", + self.name, + (_est / _budget_ceiling) * 100, + _est, + _budget_ceiling, + _remaining, + ) + if _est > _budget_ceiling: + logger.error( + "Agent %s: token budget exhausted " + "(%d / %d tokens); triggering synthesis flow", + self.name, + _est, + _budget_ceiling, + ) + _synthesis_text = ( + "You have nearly exhausted the available " + "context window. Produce your final answer " + "based on all the information already " + "gathered. If you must make one more tool " + "call to complete your answer, do it now." + ) + messages.append(HumanMessage(content=_synthesis_text)) + response = await self.chat_model.ainvoke( + messages, **invoke_kwargs + ) + _synth_tool_calls: list[dict[str, Any]] = ( + getattr(response, "tool_calls", None) or [] + ) + if ( + isinstance(_synth_tool_calls, list) + and _synth_tool_calls + and _has_tools + ): + messages.append(response) + for tc in _synth_tool_calls: + call_id = tc.get("id", "") + fn_def = tc.get("function") + if isinstance(fn_def, dict): + tool_name = fn_def.get("name", "") + arguments_raw = fn_def.get("arguments") + else: + tool_name = tc.get("name", "") + arguments_raw = tc.get("args") + if not tool_name: + messages.append( + ToolMessage( + content="Tool name is empty", + tool_call_id=call_id, + ) + ) + continue + args = {} + if isinstance(arguments_raw, str): + try: + args = json.loads(arguments_raw) + except (json.JSONDecodeError, ValueError): + args = {"_raw": arguments_raw} + elif isinstance(arguments_raw, dict): + args = arguments_raw + _bq_output_prune = args.pop("output_prune", None) + _bq_prune_ctx = args.pop("output_prune_context", None) + try: + from cleveractors.agents.tool import ( + ToolAgent as _TA, + ) + + parent_unsafe = self.config.get( + "unsafe_mode", False + ) + tool_cfg: dict[str, Any] = { + "tools": [{"name": tool_name}], + "safe_mode": not parent_unsafe, + "allow_shell": self.config.get( + "allow_shell", False + ), + "exec_python": self.config.get( + "exec_python", False + ), + "timeout": self.config.get("timeout", 1), + } + t_agent = _TA( + name=f"_tc_synth_budget_{call_id}", + config=tool_cfg, + template_renderer=self.template_renderer, + ) + t_ctx = ( + {"_unsafe_mode": True} + if parent_unsafe + else None + ) + raw_out = await t_agent.process_message( + json.dumps({"tool": tool_name, "args": args}) + if isinstance(args, dict) + else arguments_raw or "", + context=t_ctx, + ) + if ( + self._allow_tool_output_pruning + and _bq_output_prune is not False + and tool_name in self._pruning_tool_filter + and len(raw_out) > self._pruning_threshold + ): + tool_output = await self._run_pruning_pass( + tool_name, + raw_out, + list(messages), + prune_context=_bq_prune_ctx, + ) + logger.info( + "Agent %s: pruning pass for %r: " + "%d -> %d chars", + self.name, + tool_name, + len(raw_out), + len(tool_output), + ) + else: + tool_output = raw_out + messages.append( + ToolMessage( + content=tool_output, + tool_call_id=call_id, + ) + ) + except (ExecutionError, ConfigurationError) as _se: + _se_msg = str(_se) + logger.warning( + "Agent %s: budget synth tool call " + "failed for %r (tool=%s): %s", + self.name, + call_id, + tool_name, + _se_msg, + ) + messages.append( + ToolMessage( + content=f"Tool '{tool_name}' error: {_se_msg}", + tool_call_id=call_id, + ) + ) + response = await self.chat_model.ainvoke(messages) + response_text = str(response.content) + _budget_exhausted = True + break + response = await self.chat_model.ainvoke(messages, **invoke_kwargs) response_tool_calls: list[dict[str, Any]] = ( @@ -525,8 +931,6 @@ class LLMAgent(AgentWithMemory): for tc in response_tool_calls: call_id = tc.get("id", "") - # OpenAI standard: {"function": {"name": "...", "arguments": "..."}} - # HF / non-standard: {"name": "...", "args": {...}} fn_def = tc.get("function") if isinstance(fn_def, dict): tool_name = fn_def.get("name", "") @@ -550,14 +954,11 @@ class LLMAgent(AgentWithMemory): args = {"_raw": arguments_raw} elif isinstance(arguments_raw, dict): args = arguments_raw + _output_prune = args.pop("output_prune", None) + _prune_context = args.pop("output_prune_context", None) try: from cleveractors.agents.tool import ToolAgent as _TA - # Propagate tool-execution hints from the LLM - # agent's config (ยง4.5.4). The graph YAML places - # "unsafe_mode", "allow_shell", and "timeout" - # directly on the LLM agent config; we translate - # them into ToolAgent-compatible equivalents. parent_unsafe = self.config.get("unsafe_mode", False) tool_config: dict[str, Any] = { "tools": [{"name": tool_name}], @@ -566,7 +967,6 @@ class LLMAgent(AgentWithMemory): "exec_python": self.config.get("exec_python", False), "timeout": self.config.get("timeout", 1), } - agent = _TA( name=f"_tc_{call_id}", config=tool_config, @@ -575,14 +975,41 @@ class LLMAgent(AgentWithMemory): tool_ctx: dict[str, Any] | None = ( {"_unsafe_mode": True} if parent_unsafe else None ) - tool_result = await agent.process_message( + raw_tool_output = await agent.process_message( json.dumps({"tool": tool_name, "args": args}) if isinstance(args, dict) else arguments_raw or "", context=tool_ctx, ) + # โ”€โ”€ Tool-output pruning pass (ยง4.4.8) โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + # Pruning proceeds only when all conditions are met: + # allow_tool_output_pruning is true, tool is in + # pruning_tool_filter, output exceeds threshold, and + # output_prune is not explicitly false. + if ( + self._allow_tool_output_pruning + and _output_prune is not False + and tool_name in self._pruning_tool_filter + and len(raw_tool_output) > self._pruning_threshold + ): + _task_context = list(messages) + tool_output = await self._run_pruning_pass( + tool_name, + raw_tool_output, + _task_context, + prune_context=_prune_context, + ) + logger.info( + "Agent %s: pruning pass for %r: %d -> %d chars", + self.name, + tool_name, + len(raw_tool_output), + len(tool_output), + ) + else: + tool_output = raw_tool_output messages.append( - ToolMessage(content=tool_result, tool_call_id=call_id) + ToolMessage(content=tool_output, tool_call_id=call_id) ) except (ExecutionError, ConfigurationError) as _tool_err: _tool_err_msg = str(_tool_err) @@ -602,10 +1029,14 @@ class LLMAgent(AgentWithMemory): response_text: str = str(response.content) # If the tool loop exhausted but the model produced no meaningful - # content (stuck in tool-only mode), ask it to synthesize output - # from all gathered context. This happens with some models that - # never stop making tool calls on their own. - if not response_text.strip() and _has_tools and len(messages) > 2: + # content (stuck in tool-only mode), ask it to synthesize output. + # Skip when budget-exhausted synthesis already handled termination. + if ( + not _budget_exhausted + and not response_text.strip() + and _has_tools + and len(messages) > 2 + ): messages.append( HumanMessage( content=(