diff --git a/features/llm_agent_stream_tool_calls.feature b/features/llm_agent_stream_tool_calls.feature new file mode 100644 index 0000000..dd71b6e --- /dev/null +++ b/features/llm_agent_stream_tool_calls.feature @@ -0,0 +1,48 @@ +Feature: LLMAgent.stream_message() with tool-call support (issue #67) + + # ── Tools configured: single-chunk output ──────────────────────────────── + + Scenario: stream_message yields one non-empty chunk when tools configured + Given an LLMAgent with tools configured and mock that does one tool round + When I call stream_message on the tool-configured agent + Then the yielded tokens should contain exactly one non-empty string + And _last_token_usage should have non-zero accumulated counts + + Scenario: stream_message last_result token counts equal accumulated totals from tool loop + Given an LLMAgent with tools configured and two-round tool mock + When I call stream_message on the tool-configured agent + Then the accumulated prompt tokens from stream_message should equal two rounds sum + And the accumulated completion tokens from stream_message should equal two rounds sum + + Scenario: stream_message on tool actor yields final answer text as single chunk + Given an LLMAgent with tools configured and mock returning "Tool loop final answer" + When I call stream_message on the tool-configured agent + Then the single yielded chunk should be "Tool loop final answer" + + # ── No-tools configured: real astream() path unchanged ─────────────────── + + Scenario: stream_message on no-tools actor still uses real astream token-by-token + Given an LLMAgent with no tools and a mock astream yielding "tok1", "tok2", "tok3" + When I call stream_message on the no-tools agent + Then the no-tools stream yielded tokens tok1 tok2 tok3 + And _last_token_usage should be (0, 0) for no-usage mock + + Scenario: stream_message on no-tools actor captures token counts from astream final chunk + Given an LLMAgent with no tools and mock astream with usage prompt=7 completion=13 + When I call stream_message on the no-tools agent + Then _last_token_usage for no-tools stream should be (7, 13) + + # ── Billing integrity on exception ─────────────────────────────────────── + + Scenario: stream_message billing integrity preserved on tool loop exception + Given an LLMAgent with tools configured and mock that completes one round then errors + When I attempt stream_message on the tool-configured agent + Then an ExecutionError should be raised from stream_message + And _last_token_usage should have the partial accumulated counts + + # ── Memory update for tool path ─────────────────────────────────────────── + + Scenario: stream_message with tools and memory_enabled updates conversation history + Given an LLMAgent with tools and memory_enabled and mock returning "Memory answer" + When I call stream_message on the tool-configured agent + Then the conversation memory should contain the tool answer diff --git a/features/llm_agent_tool_loop.feature b/features/llm_agent_tool_loop.feature new file mode 100644 index 0000000..cbe872b --- /dev/null +++ b/features/llm_agent_tool_loop.feature @@ -0,0 +1,97 @@ +Feature: LLMAgent._execute_tool_loop() — shared tool-call orchestration helper + + # ── Basic single tool round ─────────────────────────────────────────────── + + Scenario: _execute_tool_loop returns result for single tool round + Given I have an LLMAgent with tool loop test config for tool "echo" + And I set a mock that makes one tool call then returns a final answer "Final answer from loop" + When I run _execute_tool_loop with an initial messages list + Then the loop result final_response content should be "Final answer from loop" + And the loop result accumulated_prompt should be greater than 0 + And the loop result accumulated_completion should be greater than 0 + And the loop result budget_exhausted should be false + And the loop result synthesis_was_run should be false + + # ── Multi-round (two consecutive tool-call rounds) ──────────────────────── + + Scenario: _execute_tool_loop accumulates tokens across two tool-call rounds + Given I have an LLMAgent with tool loop test config for tool "echo" + And I set a mock that makes two tool calls then returns a final answer "Final answer two rounds" + When I run _execute_tool_loop with an initial messages list + Then the loop result final_response content should be "Final answer two rounds" + And the loop result accumulated_prompt should be at least 30 + And the loop result accumulated_completion should be at least 15 + + # ── tool_max_rounds exhaustion ──────────────────────────────────────────── + + Scenario: _execute_tool_loop handles tool_max_rounds exhaustion + Given I have an LLMAgent with tool loop test config for tool "echo" and tool_max_rounds 1 + And I set a mock that always returns tool_calls never a final answer + When I run _execute_tool_loop with an initial messages list + Then the loop result should contain some content from the synthesis prompt + And the loop result budget_exhausted should be false + + # ── Token accumulation spanning all rounds ──────────────────────────────── + + Scenario: _execute_tool_loop token counts span all ainvoke rounds + Given I have an LLMAgent with tool loop test config for tool "echo" + And I set a mock that makes two tool calls then returns a final answer "Accumulated answer" + When I run _execute_tool_loop with an initial messages list + Then the loop result accumulated_prompt should equal the sum of all rounds prompt tokens + And the loop result accumulated_completion should equal the sum of all rounds completion tokens + + # ── Token-budget exhaustion ─────────────────────────────────────────────── + + Scenario: _execute_tool_loop triggers synthesis when token budget exhausted + Given I have an LLMAgent with tool loop test config for tool "echo" and tiny token_budget_percent + And I set a mock that triggers budget exhaustion then returns "Budget exhausted answer" + When I run _execute_tool_loop with an initial messages list + Then the loop result budget_exhausted should be true + And the loop result final_response content should be "Budget exhausted answer" + + # ── Stuck-model synthesis ───────────────────────────────────────────────── + + Scenario: _execute_tool_loop runs synthesis when model is stuck in tool-only mode + Given I have an LLMAgent with tool loop test config for tool "echo" and tool_max_rounds 2 + And I set a mock that returns empty content tool calls then "Synthesis answer" + When I run _execute_tool_loop with an initial messages list + Then the loop result synthesis_was_run should be true + And the loop result final_response content should be "Synthesis answer" + + # ── Tool-output pruning ─────────────────────────────────────────────────── + + Scenario: _execute_tool_loop runs pruning pass and includes its tokens in accumulated counts + Given I have an LLMAgent with tool loop test config for tool "file_read" and pruning enabled + And I set a mock that makes one file_read tool call then returns "Pruned answer" + And I set a mock pruning model that adds extra tokens + When I run _execute_tool_loop with an initial messages list + Then the loop result final_response content should be "Pruned answer" + And the loop result accumulated_prompt should include pruning tokens + + # ── Tool dispatch error ─────────────────────────────────────────────────── + + Scenario: _execute_tool_loop captures tool dispatch ExecutionError as ToolMessage + Given I have an LLMAgent with tool loop test config for tool "echo" + And I set a mock that calls a nonexistent tool then returns "After error answer" + When I run _execute_tool_loop with an initial messages list + Then the loop result final_response content should be "After error answer" + And the messages list should contain a ToolMessage with not available content + + # ── Exception path: partial token counts preserved ──────────────────────── + + Scenario: _execute_tool_loop wraps exception in _ToolLoopError with partial counts + Given I have an LLMAgent with tool loop test config for tool "echo" + And I set a mock that returns one tool round then raises a RuntimeError + When I attempt to run _execute_tool_loop with an initial messages list + Then a _ToolLoopError should be raised + And the _ToolLoopError any_invocation_made should be true + And the _ToolLoopError accumulated_prompt should be greater than 0 + + # ── Exception before any invocation ────────────────────────────────────── + + Scenario: _ToolLoopError has any_invocation_made false when error before first ainvoke + Given I have an LLMAgent with tool loop test config for tool "echo" and bad tool_max_rounds + When I attempt to run _execute_tool_loop with an initial messages list + Then a _ToolLoopError should be raised + And the _ToolLoopError any_invocation_made should be false + And the _ToolLoopError cause should be a ConfigurationError diff --git a/features/steps/llm_agent_stream_tool_calls_steps.py b/features/steps/llm_agent_stream_tool_calls_steps.py new file mode 100644 index 0000000..688849c --- /dev/null +++ b/features/steps/llm_agent_stream_tool_calls_steps.py @@ -0,0 +1,379 @@ +"""Step definitions for LLMAgent.stream_message() with tool-call support tests.""" + +from __future__ import annotations + +import asyncio +from typing import Any +from unittest.mock import Mock + +from behave import given, then, when +from langchain_core.messages import AIMessage + +from cleveractors.agents.llm import LLMAgent, last_token_usage_var +from cleveractors.core.exceptions import ExecutionError +from cleveractors.templates.renderer import TemplateRenderer + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_agent(config: dict[str, Any]) -> LLMAgent: + renderer = Mock(spec=TemplateRenderer) + renderer.render_string.return_value = "system prompt" + return LLMAgent( + name="stream_tool_agent", + config=config, + template_renderer=renderer, + ) + + +def _make_ai_message( + content: str, + tool_calls: list[dict[str, Any]] | None = None, + prompt: int = 10, + completion: int = 5, +) -> AIMessage: + return AIMessage( + content=content, + usage_metadata={ + "input_tokens": prompt, + "output_tokens": completion, + "total_tokens": prompt + completion, + }, + tool_calls=tool_calls or [], + ) + + +async def _collect_stream_message(agent: LLMAgent, message: str = "test") -> list[str]: + tokens: list[str] = [] + async for token in agent.stream_message(message): + tokens.append(token) + return tokens + + +# --------------------------------------------------------------------------- +# Given steps +# --------------------------------------------------------------------------- + + +@given("an LLMAgent with tools configured and mock that does one tool round") +def step_tool_stream_one_round(context): + call_counter = [0] + context._stream_tool_round_prompts = [15, 20] + context._stream_tool_round_completions = [5, 10] + + async def _ainvoke(messages, **kwargs): + call_counter[0] += 1 + if call_counter[0] == 1: + return _make_ai_message( + content="", + tool_calls=[ + {"id": "call_001", "name": "echo", "args": {"message": "hi"}} + ], + prompt=15, + completion=5, + ) + return _make_ai_message(content="Final stream answer", prompt=20, completion=10) + + context.stream_tool_agent_config = { + "name": "stream_tool_agent", + "provider": "openai", + "api_key": "test", + "tools": [{"name": "echo"}], + "tool_max_rounds": 5, + } + context.stream_tool_mock_ainvoke = _ainvoke + + +@given("an LLMAgent with tools configured and two-round tool mock") +def step_tool_stream_two_rounds(context): + call_counter = [0] + context._stream_two_round_prompts = [15, 15, 20] + context._stream_two_round_completions = [5, 5, 10] + + async def _ainvoke(messages, **kwargs): + call_counter[0] += 1 + if call_counter[0] <= 2: + return _make_ai_message( + content="", + tool_calls=[ + { + "id": f"call_{call_counter[0]:03d}", + "name": "echo", + "args": {"message": "input"}, + } + ], + prompt=15, + completion=5, + ) + return _make_ai_message(content="Two round answer", prompt=20, completion=10) + + context.stream_tool_agent_config = { + "name": "stream_tool_agent", + "provider": "openai", + "api_key": "test", + "tools": [{"name": "echo"}], + "tool_max_rounds": 5, + } + context.stream_tool_mock_ainvoke = _ainvoke + + +@given('an LLMAgent with tools configured and mock returning "{final_text}"') +def step_tool_stream_mock_final(context, final_text): + async def _ainvoke(messages, **kwargs): + return _make_ai_message(content=final_text, prompt=10, completion=5) + + context.stream_tool_agent_config = { + "name": "stream_tool_agent", + "provider": "openai", + "api_key": "test", + "tools": [{"name": "echo"}], + "tool_max_rounds": 1, + } + context.stream_tool_mock_ainvoke = _ainvoke + + +@given('an LLMAgent with no tools and a mock astream yielding "tok1", "tok2", "tok3"') +def step_no_tools_astream_three_tokens(context): + from types import SimpleNamespace + + async def _astream(messages, **kwargs): + for t in ["tok1", "tok2", "tok3"]: + chunk = SimpleNamespace(content=t) + yield chunk + + context.stream_no_tools_config = { + "name": "stream_no_tools_agent", + "provider": "openai", + "api_key": "test", + } + context.stream_no_tools_mock_astream = _astream + + +@given("an LLMAgent with no tools and mock astream with usage prompt=7 completion=13") +def step_no_tools_astream_with_usage(context): + async def _astream(messages, **kwargs): + yield AIMessage( + content="answer", + usage_metadata={ + "input_tokens": 7, + "output_tokens": 13, + "total_tokens": 20, + }, + ) + + context.stream_no_tools_config = { + "name": "stream_no_tools_agent", + "provider": "openai", + "api_key": "test", + } + context.stream_no_tools_mock_astream = _astream + + +@given( + "an LLMAgent with tools configured and mock that completes one round then errors" +) +def step_tool_stream_one_round_then_error(context): + call_counter = [0] + + async def _ainvoke(messages, **kwargs): + call_counter[0] += 1 + if call_counter[0] == 1: + return _make_ai_message( + content="", + tool_calls=[ + {"id": "call_001", "name": "echo", "args": {"message": "hi"}} + ], + prompt=20, + completion=8, + ) + raise RuntimeError("Simulated failure after first tool round") + + context.stream_tool_agent_config = { + "name": "stream_tool_agent", + "provider": "openai", + "api_key": "test", + "tools": [{"name": "echo"}], + "tool_max_rounds": 5, + } + context.stream_tool_mock_ainvoke = _ainvoke + context._stream_tool_partial_prompt = 20 + context._stream_tool_partial_completion = 8 + + +@given('an LLMAgent with tools and memory_enabled and mock returning "{final_text}"') +def step_tool_stream_memory_enabled(context, final_text): + async def _ainvoke(messages, **kwargs): + return _make_ai_message(content=final_text, prompt=12, completion=6) + + context.stream_tool_agent_config = { + "name": "stream_tool_agent", + "provider": "openai", + "api_key": "test", + "tools": [{"name": "echo"}], + "tool_max_rounds": 1, + "memory_enabled": True, + } + context.stream_tool_mock_ainvoke = _ainvoke + context._stream_tool_memory_answer = final_text + + +# --------------------------------------------------------------------------- +# When steps +# --------------------------------------------------------------------------- + + +@when("I call stream_message on the tool-configured agent") +def step_call_stream_message_tool(context): + agent = _make_agent(context.stream_tool_agent_config) + mock_model = Mock() + mock_model.temperature = 0.7 + mock_model.ainvoke = context.stream_tool_mock_ainvoke + agent.chat_model = mock_model + + context.stream_tool_tokens = asyncio.run(_collect_stream_message(agent)) + context.stream_tool_agent = agent + context.stream_tool_error = None + + +@when("I call stream_message on the no-tools agent") +def step_call_stream_message_no_tools(context): + agent = _make_agent(context.stream_no_tools_config) + mock_model = Mock() + mock_model.temperature = 0.7 + mock_model.astream = context.stream_no_tools_mock_astream + agent.chat_model = mock_model + + context.stream_no_tools_tokens = asyncio.run(_collect_stream_message(agent)) + context.stream_no_tools_agent = agent + + +@when("I attempt stream_message on the tool-configured agent") +def step_attempt_stream_message_tool(context): + agent = _make_agent(context.stream_tool_agent_config) + mock_model = Mock() + mock_model.temperature = 0.7 + mock_model.ainvoke = context.stream_tool_mock_ainvoke + agent.chat_model = mock_model + + context.stream_tool_agent = agent + context.stream_tool_tokens = [] + context.stream_tool_error = None + + async def _run(): + async for token in agent.stream_message("test"): + context.stream_tool_tokens.append(token) + + try: + asyncio.run(_run()) + except ExecutionError as e: + context.stream_tool_error = e + + +# --------------------------------------------------------------------------- +# Then steps +# --------------------------------------------------------------------------- + + +@then("the yielded tokens should contain exactly one non-empty string") +def step_one_non_empty_token(context): + tokens = context.stream_tool_tokens + non_empty = [t for t in tokens if t.strip()] + assert len(non_empty) == 1, ( + f"Expected exactly 1 non-empty token, got {len(non_empty)}: {tokens}" + ) + + +@then("_last_token_usage should have non-zero accumulated counts") +def step_last_token_usage_non_zero(context): + agent = context.stream_tool_agent + prompt, completion = agent._last_token_usage + assert prompt > 0 or completion > 0, ( + f"Expected non-zero token usage, got ({prompt}, {completion})" + ) + + +@then("the accumulated prompt tokens from stream_message should equal two rounds sum") +def step_stream_prompt_equals_two_rounds(context): + agent = context.stream_tool_agent + prompt, _ = agent._last_token_usage + expected = sum(context._stream_two_round_prompts) + assert prompt == expected, f"Expected prompt={expected}, got {prompt}" + + +@then( + "the accumulated completion tokens from stream_message should equal two rounds sum" +) +def step_stream_completion_equals_two_rounds(context): + agent = context.stream_tool_agent + _, completion = agent._last_token_usage + expected = sum(context._stream_two_round_completions) + assert completion == expected, f"Expected completion={expected}, got {completion}" + + +@then('the single yielded chunk should be "{expected}"') +def step_single_chunk_equals(context, expected): + tokens = context.stream_tool_tokens + combined = "".join(tokens) + assert combined == expected, ( + f"Expected combined tokens={expected!r}, got {combined!r}" + ) + + +@then("the no-tools stream yielded tokens tok1 tok2 tok3") +def step_three_tokens(context): + assert context.stream_no_tools_tokens == ["tok1", "tok2", "tok3"], ( + f"Got {context.stream_no_tools_tokens}" + ) + + +@then("_last_token_usage should be (0, 0) for no-usage mock") +def step_no_usage_zero(context): + agent = context.stream_no_tools_agent + assert agent._last_token_usage == (0, 0), ( + f"Expected (0,0), got {agent._last_token_usage}" + ) + + +@then("_last_token_usage for no-tools stream should be (7, 13)") +def step_no_tools_token_counts(context): + agent = context.stream_no_tools_agent + assert agent._last_token_usage == (7, 13), ( + f"Expected (7,13), got {agent._last_token_usage}" + ) + + +@then("an ExecutionError should be raised from stream_message") +def step_execution_error_raised(context): + assert isinstance(context.stream_tool_error, ExecutionError), ( + f"Expected ExecutionError, got {type(context.stream_tool_error)}" + ) + + +@then("_last_token_usage should have the partial accumulated counts") +def step_partial_accumulated_counts(context): + agent = context.stream_tool_agent + prompt, completion = agent._last_token_usage + expected_prompt = context._stream_tool_partial_prompt + expected_completion = context._stream_tool_partial_completion + assert prompt == expected_prompt and completion == expected_completion, ( + f"Expected ({expected_prompt},{expected_completion}), got ({prompt},{completion})" + ) + + +@then("the conversation memory should contain the tool answer") +def step_memory_contains_answer(context): + agent = context.stream_tool_agent + + async def _check_memory(): + history = await agent.get_memory("conversation_history", []) + return history + + history = asyncio.run(_check_memory()) + answer = context._stream_tool_memory_answer + found = any( + msg.get("role") == "assistant" and msg.get("content") == answer + for msg in history + ) + assert found, f"Expected assistant message with {answer!r} in memory, got {history}" diff --git a/features/steps/llm_agent_tool_loop_steps.py b/features/steps/llm_agent_tool_loop_steps.py new file mode 100644 index 0000000..7ce502c --- /dev/null +++ b/features/steps/llm_agent_tool_loop_steps.py @@ -0,0 +1,601 @@ +"""Step definitions for LLMAgent._execute_tool_loop() BDD tests.""" + +from __future__ import annotations + +import asyncio +import json +from typing import Any +from unittest.mock import Mock, patch + +from behave import given, then, when +from langchain_core.messages import AIMessage, ToolMessage + +from cleveractors.agents.llm import LLMAgent, _ToolLoopError +from cleveractors.core.exceptions import ConfigurationError +from cleveractors.templates.renderer import TemplateRenderer + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_agent(config: dict[str, Any]) -> LLMAgent: + """Create an LLMAgent with a mock TemplateRenderer.""" + renderer = Mock(spec=TemplateRenderer) + renderer.render_string.return_value = "system prompt" + return LLMAgent( + name="tool_loop_agent", + config=config, + template_renderer=renderer, + ) + + +def _make_ai_message( + content: str, + tool_calls: list[dict[str, Any]] | None = None, + prompt: int = 10, + completion: int = 5, +) -> AIMessage: + """Return an AIMessage with token usage metadata.""" + return AIMessage( + content=content, + usage_metadata={ + "input_tokens": prompt, + "output_tokens": completion, + "total_tokens": prompt + completion, + }, + tool_calls=tool_calls or [], + ) + + +def _initial_messages() -> list[Any]: + """Return a minimal messages list (system + user).""" + from langchain_core.messages import HumanMessage, SystemMessage + + return [ + SystemMessage(content="You are a helpful assistant."), + HumanMessage(content="Test message"), + ] + + +# --------------------------------------------------------------------------- +# Given steps +# --------------------------------------------------------------------------- + + +@given('I have an LLMAgent with tool loop test config for tool "{tool_name}"') +def step_basic_tool_loop_config(context, tool_name): + context.tool_loop_config = { + "name": "tool_loop_agent", + "provider": "openai", + "api_key": "test", + "tools": [{"name": tool_name}], + } + context.tool_loop_tool_name = tool_name + + +@given( + 'I have an LLMAgent with tool loop test config for tool "{tool_name}" and ' + "tool_max_rounds {n:d}" +) +def step_tool_loop_config_with_max_rounds(context, tool_name, n): + context.tool_loop_config = { + "name": "tool_loop_agent", + "provider": "openai", + "api_key": "test", + "tools": [{"name": tool_name}], + "tool_max_rounds": n, + } + context.tool_loop_tool_name = tool_name + + +@given( + 'I have an LLMAgent with tool loop test config for tool "{tool_name}" and ' + "tiny token_budget_percent" +) +def step_tool_loop_config_tiny_budget(context, tool_name): + context.tool_loop_config = { + "name": "tool_loop_agent", + "provider": "openai", + "api_key": "test", + "tools": [{"name": tool_name}], + "token_budget_percent": 0.000001, + "tool_max_rounds": 5, + } + context.tool_loop_tool_name = tool_name + + +@given( + 'I have an LLMAgent with tool loop test config for tool "{tool_name}" and ' + "pruning enabled" +) +def step_tool_loop_config_pruning_enabled(context, tool_name): + context.tool_loop_config = { + "name": "tool_loop_agent", + "provider": "openai", + "api_key": "test", + "tools": [{"name": tool_name}], + "allow_tool_output_pruning": True, + "pruning_threshold": 1, + "pruning_tool_filter": [tool_name], + "unsafe_mode": True, + } + context.tool_loop_tool_name = tool_name + + +@given( + 'I have an LLMAgent with tool loop test config for tool "{tool_name}" and ' + "bad tool_max_rounds" +) +def step_tool_loop_config_bad_max_rounds(context, tool_name): + context.tool_loop_config = { + "name": "tool_loop_agent", + "provider": "openai", + "api_key": "test", + "tools": [{"name": tool_name}], + "tool_max_rounds": "not_a_number", + } + context.tool_loop_tool_name = tool_name + + +@given( + 'I set a mock that makes one tool call then returns a final answer "{final_text}"' +) +def step_mock_one_tool_call_then_final(context, final_text): + tool_name = getattr(context, "tool_loop_tool_name", "echo") + call_counter = [0] + context._tool_loop_round_prompts = [] + context._tool_loop_round_completions = [] + + async def _ainvoke(messages, **kwargs): + call_counter[0] += 1 + if call_counter[0] == 1: + msg = _make_ai_message( + content="", + tool_calls=[ + { + "id": "call_001", + "name": tool_name, + "args": {"message": "tool input"}, + } + ], + prompt=15, + completion=5, + ) + context._tool_loop_round_prompts.append(15) + context._tool_loop_round_completions.append(5) + return msg + msg = _make_ai_message(content=final_text, prompt=20, completion=10) + context._tool_loop_round_prompts.append(20) + context._tool_loop_round_completions.append(10) + return msg + + context.tool_loop_mock_ainvoke = _ainvoke + context.tool_loop_config.setdefault("tool_max_rounds", 5) + + +@given( + 'I set a mock that makes two tool calls then returns a final answer "{final_text}"' +) +def step_mock_two_tool_calls_then_final(context, final_text): + tool_name = getattr(context, "tool_loop_tool_name", "echo") + call_counter = [0] + context._tool_loop_round_prompts = [] + context._tool_loop_round_completions = [] + + async def _ainvoke(messages, **kwargs): + call_counter[0] += 1 + if call_counter[0] <= 2: + msg = _make_ai_message( + content="", + tool_calls=[ + { + "id": f"call_{call_counter[0]:03d}", + "name": tool_name, + "args": {"message": "tool input"}, + } + ], + prompt=15, + completion=5, + ) + context._tool_loop_round_prompts.append(15) + context._tool_loop_round_completions.append(5) + return msg + msg = _make_ai_message(content=final_text, prompt=20, completion=10) + context._tool_loop_round_prompts.append(20) + context._tool_loop_round_completions.append(10) + return msg + + context.tool_loop_mock_ainvoke = _ainvoke + context.tool_loop_config.setdefault("tool_max_rounds", 5) + + +@given("I set a mock that always returns tool_calls never a final answer") +def step_mock_always_tool_calls(context): + tool_name = getattr(context, "tool_loop_tool_name", "echo") + call_counter = [0] + + async def _ainvoke(messages, **kwargs): + call_counter[0] += 1 + return _make_ai_message( + content="", + tool_calls=[ + { + "id": f"call_{call_counter[0]:03d}", + "name": tool_name, + "args": {"message": "always calling"}, + } + ], + prompt=10, + completion=5, + ) + + context.tool_loop_mock_ainvoke = _ainvoke + + +@given('I set a mock that triggers budget exhaustion then returns "{final_text}"') +def step_mock_budget_exhaustion(context, final_text): + tool_name = getattr(context, "tool_loop_tool_name", "echo") + call_counter = [0] + + async def _ainvoke(messages, **kwargs): + call_counter[0] += 1 + if call_counter[0] == 1: + # First call: return tool calls so something accumulates + return _make_ai_message( + content="", + tool_calls=[ + { + "id": "call_001", + "name": tool_name, + "args": {"message": "input"}, + } + ], + prompt=30, + completion=20, + ) + # Synthesis call after budget exhausted + return _make_ai_message(content=final_text, prompt=10, completion=5) + + context.tool_loop_mock_ainvoke = _ainvoke + + +@given('I set a mock that returns empty content tool calls then "{synthesis_answer}"') +def step_mock_stuck_then_synthesis(context, synthesis_answer): + tool_name = getattr(context, "tool_loop_tool_name", "echo") + call_counter = [0] + + async def _ainvoke(messages, **kwargs): + call_counter[0] += 1 + # Check if this is the post-loop synthesis call (no tools in kwargs) + if not kwargs.get("tools"): + return _make_ai_message(content=synthesis_answer, prompt=12, completion=8) + # During the loop: return tool_calls with empty content + return _make_ai_message( + content="", + tool_calls=[ + { + "id": f"call_{call_counter[0]:03d}", + "name": tool_name, + "args": {"message": "loop call"}, + } + ], + prompt=10, + completion=5, + ) + + context.tool_loop_mock_ainvoke = _ainvoke + + +@given('I set a mock that makes one file_read tool call then returns "{final_text}"') +def step_mock_file_read_then_final(context, final_text): + call_counter = [0] + + async def _ainvoke(messages, **kwargs): + call_counter[0] += 1 + if call_counter[0] == 1: + return _make_ai_message( + content="", + tool_calls=[ + { + "id": "call_001", + "name": "file_read", + "args": {"path": "/tmp/big_file.txt"}, + } + ], + prompt=15, + completion=5, + ) + return _make_ai_message(content=final_text, prompt=18, completion=7) + + context.tool_loop_mock_ainvoke = _ainvoke + + +@given("I set a mock pruning model that adds extra tokens") +def step_mock_pruning_model(context): + from langchain_core.messages import AIMessage as LC_AIMessage + + context.tool_loop_pruning_response = LC_AIMessage( + content="[PRUNE_OUTPUT_START]\npruned content\n[PRUNE_OUTPUT_END]", + usage_metadata={ + "input_tokens": 25, + "output_tokens": 10, + "total_tokens": 35, + }, + ) + context.tool_loop_extra_pruning_prompt = 25 + context.tool_loop_extra_pruning_completion = 10 + + +@given('I set a mock that calls a nonexistent tool then returns "{final_text}"') +def step_mock_nonexistent_tool(context, final_text): + call_counter = [0] + + async def _ainvoke(messages, **kwargs): + call_counter[0] += 1 + if call_counter[0] == 1: + return _make_ai_message( + content="", + tool_calls=[ + { + "id": "call_001", + "name": "nonexistent_tool_xyz", + "args": {}, + } + ], + prompt=12, + completion=6, + ) + return _make_ai_message(content=final_text, prompt=10, completion=4) + + context.tool_loop_mock_ainvoke = _ainvoke + context.tool_loop_config.setdefault("tool_max_rounds", 5) + + +@given("I set a mock that returns one tool round then raises a RuntimeError") +def step_mock_one_round_then_error(context): + tool_name = getattr(context, "tool_loop_tool_name", "echo") + call_counter = [0] + + async def _ainvoke(messages, **kwargs): + call_counter[0] += 1 + if call_counter[0] == 1: + return _make_ai_message( + content="", + tool_calls=[ + { + "id": "call_001", + "name": tool_name, + "args": {"message": "input"}, + } + ], + prompt=20, + completion=8, + ) + raise RuntimeError("Simulated LLM failure after first round") + + context.tool_loop_mock_ainvoke = _ainvoke + context.tool_loop_config.setdefault("tool_max_rounds", 5) + + +# --------------------------------------------------------------------------- +# When steps +# --------------------------------------------------------------------------- + + +@when("I run _execute_tool_loop with an initial messages list") +def step_run_execute_tool_loop(context): + agent = _make_agent(context.tool_loop_config) + + # Patch build_chat_model to return a mock with custom ainvoke + ainvoke_fn = getattr(context, "tool_loop_mock_ainvoke", None) + if ainvoke_fn is not None: + mock_model = Mock() + mock_model.temperature = 0.7 + mock_model.ainvoke = ainvoke_fn + agent.chat_model = mock_model + + # If pruning model expected, also patch _run_pruning_pass + if hasattr(context, "tool_loop_pruning_response"): + pruning_resp = context.tool_loop_pruning_response + + async def _mock_pruning_pass( + tool_name, raw_output, task_context, prune_context=None + ): + return ( + str(pruning_resp.content), + context.tool_loop_extra_pruning_prompt, + context.tool_loop_extra_pruning_completion, + ) + + agent._run_pruning_pass = _mock_pruning_pass # type: ignore[method-assign] + + msgs = _initial_messages() + context.tool_loop_messages = msgs + + try: + context.tool_loop_result = asyncio.run(agent._execute_tool_loop(msgs)) + context.tool_loop_error = None + except _ToolLoopError as e: + context.tool_loop_result = None + context.tool_loop_error = e + except Exception as e: + context.tool_loop_result = None + context.tool_loop_error = e + + +@when("I attempt to run _execute_tool_loop with an initial messages list") +def step_attempt_run_execute_tool_loop(context): + """Same as run but expects an error.""" + step_run_execute_tool_loop(context) + + +# --------------------------------------------------------------------------- +# Then steps +# --------------------------------------------------------------------------- + + +@then('the loop result final_response content should be "{expected}"') +def step_loop_result_content(context, expected): + assert context.tool_loop_result is not None, ( + f"Expected loop result but got error: {context.tool_loop_error}" + ) + actual = str(context.tool_loop_result.final_response.content) + assert actual == expected, f"Expected {expected!r}, got {actual!r}" + + +@then("the loop result accumulated_prompt should be greater than 0") +def step_loop_result_prompt_gt_zero(context): + assert context.tool_loop_result is not None + assert context.tool_loop_result.accumulated_prompt > 0, ( + f"Expected prompt > 0, got {context.tool_loop_result.accumulated_prompt}" + ) + + +@then("the loop result accumulated_completion should be greater than 0") +def step_loop_result_completion_gt_zero(context): + assert context.tool_loop_result is not None + assert context.tool_loop_result.accumulated_completion > 0, ( + f"Expected completion > 0, got {context.tool_loop_result.accumulated_completion}" + ) + + +@then("the loop result budget_exhausted should be false") +def step_loop_result_budget_not_exhausted(context): + assert context.tool_loop_result is not None + assert not context.tool_loop_result.budget_exhausted, ( + "Expected budget_exhausted=False" + ) + + +@then("the loop result synthesis_was_run should be false") +def step_loop_result_synthesis_not_run(context): + assert context.tool_loop_result is not None + assert not context.tool_loop_result.synthesis_was_run, ( + "Expected synthesis_was_run=False" + ) + + +@then("the loop result accumulated_prompt should be at least {n:d}") +def step_loop_result_prompt_at_least(context, n): + assert context.tool_loop_result is not None + actual = context.tool_loop_result.accumulated_prompt + assert actual >= n, f"Expected accumulated_prompt >= {n}, got {actual}" + + +@then("the loop result accumulated_completion should be at least {n:d}") +def step_loop_result_completion_at_least(context, n): + assert context.tool_loop_result is not None + actual = context.tool_loop_result.accumulated_completion + assert actual >= n, f"Expected accumulated_completion >= {n}, got {actual}" + + +@then( + "the loop result accumulated_prompt should equal the sum of all rounds prompt tokens" +) +def step_loop_result_prompt_sum(context): + assert context.tool_loop_result is not None + expected = sum(context._tool_loop_round_prompts) + actual = context.tool_loop_result.accumulated_prompt + assert actual == expected, ( + f"Expected accumulated_prompt={expected} (sum of rounds {context._tool_loop_round_prompts}), " + f"got {actual}" + ) + + +@then( + "the loop result accumulated_completion should equal the sum of all rounds completion tokens" +) +def step_loop_result_completion_sum(context): + assert context.tool_loop_result is not None + expected = sum(context._tool_loop_round_completions) + actual = context.tool_loop_result.accumulated_completion + assert actual == expected, ( + f"Expected accumulated_completion={expected} (sum of rounds {context._tool_loop_round_completions}), " + f"got {actual}" + ) + + +@then("the loop result budget_exhausted should be true") +def step_loop_result_budget_exhausted(context): + assert context.tool_loop_result is not None + assert context.tool_loop_result.budget_exhausted, "Expected budget_exhausted=True" + + +@then("the loop result synthesis_was_run should be true") +def step_loop_result_synthesis_run(context): + assert context.tool_loop_result is not None + assert context.tool_loop_result.synthesis_was_run, "Expected synthesis_was_run=True" + + +@then("the loop result should contain some content from the synthesis prompt") +def step_loop_result_synthesis_content(context): + assert context.tool_loop_result is not None + # After max_rounds exhaustion, synthesis runs; result content may be empty or + # contain synthesis output. We just check that the loop completed normally. + assert context.tool_loop_result.final_response is not None, ( + "Expected final_response to be set" + ) + + +@then("the loop result accumulated_prompt should include pruning tokens") +def step_loop_result_includes_pruning_tokens(context): + assert context.tool_loop_result is not None + expected_min = context.tool_loop_extra_pruning_prompt + actual = context.tool_loop_result.accumulated_prompt + assert actual >= expected_min, ( + f"Expected accumulated_prompt to include pruning tokens ({expected_min}), " + f"got {actual}" + ) + + +@then("the messages list should contain a ToolMessage with not available content") +def step_messages_contain_not_available(context): + msgs = context.tool_loop_messages + found = any( + isinstance(m, ToolMessage) and "is not available" in str(m.content) + for m in msgs + ) + assert found, ( + f"Expected a ToolMessage with 'is not available' in {[type(m).__name__ for m in msgs]}" + ) + + +@then("a _ToolLoopError should be raised") +def step_tool_loop_error_raised(context): + assert isinstance(context.tool_loop_error, _ToolLoopError), ( + f"Expected _ToolLoopError, got {type(context.tool_loop_error).__name__}: " + f"{context.tool_loop_error}" + ) + + +@then("the _ToolLoopError any_invocation_made should be true") +def step_tool_loop_error_any_invocation_true(context): + assert isinstance(context.tool_loop_error, _ToolLoopError) + assert context.tool_loop_error.any_invocation_made, ( + "Expected any_invocation_made=True" + ) + + +@then("the _ToolLoopError any_invocation_made should be false") +def step_tool_loop_error_any_invocation_false(context): + assert isinstance(context.tool_loop_error, _ToolLoopError) + assert not context.tool_loop_error.any_invocation_made, ( + "Expected any_invocation_made=False" + ) + + +@then("the _ToolLoopError accumulated_prompt should be greater than 0") +def step_tool_loop_error_prompt_gt_zero(context): + assert isinstance(context.tool_loop_error, _ToolLoopError) + assert context.tool_loop_error.accumulated_prompt > 0, ( + f"Expected accumulated_prompt > 0, got {context.tool_loop_error.accumulated_prompt}" + ) + + +@then("the _ToolLoopError cause should be a ConfigurationError") +def step_tool_loop_error_cause_config_error(context): + assert isinstance(context.tool_loop_error, _ToolLoopError) + assert isinstance(context.tool_loop_error.cause, ConfigurationError), ( + f"Expected cause=ConfigurationError, got {type(context.tool_loop_error.cause).__name__}" + ) diff --git a/robot/ToolCallingTestLib.py b/robot/ToolCallingTestLib.py index 66e623d..5eea226 100644 --- a/robot/ToolCallingTestLib.py +++ b/robot/ToolCallingTestLib.py @@ -25,6 +25,7 @@ class ToolCallingTestLib: self._last_result: ActorResult | None = None self._mock_ainvoke_calls: list[dict[str, Any]] = [] self._patches: list[Any] = [] + self._multi_round_expected_prompt: int = 0 def _teardown_patches(self) -> None: for p in self._patches: @@ -198,3 +199,119 @@ class ToolCallingTestLib: assert len(self._last_result.response.strip()) > 0, ( "Expected non-empty response after tool calls" ) + + # ------------------------------------------------------------------ + # New keywords for issue #67 (stream_message tool-loop integration) + # ------------------------------------------------------------------ + + def stream_result_response_is_non_empty(self) -> None: # noqa: D401 + """Verify that the stream response is non empty.""" + assert self._last_result is not None + assert len(self._last_result.response.strip()) > 0, ( + f"Expected non-empty stream response, got: {self._last_result.response!r}" + ) + + def stream_result_prompt_tokens_greater_than_zero(self) -> None: + """Verify stream result has non-zero accumulated prompt tokens.""" + assert self._last_result is not None + assert self._last_result.prompt_tokens > 0, ( + f"Expected stream prompt_tokens > 0, got {self._last_result.prompt_tokens}" + ) + + def stream_result_response_contains(self, text: str) -> None: + """Verify the stream response contains *text*.""" + assert self._last_result is not None + assert text in self._last_result.response, ( + f"Expected {text!r} in stream response, got: {self._last_result.response[:200]!r}" + ) + + def stream_result_prompt_tokens_equals_multi_round_sum(self) -> None: + """Verify accumulated prompt tokens equal the sum across all mock rounds.""" + assert self._last_result is not None + expected = self._multi_round_expected_prompt + actual = self._last_result.prompt_tokens + assert actual == expected, ( + f"Expected prompt_tokens={expected} (sum of all rounds), got {actual}" + ) + + def create_executor_with_multi_round_tool_calling_agent( + self, tool_name: str = "echo", rounds: int = 2 + ) -> None: + """Create an Executor whose mock model returns tool_calls for *rounds* rounds + then a final answer. Tracks per-round token counts for assertion. + """ + self._teardown_patches() + self._mock_ainvoke_calls = [] + + # Per-round prompt tokens: 50 per tool round + 30 for final answer + tool_round_prompt = 50 + final_prompt = 30 + tool_round_completion = 10 + final_completion = 15 + total_rounds = int(rounds) + + self._multi_round_expected_prompt = ( + tool_round_prompt * total_rounds + final_prompt + ) + + final_response_text = "Multi round final answer" + + def _build_mock_model(*args, **kwargs): # type: ignore[no-untyped-def] + mock_model = MagicMock() + mock_model.temperature = 0.7 + + async def _mock_ainvoke(messages, **invoke_kwargs): # type: ignore[no-untyped-def] + self._mock_ainvoke_calls.append({"kwargs": invoke_kwargs}) + call_count = len(self._mock_ainvoke_calls) + if call_count <= total_rounds and invoke_kwargs.get("tools"): + return AIMessage( + content="", + usage_metadata={ + "input_tokens": tool_round_prompt, + "output_tokens": tool_round_completion, + "total_tokens": tool_round_prompt + tool_round_completion, + }, + tool_calls=[ + { + "id": f"call_{call_count:03d}", + "name": tool_name, + "args": {"message": f"round_{call_count}"}, + } + ], + ) + return AIMessage( + content=final_response_text, + usage_metadata={ + "input_tokens": final_prompt, + "output_tokens": final_completion, + "total_tokens": final_prompt + final_completion, + }, + ) + + mock_model.ainvoke = _mock_ainvoke + mock_model.astream = MagicMock() # not called in tool path + return mock_model + + patcher = patch( + "cleveractors.agents.llm.build_chat_model", + side_effect=_build_mock_model, + ) + patcher.start() + self._patches.append(patcher) + + config = { + "type": "llm", + "name": "multi_round_tool_agent", + "provider": "openai", + "model": "gpt-3.5-turbo", + "config": { + "tools": [{"name": tool_name}], + }, + } + + self._executor = create_executor( + config_dict=config, + credentials={"openai": {"api_key": "mock-key"}}, + limits={}, + pricing={}, + ) diff --git a/robot/llm_tool_calling.robot b/robot/llm_tool_calling.robot index b52c50d..feaa630 100644 --- a/robot/llm_tool_calling.robot +++ b/robot/llm_tool_calling.robot @@ -41,12 +41,24 @@ Tool Calling Produces Non-Empty Final Response Tool Call Result Is In Response Result Prompt Tokens Greater Than Zero -Streaming Path Works With Tool-Configured Agent - [Documentation] execute_stream produces a valid result with a tool-configured - ... agent. Tool calling is not exercised in the streaming path — tools are - ... not passed to the model during stream_message. This test verifies the - ... streaming path itself works and does not crash with a tool-configured agent. +Streaming Path Exercises Tool Loop Via Execute Stream + [Documentation] execute_stream on a tool-declaring single-LLM actor exercises + ... the _execute_tool_loop() path (issue #67): the tool-call loop runs via + ... ainvoke(), the final response is yielded as a single content chunk, and + ... executor.last_result reports accumulated token counts covering all rounds. Create Executor With Tool Calling Agent tool_name=echo message=from_stream Execute Stream With Message Tell me about yourself Result Is Valid Actor Result Result Has Nodes + Stream Result Response Is Non Empty + Stream Result Prompt Tokens Greater Than Zero + +Streaming Path Accumulates Tokens Across Two Tool Rounds + [Documentation] When the mock LLM returns tool_calls for two rounds before the + ... final answer, execute_stream yields the final answer and the accumulated + ... prompt_tokens equals the sum across all three invocations. + Create Executor With Multi Round Tool Calling Agent tool_name=echo rounds=2 + Execute Stream With Message Multi round test + Result Is Valid Actor Result + Stream Result Response Contains Multi round final answer + Stream Result Prompt Tokens Equals Multi Round Sum diff --git a/src/cleveractors/agents/llm.py b/src/cleveractors/agents/llm.py index 5b0f3d0..259f65d 100644 --- a/src/cleveractors/agents/llm.py +++ b/src/cleveractors/agents/llm.py @@ -24,6 +24,7 @@ Extended Provider Routing (ADR-2028): from __future__ import annotations import contextvars +import dataclasses import json import logging import os @@ -111,6 +112,82 @@ ChatGoogleGenerativeAI: Any = None ChatOpenAI: Any = None +# --------------------------------------------------------------------------- +# _ToolLoopResult — return type for LLMAgent._execute_tool_loop() +# --------------------------------------------------------------------------- + + +@dataclasses.dataclass +class _ToolLoopResult: + """Return value from ``LLMAgent._execute_tool_loop()``. + + Carries everything both callers (``process_message`` and + ``stream_message``) need after the tool-call loop completes: + + Attributes: + final_response: The last AIMessage produced by the loop (the model's + final answer, or the synthesis response when stuck/budget-exhausted). + accumulated_prompt: Sum of prompt tokens across all ``ainvoke()`` + calls in the loop (main rounds + synthesis + pruning passes). + accumulated_completion: Corresponding completion-token sum. + budget_exhausted: True when the token-budget synthesis flow was + triggered mid-loop (§4.4.7 D-4). + synthesis_was_run: True when the stuck-model synthesis prompt was + appended and a final ``ainvoke()`` issued (post-loop path). + """ + + final_response: Any # AIMessage from langchain_core + accumulated_prompt: int + accumulated_completion: int + budget_exhausted: bool + synthesis_was_run: bool + + +# --------------------------------------------------------------------------- +# _ToolLoopError — internal exception wrapper for _execute_tool_loop() +# --------------------------------------------------------------------------- + + +class _ToolLoopError(Exception): + """Wraps an exception raised inside ``_execute_tool_loop()``. + + Carries the partial accumulated token counts collected up to the point + of failure so callers can preserve billing integrity even when the loop + terminates abnormally. + + The caller pattern:: + + try: + result = await self._execute_tool_loop(messages) + except _ToolLoopError as tle: + if tle.any_invocation_made: + _captured_prompt = tle.accumulated_prompt + _captured_completion = tle.accumulated_completion + raise tle.cause + + Attributes: + cause: The original exception that caused the loop to abort. + accumulated_prompt: Prompt tokens accumulated before failure. + accumulated_completion: Completion tokens accumulated before failure. + any_invocation_made: True when at least one ``ainvoke()`` completed + successfully before the exception was raised. Used by callers + to decide whether to preserve or reset billing counters. + """ + + def __init__( + self, + cause: BaseException, + accumulated_prompt: int, + accumulated_completion: int, + any_invocation_made: bool, + ) -> None: + super().__init__(str(cause)) + self.cause = cause + self.accumulated_prompt = accumulated_prompt + self.accumulated_completion = accumulated_completion + self.any_invocation_made = any_invocation_made + + class LLMAgent(AgentWithMemory): """ Reactive agent powered by a large language model (LLM) using LangChain. @@ -605,6 +682,585 @@ class LLMAgent(AgentWithMemory): ) return (parsed, _pp, _pc) + # ------------------------------------------------------------------ + # _execute_tool_loop — shared tool-call orchestration helper + # ------------------------------------------------------------------ + + async def _execute_tool_loop( + self, + messages: list[Any], + ) -> _ToolLoopResult: + """Execute the complete multi-turn tool-call loop via ``ainvoke()``. + + This helper encapsulates the entire tool-call orchestration so that + both :meth:`process_message` and :meth:`stream_message` can share + the same implementation without duplication. + + The method: + + 1. Parses ``tool_max_rounds`` from config/environment. + 2. Builds ``invoke_kwargs`` with the tool schemas (with optional + pruning augmentation per §4.4.8). + 3. Runs the multi-turn ``ainvoke()`` loop until the model stops + producing tool calls, ``tool_max_rounds`` is exhausted, or the + token budget is exceeded (§4.4.7). + 4. Dispatches each ``tool_call`` to an ephemeral :class:`ToolAgent`. + 5. Runs tool-output pruning when enabled (§4.4.8). + 6. On token-budget exhaustion, appends a synthesis prompt and runs + one more round (D-4 of ADR-2031). + 7. On stuck-model (empty final content), appends the synthesis prompt + post-loop and runs a final ``ainvoke()``. + + *messages* is modified in-place: ``AIMessage`` and ``ToolMessage`` + entries are appended as the conversation progresses. + + Args: + messages: The initial message list (system prompt + conversation + history + current user message). Modified in-place. + + Returns: + A :class:`_ToolLoopResult` carrying the final AIMessage, the + accumulated prompt and completion token counts spanning all + ``ainvoke()`` calls (main rounds + pruning passes), the + ``budget_exhausted`` flag, and the ``synthesis_was_run`` flag. + + Raises: + _ToolLoopError: Wraps any exception raised during the loop, + carrying partial accumulated token counts for billing + integrity. Callers must catch this, extract the counts, + then re-raise ``cause``. + """ + _accumulated_prompt: int = 0 + _accumulated_completion: int = 0 + _any_invocation_made: bool = False + _budget_exhausted: bool = False + _synthesis_was_run: bool = False + response: Any = None + + try: + # ── Parse tool_max_rounds ────────────────────────────────────── + _raw_max_rounds: object = self.config.get( + "tool_max_rounds" + ) or os.environ.get("TOOL_MAX_ROUNDS", "20") + try: + _TOOL_MAX_ROUNDS = max(1, int(str(_raw_max_rounds))) + except (TypeError, ValueError) as _mre: + raise ConfigurationError( + f"tool_max_rounds must be an integer, got {_raw_max_rounds!r}" + ) from _mre + + # ── Build invoke_kwargs with tools ───────────────────────────── + invoke_kwargs: dict[str, Any] = {} + if self._allow_tool_output_pruning: + _all_functions = all( + isinstance(t, dict) and "function" in t + for t in (self._lc_tools or []) + ) + if _all_functions: + invoke_kwargs["tools"] = self._augment_tool_schemas_for_pruning( + self._lc_tools or [] + ) + else: + invoke_kwargs["tools"] = self._lc_tools + else: + invoke_kwargs["tools"] = self._lc_tools + + # Set of declared tool names for validation + _declared_names = { + t.get("function", {}).get("name", "") for t in (self._lc_tools or []) + } + + # ── Multi-turn tool-call loop ────────────────────────────────── + for _tool_round in range(_TOOL_MAX_ROUNDS): + # ── Token-budget check (§4.4.7) ──────────────────────────── + if self._token_budget_percent is not None: + _budget_ceiling = max( + 1, + int( + self._token_budget_percent + * self._get_model_context_window() + ), + ) + _est = _accumulated_prompt + _accumulated_completion + _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 + ) + _any_invocation_made = True + _bp, _bc, _ = self._extract_token_counts(response) + _accumulated_prompt += _bp + _accumulated_completion += _bc + _synth_tool_calls: list[dict[str, Any]] = ( + getattr(response, "tool_calls", None) or [] + ) + if isinstance(_synth_tool_calls, list) and _synth_tool_calls: + messages.append(response) + for tc in _synth_tool_calls: + call_id = tc.get("id", "") + fn_def = tc.get("function") + if isinstance(fn_def, dict): + tool_name = fn_def.get("name", "") + arguments_raw = fn_def.get("arguments") + else: + tool_name = tc.get("name", "") + arguments_raw = tc.get("args") + if not tool_name: + messages.append( + ToolMessage( + content="Tool name is empty", + tool_call_id=call_id, + ) + ) + continue + args: dict[str, Any] = {} + if isinstance(arguments_raw, str): + try: + args = json.loads(arguments_raw) + except (json.JSONDecodeError, ValueError): + args = {"_raw": arguments_raw} + elif isinstance(arguments_raw, dict): + args = arguments_raw + _bq_output_prune = args.pop("output_prune", None) + _bq_prune_ctx = args.pop("output_prune_context", None) + try: + from cleveractors.agents.tool import ( + ToolAgent as _TA, + ) + + parent_unsafe = self.config.get( + "unsafe_mode", False + ) + tool_cfg: dict[str, Any] = { + "tools": [{"name": tool_name}], + "safe_mode": not parent_unsafe, + "allow_shell": self.config.get( + "allow_shell", False + ), + "exec_python": self.config.get( + "exec_python", False + ), + "timeout": self.config.get("timeout", 1), + } + t_agent = _TA( + name=f"_tc_synth_budget_{call_id}", + config=tool_cfg, + template_renderer=self.template_renderer, + ) + t_ctx: dict[str, Any] | None = ( + {"_unsafe_mode": True} + if parent_unsafe + else None + ) + if tool_name not in _declared_names: + messages.append( + ToolMessage( + content=f"Tool '{tool_name}' is not available.", + tool_call_id=call_id, + ) + ) + break + 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 + ): + try: + ( + tool_output, + _bp_tok, + _bc_tok, + ) = await self._run_pruning_pass( + tool_name, + raw_out, + list(messages), + prune_context=_bq_prune_ctx, + ) + except MissingUsageMetadataError: + logger.warning( + "Agent %s: pruning pass for %r " + "raised MissingUsageMetadataError; " + "%d prompt + %d completion tokens " + "accumulated before failure will be " + "preserved for billing", + self.name, + tool_name, + _accumulated_prompt, + _accumulated_completion, + ) + raise + _accumulated_prompt += _bp_tok + _accumulated_completion += _bc_tok + 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) + _any_invocation_made = True + _bp3, _bc3, _ = self._extract_token_counts(response) + _accumulated_prompt += _bp3 + _accumulated_completion += _bc3 + _budget_exhausted = True + break + + # ── Regular ainvoke ──────────────────────────────────────── + response = await self.chat_model.ainvoke(messages, **invoke_kwargs) + _any_invocation_made = True + _mp, _mc, _ = self._extract_token_counts(response) + _accumulated_prompt += _mp + _accumulated_completion += _mc + + response_tool_calls: list[dict[str, Any]] = ( + getattr(response, "tool_calls", None) or [] + ) + if not isinstance(response_tool_calls, list) or not response_tool_calls: + break + + messages.append(response) + + for tc in response_tool_calls: + call_id = tc.get("id", "") + fn_def = tc.get("function") + if isinstance(fn_def, dict): + tool_name = fn_def.get("name", "") + arguments_raw = fn_def.get("arguments") + else: + tool_name = tc.get("name", "") + arguments_raw = tc.get("args") + if not tool_name: + messages.append( + ToolMessage( + content="Tool name is empty", + tool_call_id=call_id, + ) + ) + continue + + args_m: dict[str, Any] = {} + if isinstance(arguments_raw, str): + try: + args_m = json.loads(arguments_raw) + except (json.JSONDecodeError, ValueError): + args_m = {"_raw": arguments_raw} + elif isinstance(arguments_raw, dict): + args_m = arguments_raw + _output_prune = args_m.pop("output_prune", None) + _prune_context = args_m.pop("output_prune_context", None) + try: + from cleveractors.agents.tool import ToolAgent as _TA + + parent_unsafe = self.config.get("unsafe_mode", False) + tool_config: dict[str, Any] = { + "tools": [{"name": tool_name}], + "safe_mode": not parent_unsafe, + "allow_shell": self.config.get("allow_shell", False), + "exec_python": self.config.get("exec_python", False), + "timeout": self.config.get("timeout", 1), + } + + if tool_name not in _declared_names: + messages.append( + ToolMessage( + content=f"Tool '{tool_name}' is not available.", + tool_call_id=call_id, + ) + ) + continue + + agent = _TA( + name=f"_tc_{call_id}", + config=tool_config, + template_renderer=self.template_renderer, + ) + tool_ctx: dict[str, Any] | None = ( + {"_unsafe_mode": True} if parent_unsafe else None + ) + raw_tool_output = await agent.process_message( + json.dumps({"tool": tool_name, "args": args_m}) + if isinstance(args_m, dict) + else arguments_raw or "", + context=tool_ctx, + ) + # ── Tool-output pruning pass (§4.4.8) ──────────── + 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) + try: + ( + tool_output, + _mlp_tok, + _mlc_tok, + ) = await self._run_pruning_pass( + tool_name, + raw_tool_output, + _task_context, + prune_context=_prune_context, + ) + except MissingUsageMetadataError: + logger.warning( + "Agent %s: pruning pass for %r " + "raised MissingUsageMetadataError; " + "%d prompt + %d completion tokens " + "accumulated before failure will be " + "preserved for billing", + self.name, + tool_name, + _accumulated_prompt, + _accumulated_completion, + ) + raise + _accumulated_prompt += _mlp_tok + _accumulated_completion += _mlc_tok + 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_output, tool_call_id=call_id) + ) + except (ExecutionError, ConfigurationError) as _tool_err: + _tool_err_msg = str(_tool_err) + logger.warning( + "Agent %s: tool call failed for %r (tool=%s): %s", + self.name, + call_id, + tool_name, + _tool_err_msg, + ) + messages.append( + ToolMessage( + content=f"Tool '{tool_name}' error: {_tool_err_msg}", + tool_call_id=call_id, + ) + ) + + # ── Post-loop: stuck-model synthesis ────────────────────────── + # When the loop exhausted all rounds but the model never produced + # meaningful text (stuck in tool-only mode), inject a synthesis + # prompt and issue one final ainvoke(). Skip when the budget- + # exhausted synthesis flow already handled termination. + _response_text_check = str(response.content) if response is not None else "" + if ( + not _budget_exhausted + and not _response_text_check.strip() + and len(messages) > 2 + ): + _synthesis_was_run = True + messages.append( + HumanMessage( + content=( + "You have finished gathering information. " + "Now produce your final answer based on all the " + "data collected. Do NOT make any more tool calls. " + "If the answer requires writing a file, call " + "file_write in this very response." + ) + ) + ) + response = await self.chat_model.ainvoke(messages, tools=self._lc_tools) + _any_invocation_made = True + _sp, _sc, _ = self._extract_token_counts(response) + _accumulated_prompt += _sp + _accumulated_completion += _sc + + # Allow *one* final tool-call round so the model can + # write files or perform other last-minute operations. + synth_tool_calls: object = getattr(response, "tool_calls", None) + + if isinstance(synth_tool_calls, list) and synth_tool_calls: + messages.append(response) + for tc in synth_tool_calls: + call_id = tc.get("id", "") + fn_def = tc.get("function") + if isinstance(fn_def, dict): + tool_name = fn_def.get("name", "") + arguments_raw = fn_def.get("arguments") + else: + tool_name = tc.get("name", "") + arguments_raw = tc.get("args") + if not tool_name: + messages.append( + ToolMessage( + content="Tool name is empty", + tool_call_id=call_id, + ) + ) + continue + + if tool_name not in _declared_names: + messages.append( + ToolMessage( + content=f"Tool '{tool_name}' is not available.", + tool_call_id=call_id, + ) + ) + continue + + args_s: dict[str, Any] = {} + if isinstance(arguments_raw, str): + try: + args_s = json.loads(arguments_raw) + except (json.JSONDecodeError, ValueError): + args_s = {"_raw": arguments_raw} + elif isinstance(arguments_raw, dict): + args_s = arguments_raw + try: + from cleveractors.agents.tool import ( + ToolAgent as _TA, + ) + + parent_unsafe = self.config.get("unsafe_mode", False) + tool_config_s: dict[str, Any] = { + "tools": [{"name": tool_name}], + "safe_mode": not parent_unsafe, + "allow_shell": self.config.get("allow_shell", False), + "exec_python": self.config.get("exec_python", False), + "timeout": self.config.get("timeout", 1), + } + agent_s = _TA( + name=f"_tc_synth_{call_id}", + config=tool_config_s, + template_renderer=self.template_renderer, + ) + tool_ctx_s: dict[str, Any] | None = ( + {"_unsafe_mode": True} if parent_unsafe else None + ) + tool_result = await agent_s.process_message( + json.dumps({"tool": tool_name, "args": args_s}) + if isinstance(args_s, dict) + else arguments_raw or "", + context=tool_ctx_s, + ) + messages.append( + ToolMessage(content=tool_result, tool_call_id=call_id) + ) + except (ExecutionError, ConfigurationError) as _se: + _se_msg = str(_se) + logger.warning( + "Agent %s: 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) + _any_invocation_made = True + _sfp, _sfc, _ = self._extract_token_counts(response) + _accumulated_prompt += _sfp + _accumulated_completion += _sfc + + # Log accumulated token usage across all rounds. + if _accumulated_prompt > 0 or _accumulated_completion > 0: + logger.info( + "Agent %s: accumulated token usage: " + "%d prompt + %d completion tokens across all rounds", + self.name, + _accumulated_prompt, + _accumulated_completion, + ) + + # response must be set by at least one ainvoke() call above. + if response is None: + raise RuntimeError( + "_execute_tool_loop completed without making any LLM call. " + "This is a bug — please report it." + ) + + return _ToolLoopResult( + final_response=response, + accumulated_prompt=_accumulated_prompt, + accumulated_completion=_accumulated_completion, + budget_exhausted=_budget_exhausted, + synthesis_was_run=_synthesis_was_run, + ) + + except Exception as _e: # pylint: disable=broad-exception-caught + # Do not double-wrap a _ToolLoopError that propagated through + # a recursive or nested call (guard against future refactors). + if isinstance(_e, _ToolLoopError): + raise + raise _ToolLoopError( + _e, + _accumulated_prompt, + _accumulated_completion, + _any_invocation_made, + ) from _e + # ------------------------------------------------------------------ # Public credentials property # ------------------------------------------------------------------ @@ -781,546 +1437,35 @@ class LLMAgent(AgentWithMemory): # Add current user message messages.append(HumanMessage(content=processed_message)) - # -- Tool calling support (issue #59) ----------------------------------- - # When the agent config declares tools, convert them to a LangChain- - # compatible format and pass them to the LLM so it can produce - # structured tool calls instead of hallucinating tool names in plain - # text. - # - # Multi-turn tool-use loop: when the model returns tool_calls, - # execute each tool, append the AIMessage (with tool_calls) plus - # corresponding ToolMessages to the conversation, then re-invoke the - # model. The loop repeats until the model returns a final answer - # (no more tool_calls) or the maximum number of tool rounds is - # reached. - invoke_kwargs: dict[str, Any] = {} - _raw_max_rounds: object = self.config.get( - "tool_max_rounds" - ) or os.environ.get("TOOL_MAX_ROUNDS", "20") - try: - _TOOL_MAX_ROUNDS = max(1, int(str(_raw_max_rounds))) - except (TypeError, ValueError) as _mre: - raise ConfigurationError( - f"tool_max_rounds must be an integer, got {_raw_max_rounds!r}" - ) from _mre - _has_tools = self._lc_tools is not None and LANGCHAIN_AVAILABLE - if _has_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 + # ── Tool-call loop or plain ainvoke ────────────────────────── + # When tools are configured, delegate to _execute_tool_loop() + # (multi-turn ainvoke, ToolAgent dispatch, pruning, synthesis). + # When no tools are configured, issue a single plain ainvoke(). + _prompt_tokens: int + _completion_tokens: int + response_text: str - _budget_exhausted: bool = False - _accumulated_prompt: int = 0 - _accumulated_completion: int = 0 - - _declared_names = { - t.get("function", {}).get("name", "") for t in (self._lc_tools or []) - } - - for _tool_round in range(_TOOL_MAX_ROUNDS): - if _tool_round > 0 and not _has_tools: - break - - # ── Token-budget check (§4.4.7) ───────────────────────── - # Uses accumulated actual token counts from prior rounds - # instead of heuristic estimation (issue #65). - 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 = _accumulated_prompt + _accumulated_completion - _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 - ) - _bp, _bc, _ = self._extract_token_counts(response) - _accumulated_prompt += _bp - _accumulated_completion += _bc - _synth_tool_calls: list[dict[str, Any]] = ( - getattr(response, "tool_calls", None) or [] - ) - if ( - isinstance(_synth_tool_calls, list) - and _synth_tool_calls - and _has_tools - ): - messages.append(response) - for tc in _synth_tool_calls: - call_id = tc.get("id", "") - fn_def = tc.get("function") - if isinstance(fn_def, dict): - tool_name = fn_def.get("name", "") - arguments_raw = fn_def.get("arguments") - else: - tool_name = tc.get("name", "") - arguments_raw = tc.get("args") - if not tool_name: - messages.append( - ToolMessage( - content="Tool name is empty", - tool_call_id=call_id, - ) - ) - continue - args = {} - if isinstance(arguments_raw, str): - try: - args = json.loads(arguments_raw) - except (json.JSONDecodeError, ValueError): - args = {"_raw": arguments_raw} - elif isinstance(arguments_raw, dict): - args = arguments_raw - _bq_output_prune = args.pop("output_prune", None) - _bq_prune_ctx = args.pop("output_prune_context", None) - try: - from cleveractors.agents.tool import ( - ToolAgent as _TA, - ) - - parent_unsafe = self.config.get( - "unsafe_mode", False - ) - tool_cfg: dict[str, Any] = { - "tools": [{"name": tool_name}], - "safe_mode": not parent_unsafe, - "allow_shell": self.config.get( - "allow_shell", False - ), - "exec_python": self.config.get( - "exec_python", False - ), - "timeout": self.config.get("timeout", 1), - } - t_agent = _TA( - name=f"_tc_synth_budget_{call_id}", - config=tool_cfg, - template_renderer=self.template_renderer, - ) - t_ctx = ( - {"_unsafe_mode": True} - if parent_unsafe - else None - ) - if tool_name not in _declared_names: - messages.append( - ToolMessage( - content=f"Tool '{tool_name}' is not available.", - tool_call_id=call_id, - ) - ) - break - 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 - ): - try: - ( - tool_output, - _bp_tok, - _bc_tok, - ) = await self._run_pruning_pass( - tool_name, - raw_out, - list(messages), - prune_context=_bq_prune_ctx, - ) - except MissingUsageMetadataError: - logger.warning( - "Agent %s: pruning pass for %r " - "raised MissingUsageMetadataError; " - "%d prompt + %d completion tokens " - "accumulated before failure will be " - "preserved for billing", - self.name, - tool_name, - _accumulated_prompt, - _accumulated_completion, - ) - _captured_prompt = _accumulated_prompt - _captured_completion = ( - _accumulated_completion - ) - raise - _accumulated_prompt += _bp_tok - _accumulated_completion += _bc_tok - 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) - _bp3, _bc3, _ = self._extract_token_counts(response) - _accumulated_prompt += _bp3 - _accumulated_completion += _bc3 - response_text = str(response.content) - _budget_exhausted = True - break - - response = await self.chat_model.ainvoke(messages, **invoke_kwargs) - _mp, _mc, _ = self._extract_token_counts(response) - _accumulated_prompt += _mp - _accumulated_completion += _mc - - response_tool_calls: list[dict[str, Any]] = ( - getattr(response, "tool_calls", None) or [] - ) - if not isinstance(response_tool_calls, list) or not response_tool_calls: - break - - if not _has_tools: - break - - messages.append(response) - - for tc in response_tool_calls: - call_id = tc.get("id", "") - fn_def = tc.get("function") - if isinstance(fn_def, dict): - tool_name = fn_def.get("name", "") - arguments_raw = fn_def.get("arguments") - else: - tool_name = tc.get("name", "") - arguments_raw = tc.get("args") - if not tool_name: - messages.append( - ToolMessage( - content="Tool name is empty", - tool_call_id=call_id, - ) - ) - continue - - args: dict[str, Any] = {} - 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 - _output_prune = args.pop("output_prune", None) - _prune_context = args.pop("output_prune_context", None) - try: - from cleveractors.agents.tool import ToolAgent as _TA - - parent_unsafe = self.config.get("unsafe_mode", False) - tool_config: dict[str, Any] = { - "tools": [{"name": tool_name}], - "safe_mode": not parent_unsafe, - "allow_shell": self.config.get("allow_shell", False), - "exec_python": self.config.get("exec_python", False), - "timeout": self.config.get("timeout", 1), - } - - if tool_name not in _declared_names: - messages.append( - ToolMessage( - content=f"Tool '{tool_name}' is not available.", - tool_call_id=call_id, - ) - ) - continue - - agent = _TA( - name=f"_tc_{call_id}", - config=tool_config, - template_renderer=self.template_renderer, - ) - tool_ctx: dict[str, Any] | None = ( - {"_unsafe_mode": True} if parent_unsafe else None - ) - 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) - try: - ( - tool_output, - _mlp_tok, - _mlc_tok, - ) = await self._run_pruning_pass( - tool_name, - raw_tool_output, - _task_context, - prune_context=_prune_context, - ) - except MissingUsageMetadataError: - logger.warning( - "Agent %s: pruning pass for %r " - "raised MissingUsageMetadataError; " - "%d prompt + %d completion tokens " - "accumulated before failure will be " - "preserved for billing", - self.name, - tool_name, - _accumulated_prompt, - _accumulated_completion, - ) - _captured_prompt = _accumulated_prompt - _captured_completion = _accumulated_completion - raise - _accumulated_prompt += _mlp_tok - _accumulated_completion += _mlc_tok - 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_output, tool_call_id=call_id) - ) - except (ExecutionError, ConfigurationError) as _tool_err: - _tool_err_msg = str(_tool_err) - logger.warning( - "Agent %s: tool call failed for %r (tool=%s): %s", - self.name, - call_id, - tool_name, - _tool_err_msg, - ) - messages.append( - ToolMessage( - content=f"Tool '{tool_name}' error: {_tool_err_msg}", - tool_call_id=call_id, - ) - ) - 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. - # 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=( - "You have finished gathering information. " - "Now produce your final answer based on all the " - "data collected. Do NOT make any more tool calls. " - "If the answer requires writing a file, call " - "file_write in this very response." - ) - ) - ) - response = await self.chat_model.ainvoke(messages, tools=self._lc_tools) - _sp, _sc, _ = self._extract_token_counts(response) - _accumulated_prompt += _sp - _accumulated_completion += _sc - - # Allow *one* final tool-call round so the model can - # write files or perform other last-minute operations - # requested by the synthesis prompt. - synth_tool_calls: object = getattr(response, "tool_calls", None) - - if ( - isinstance(synth_tool_calls, list) - and synth_tool_calls - and _has_tools - ): - messages.append(response) - for tc in synth_tool_calls: - call_id = tc.get("id", "") - fn_def = tc.get("function") - if isinstance(fn_def, dict): - tool_name = fn_def.get("name", "") - arguments_raw = fn_def.get("arguments") - else: - tool_name = tc.get("name", "") - arguments_raw = tc.get("args") - if not tool_name: - messages.append( - ToolMessage( - content="Tool name is empty", - tool_call_id=call_id, - ) - ) - continue - - if tool_name not in _declared_names: - messages.append( - ToolMessage( - content=f"Tool '{tool_name}' is not available.", - tool_call_id=call_id, - ) - ) - continue - - args: dict[str, Any] = {} - 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 - try: - from cleveractors.agents.tool import ( - ToolAgent as _TA, - ) - - parent_unsafe = self.config.get("unsafe_mode", False) - tool_config: dict[str, Any] = { - "tools": [{"name": tool_name}], - "safe_mode": not parent_unsafe, - "allow_shell": self.config.get("allow_shell", False), - "exec_python": self.config.get("exec_python", False), - "timeout": self.config.get("timeout", 1), - } - agent = _TA( - name=f"_tc_synth_{call_id}", - config=tool_config, - template_renderer=self.template_renderer, - ) - tool_ctx: dict[str, Any] | None = ( - {"_unsafe_mode": True} if parent_unsafe else None - ) - tool_result = await agent.process_message( - json.dumps({"tool": tool_name, "args": args}) - if isinstance(args, dict) - else arguments_raw or "", - context=tool_ctx, - ) - messages.append( - ToolMessage(content=tool_result, tool_call_id=call_id) - ) - except (ExecutionError, ConfigurationError) as _se: - _se_msg = str(_se) - logger.warning( - "Agent %s: 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) - _sfp, _sfc, _ = self._extract_token_counts(response) - _accumulated_prompt += _sfp - _accumulated_completion += _sfc - - response_text = str(response.content) - - # Use accumulated token counts from all ainvoke() calls - # (main model rounds + pruning passes) instead of only the - # final response (fix for issue #65 Gap 1 & Gap 2). - # If no ainvoke() calls were made (pre-ainvoke failure), - # _accumulated_prompt / _accumulated_completion remain 0 - # which is correct (no tokens consumed). - _prompt_tokens = _accumulated_prompt - _completion_tokens = _accumulated_completion - - # Log a summary of accumulated token usage for observability. - if _accumulated_prompt > 0 or _accumulated_completion > 0: - logger.info( - "Agent %s: accumulated token usage: " - "%d prompt + %d completion tokens across all rounds", - self.name, - _accumulated_prompt, - _accumulated_completion, - ) + if self._lc_tools is not None and LANGCHAIN_AVAILABLE: + # Tools path: full multi-turn loop via _execute_tool_loop(). + # Billing-integrity: on _ToolLoopError, capture partial counts + # before re-raising so the outer exception handlers preserve them. + try: + _loop_result = await self._execute_tool_loop(messages) + except _ToolLoopError as _tle: + if _tle.any_invocation_made: + _captured_prompt = _tle.accumulated_prompt + _captured_completion = _tle.accumulated_completion + raise _tle.cause from None + response_text = str(_loop_result.final_response.content) + _prompt_tokens = _loop_result.accumulated_prompt + _completion_tokens = _loop_result.accumulated_completion + else: + # No tools: single plain ainvoke() — no tool dispatch or synthesis. + _no_tools_resp = await self.chat_model.ainvoke(messages) + _pt, _ct, _ = self._extract_token_counts(_no_tools_resp) + _prompt_tokens = _pt + _completion_tokens = _ct + response_text = str(_no_tools_resp.content) # Capture token counts immediately after ainvoke() succeeds. # Setting the sentinel variables here marks the "ainvoke succeeded" @@ -1566,6 +1711,58 @@ class LLMAgent(AgentWithMemory): # Add current user message lc_messages.append(HumanMessage(content=processed_message)) + _memory_enabled: bool = self.config.get("memory_enabled", False) + + if self._lc_tools is not None and LANGCHAIN_AVAILABLE: + # ── Tools path: single-chunk output via _execute_tool_loop ── + # When tools are configured, the full multi-turn ainvoke() + # loop runs synchronously via _execute_tool_loop() and the + # complete final answer is yielded as a single content chunk. + # Real token-by-token astream() is preserved for the no-tools + # path below. See issue #67 for the accepted single-chunk + # trade-off rationale. + try: + _loop_result = await self._execute_tool_loop(lc_messages) + except _ToolLoopError as _tle: + # Billing-integrity: capture partial accumulated counts + # before re-raising so the outer exception handlers can + # preserve them (mirrors process_message() sentinel logic). + if _tle.any_invocation_made: + _captured_prompt = _tle.accumulated_prompt + _captured_completion = _tle.accumulated_completion + raise _tle.cause from None + _tool_token_str = str(_loop_result.final_response.content) + _captured_prompt = _loop_result.accumulated_prompt + _captured_completion = _loop_result.accumulated_completion + self._last_token_usage = (_captured_prompt, _captured_completion) + last_token_usage_var.set((_captured_prompt, _captured_completion)) + + yield _tool_token_str + + # Update memory if enabled (same logic as no-tools path). + if _memory_enabled: + await self.update_memory("last_message", message) + await self.update_memory("last_response", _tool_token_str) + + _tool_mem_history: list[dict[str, str]] = await self.get_memory( + "conversation_history", [] + ) + _tool_mem_history.append( + {"role": "user", "content": processed_message} + ) + _tool_mem_history.append( + {"role": "assistant", "content": _tool_token_str} + ) + _tool_max_history: int = self.config.get( + "max_history", DEFAULT_MAX_HISTORY + ) + if len(_tool_mem_history) > _tool_max_history: + _tool_mem_history = _tool_mem_history[-_tool_max_history:] + await self.update_memory("conversation_history", _tool_mem_history) + + return # generator done; finally block still runs + + # ── No-tools path: real token-by-token astream() ───────────── # Stream tokens — accumulate full response for memory update only # when memory_enabled is True. Accumulating unconditionally would # grow agent_response_parts to 100K+ entries for long responses @@ -1573,7 +1770,6 @@ class LLMAgent(AgentWithMemory): # consumed (m1 fix). # Use a list accumulator and join at the end to avoid O(n²) string # allocations from repeated += on immutable Python strings. - _memory_enabled: bool = self.config.get("memory_enabled", False) last_chunk: Any = None agent_response_parts: list[str] = [] async for chunk in self.chat_model.astream(lc_messages):