CoreRasurae
964e87cc0f
CI / lint (pull_request) Successful in 1m31s
CI / typecheck (pull_request) Successful in 50s
CI / security (pull_request) Successful in 50s
CI / quality (pull_request) Successful in 33s
CI / unit_tests (pull_request) Successful in 3m5s
CI / integration_tests (pull_request) Successful in 1m7s
CI / build (pull_request) Successful in 34s
CI / coverage (pull_request) Successful in 3m7s
CI / status-check (pull_request) Successful in 3s
CI / lint (push) Successful in 33s
CI / typecheck (push) Successful in 50s
CI / security (push) Successful in 50s
CI / quality (push) Successful in 33s
CI / integration_tests (push) Successful in 1m12s
CI / build (push) Successful in 33s
CI / unit_tests (push) Successful in 3m5s
CI / coverage (push) Successful in 3m6s
CI / status-check (push) Successful in 6s
Implement token-budget awareness (§4.4.7) and tool output pruning (§4.4.8-9) for the multi-turn tool-call loop to prevent context-window exhaustion. Token budget: tracks estimated token count before each ainvoke(), emits warning at 75% budget, injects synthesis prompt at exhaustion with one final tool-call round. Estimation via chat_model.get_num_tokens() with len(content)//4 heuristic fallback. Config key: token_budget_percent. Tool output pruning: implicit LLM extraction pass scoped to file_read calls. Schema augmentation injects output_prune (bool) and output_prune_context (string) meta-arguments. Pruning model responds with [PRUNE_INFO_START/END] and [PRUNE_OUTPUT_START/END] markers. Config keys: allow_tool_output_pruning, pruning_model. - Capture output_prune value in budget synthesis flow with proper guard - Fix ADR D-4: synthesis prompt is HumanMessage not SystemMessage - Fix Robot tests: replace /dev/null with temp file path ADR-2031 documents all specification extensions. 21 new BDD scenarios, 5 Robot integration tests, 16 ASV benchmarks. ISSUES CLOSED: #61
382 lines
14 KiB
Python
382 lines
14 KiB
Python
"""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)}"
|
|
)
|