feat(tool-loop): add token-budget awareness and tool output pruning #62

Merged
CoreRasurae merged 1 commits from feature/m2-tool-budget into master 2026-06-26 17:41:56 +00:00
9 changed files with 2447 additions and 15 deletions
+10
View File
@@ -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.
+249
View File
@@ -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)
)
@@ -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.
+8
View File
@@ -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
+265
View File
@@ -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
Outdated
Review

Major (test quality): This scenario is tautological. The step And I augment the tool schemas for pruning calls _augment_tool_schemas_for_pruning directly, which always augments regardless of the agent's allow_tool_output_pruning config. The assertion would pass even if the agent's process_message ignored the allow_tool_output_pruning flag entirely. To meaningfully test the agent's behavior, mock the chat model, call process_message, and inspect the tools kwarg actually passed to ainvoke() — assert that it contains the output_prune parameter when pruning is enabled.

**Major (test quality):** This scenario is tautological. The step `And I augment the tool schemas for pruning` calls `_augment_tool_schemas_for_pruning` directly, which always augments regardless of the agent's `allow_tool_output_pruning` config. The assertion would pass even if the agent's `process_message` ignored the `allow_tool_output_pruning` flag entirely. To meaningfully test the agent's behavior, mock the chat model, call `process_message`, and inspect the `tools` kwarg actually passed to `ainvoke()` — assert that it contains the `output_prune` parameter when pruning is enabled.
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
Outdated
Review

Major (test quality): This scenario is contradictory and tautological. The title says "not augmented when pruning disabled" but the assertion (line 137) checks that the augmented schemas DO contain output_prune. Furthermore, the step I augment the tool schemas for pruning calls _augment_tool_schemas_for_pruning directly — the function always augments, so the test passes whether allow_tool_output_pruning is true or false. Replace with a test that mocks the chat model, calls process_message, and inspects the tools kwarg to assert the schemas are NOT augmented when allow_tool_output_pruning is false (or omitted).

**Major (test quality):** This scenario is contradictory and tautological. The title says "not augmented when pruning disabled" but the assertion (line 137) checks that the augmented schemas DO contain `output_prune`. Furthermore, the step `I augment the tool schemas for pruning` calls `_augment_tool_schemas_for_pruning` directly — the function always augments, so the test passes whether `allow_tool_output_pruning` is true or false. Replace with a test that mocks the chat model, calls `process_message`, and inspects the `tools` kwarg to assert the schemas are NOT augmented when `allow_tool_output_pruning` is `false` (or omitted).
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
Outdated
Review

Major (test quality): This scenario does not exercise the opt-out code path. The mock returns a tool call for echo, but the pruning check at line 965 already gates on tool_name == "file_read", so the pruning pass is skipped regardless of the output_prune: false arg. The test would still pass if the _output_prune is not False guard were removed entirely. Change the tool to file_read (or another in-scope tool) so the opt-out branch is actually exercised.

**Major (test quality):** This scenario does not exercise the opt-out code path. The mock returns a tool call for `echo`, but the pruning check at line 965 already gates on `tool_name == "file_read"`, so the pruning pass is skipped regardless of the `output_prune: false` arg. The test would still pass if the `_output_prune is not False` guard were removed entirely. Change the tool to `file_read` (or another in-scope tool) so the opt-out branch is actually exercised.
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"
@@ -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
+381
View File
@@ -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)}"
)
+53
View File
@@ -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
+446 -15
View File
@@ -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<describe what was pruned>\n"
f"{self._PRUNE_INFO_END}\n"
f"{self._PRUNE_OUTPUT_START}\n<extracted content>\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")
Outdated
Review

Major: This args.pop("output_prune", None) discards the return value, so the LLM's output_prune: false opt-out cannot be detected downstream. The main tool loop at lines 932-933 correctly captures _output_prune = args.pop("output_prune", None) and the pruning check at line 964 gates on _output_prune is not False. Mirror that pattern here to satisfy ticket AC6.

**Major:** This `args.pop("output_prune", None)` discards the return value, so the LLM's `output_prune: false` opt-out cannot be detected downstream. The main tool loop at lines 932-933 correctly captures `_output_prune = args.pop("output_prune", None)` and the pruning check at line 964 gates on `_output_prune is not False`. Mirror that pattern here to satisfy ticket AC6.
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
Outdated
Review

Major: The pruning condition here is missing the _output_prune is not False guard that the main tool loop applies at line 964. As a result, when the LLM calls file_read with output_prune: false during the budget-exhausted synthesis round, the pruning pass still runs — silently overriding the LLM's opt-out and violating ticket AC6. Fix: capture _bq_output_prune = args.pop("output_prune", None) (see comment on line 810) and add _bq_output_prune is not False to this condition.

**Major:** The pruning condition here is missing the `_output_prune is not False` guard that the main tool loop applies at line 964. As a result, when the LLM calls `file_read` with `output_prune: false` during the budget-exhausted synthesis round, the pruning pass still runs — silently overriding the LLM's opt-out and violating ticket AC6. Fix: capture `_bq_output_prune = args.pop("output_prune", None)` (see comment on line 810) and add `_bq_output_prune is not False` to this condition.
),
"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=(