From 5535415e8982066feee3695912cc3c2de0c645c2 Mon Sep 17 00:00:00 2001 From: Tomas Sereikis Date: Sat, 29 Aug 2026 13:21:54 +0300 Subject: [PATCH] perf(llm): run one turn's tool calls concurrently instead of serially MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A single assistant turn routinely requests several independent reads — a dialectic turn typically emits `search_memory` and `search_messages` together — but the loop awaited them one at a time, so the turn cost their sum when it only needed to cost the slowest. Measured on a production `minimal` dialectic: 1.2s in `search_memory`, then a further 1.9s in `search_messages`, both pure reads. Measured across 42 production dialectic requests (103 iterations, 77 of them multi-tool, 3.34 tool calls per iteration on average), running each turn's calls concurrently would cut 11.3% of total wall-clock: median 9.3% per request, p90 29.7%, best case 51.0%. Only 4 of the 42 requests gain nothing. This is safe because tool handlers own their sessions: each opens a short-lived one via `tracked_db()`, and only the mutating handlers (`create_observations`, `update_peer_card`, `delete_observations`) take `ctx.db_lock`, so concurrent reads do not contend on shared state. The per-call telemetry ContextVars move inside the task. `asyncio` copies the context per task, so `set_current_tool_call_seq` and `set_last_tool_metadata` now bind to their own call instead of being written and read across one shared context — which the previous code could only keep straight by never overlapping. `gather` preserves argument order, so `tool_results` and `all_tool_calls` stay in the order the model asked for them. Fan-out is capped at MAX_CONCURRENT_TOOL_CALLS (4). Production has produced 18 tool calls in one iteration, and firing all of them at once would mean that many simultaneous embedding + pgvector queries on a single instance. The cap leaves the measured common case fully parallel while bounding the tail. --- src/llm/tool_loop.py | 68 ++++++++--- tests/llm/test_tool_loop_parallel.py | 173 +++++++++++++++++++++++++++ 2 files changed, 224 insertions(+), 17 deletions(-) create mode 100644 tests/llm/test_tool_loop_parallel.py diff --git a/src/llm/tool_loop.py b/src/llm/tool_loop.py index 783d4965..7a30b960 100644 --- a/src/llm/tool_loop.py +++ b/src/llm/tool_loop.py @@ -11,6 +11,7 @@ from __future__ import annotations +import asyncio import dataclasses import functools import logging @@ -201,6 +202,13 @@ def _emit_agent_iteration( logger = logging.getLogger(__name__) +# A single assistant turn can request many tools at once — production has seen +# 18 in one iteration. Running every one concurrently would put that many +# simultaneous embedding + pgvector queries on a single instance, so the fan-out +# is capped: the measured common case is 3.3 calls per iteration, which stays +# fully parallel, while the tail is bounded instead of stampeding. +MAX_CONCURRENT_TOOL_CALLS = 4 + # Bounds for max_tool_iterations to prevent runaway loops. MIN_TOOL_ITERATIONS = 1 MAX_TOOL_ITERATIONS = 100 @@ -391,6 +399,9 @@ async def execute_tool_loop( if telemetry is not None and telemetry.hash_memo is None: telemetry = dataclasses.replace(telemetry, hash_memo={}) + # One semaphore for the whole loop: iterations run one after another, so a + # single cap bounds the fan-out without being rebound per iteration. + tool_call_semaphore = asyncio.Semaphore(MAX_CONCURRENT_TOOL_CALLS) iteration = 0 all_tool_calls: list[dict[str, Any]] = [] total_input_tokens = 0 @@ -582,54 +593,77 @@ async def execute_tool_loop( # Telemetry context — 1-indexed iteration. set_current_iteration(iteration + 1) - tool_results: list[dict[str, Any]] = [] - for seq, tool_call in enumerate(response.tool_calls_made): + # Tools requested in one assistant turn are independent of each + # other, so run them concurrently and pay the slowest instead of the + # sum. A dialectic turn routinely asks for `search_memory` and + # `search_messages` together; serially that was 1.2s + 1.9s of pure + # read. Handlers open their own short-lived sessions via + # `tracked_db()` and only the mutating ones (`create_observations`, + # `update_peer_card`, `delete_observations`) take `ctx.db_lock`, so + # concurrent reads are safe. + async def run_tool_call( + seq: int, tool_call: dict[str, Any] + ) -> tuple[dict[str, Any], dict[str, Any] | None]: tool_name = tool_call["name"] tool_input = tool_call["input"] tool_id = tool_call.get("id", "") logger.debug(f"Executing tool: {tool_name}") - # the executor closure reads these from - # ContextVars to populate AgentToolCallCompletedEvent. Set BEFORE - # the executor call so two calls to the same tool in one iteration - # get distinct seq values. Reset last-tool metadata so we never - # observe stale state from a prior call. + # The executor closure reads these from ContextVars to populate + # AgentToolCallCompletedEvent. They are set INSIDE the task: + # asyncio copies the context per task, so two calls to the same + # tool in one iteration keep their own seq and their own + # last-tool metadata rather than racing on one shared context. set_current_tool_call_seq(seq, tool_id or None) set_last_tool_metadata({}) try: - tool_result = await tool_executor(tool_name, tool_input) + async with tool_call_semaphore: + tool_result = await tool_executor(tool_name, tool_input) # Stash ToolResult.metadata on all_tool_calls so # specialist rollups can read created/deleted observation # counts without round-tripping through the event store. - tool_result_metadata = get_last_tool_metadata() - tool_results.append( + return ( { "tool_id": tool_id, "tool_name": tool_name, "result": tool_result, - } - ) - all_tool_calls.append( + }, { "tool_name": tool_name, "tool_input": tool_input, "tool_result": tool_result, - "tool_result_metadata": tool_result_metadata, - } + "tool_result_metadata": get_last_tool_metadata(), + }, ) except Exception as e: logger.error(f"Tool execution failed for {tool_name}: {e}") - tool_results.append( + return ( { "tool_id": tool_id, "tool_name": tool_name, "result": f"Error: {str(e)}", "is_error": True, - } + }, + None, ) + # gather preserves argument order, so tool_results and + # all_tool_calls stay in the order the model asked for them. + outcomes = await asyncio.gather( + *( + asyncio.create_task(run_tool_call(seq, tool_call)) + for seq, tool_call in enumerate(response.tool_calls_made) + ) + ) + + tool_results: list[dict[str, Any]] = [] + for result_entry, call_entry in outcomes: + tool_results.append(result_entry) + if call_entry is not None: + all_tool_calls.append(call_entry) + append_tool_results(current_provider, tool_results, conversation_messages) finally: if step is not None: diff --git a/tests/llm/test_tool_loop_parallel.py b/tests/llm/test_tool_loop_parallel.py new file mode 100644 index 00000000..ac1408d9 --- /dev/null +++ b/tests/llm/test_tool_loop_parallel.py @@ -0,0 +1,173 @@ +# pyright: reportPrivateUsage=false, reportUnknownLambdaType=false, reportArgumentType=false +"""Tool calls emitted in a single assistant turn execute concurrently. + +The model routinely asks for several independent reads at once — a dialectic +turn typically emits `search_memory` and `search_messages` together. Running +them one after another makes the turn cost the sum of their latencies when it +only needs to cost the slowest: measured against production, a `minimal` +dialectic spent 1.2s in `search_memory` and then a further 1.9s in +`search_messages`, both pure reads. + +Tool handlers open their own short-lived sessions via `tracked_db()` and only +the mutating handlers (`create_observations`, `update_peer_card`, +`delete_observations`) take `ctx.db_lock`, so concurrent reads are safe. The +per-call telemetry ContextVars are set inside each task, and `asyncio` copies +the context per task, so sequence numbers and last-tool metadata stay bound to +their own call instead of racing on one shared context. +""" + +from __future__ import annotations + +import asyncio +from typing import Any +from unittest.mock import patch + +import pytest + +from src.llm import tool_loop +from src.llm.runtime import AttemptPlan +from src.llm.tool_loop import MAX_CONCURRENT_TOOL_CALLS, execute_tool_loop +from src.llm.types import HonchoLLMCallResponse + +TOOL_DELAY_SECONDS = 0.2 + + +def _make_plan() -> AttemptPlan: + return AttemptPlan( + provider="anthropic", + model="claude-sonnet-4-5", + client=object(), + thinking_budget_tokens=None, + reasoning_effort=None, + selected_config=None, + attempt=1, + retry_attempts=1, + is_fallback=False, + ) + + +def _response(tool_calls: list[dict[str, Any]]) -> HonchoLLMCallResponse[Any]: + return HonchoLLMCallResponse( + content="done", + input_tokens=10, + output_tokens=5, + cache_creation_input_tokens=0, + cache_read_input_tokens=0, + finish_reasons=["stop"], + tool_calls_made=tool_calls, + ) + + +class _Tracker: + """Counts how many tool executions are in flight at once.""" + + def __init__(self) -> None: + self.in_flight = 0 + self.peak = 0 + + async def execute(self, _name: str, _input: dict[str, Any]) -> str: + self.in_flight += 1 + self.peak = max(self.peak, self.in_flight) + try: + await asyncio.sleep(TOOL_DELAY_SECONDS) + finally: + self.in_flight -= 1 + return "ok" + + +async def _run(tracker: _Tracker, n_calls: int = 2) -> HonchoLLMCallResponse[Any]: + names = ["search_memory", "search_messages"] + calls = [ + {"name": names[i % len(names)], "input": {"query": str(i)}, "id": f"t{i}"} + for i in range(n_calls) + ] + responses = iter([_response(calls), _response([])]) + + async def _call(*_args: Any, **_kwargs: Any) -> HonchoLLMCallResponse[Any]: + return next(responses) + + with ( + patch.object(tool_loop, "honcho_llm_call_inner", new=_call), + patch("src.llm.conversation.count_message_tokens", return_value=10), + patch( + "src.llm.conversation.truncate_messages_to_fit", + side_effect=lambda msgs, _cap: msgs, + ), + ): + return await execute_tool_loop( + prompt="hi", + max_tokens=64, + messages=[{"role": "user", "content": "q"}], + tools=[ + { + "name": "search_memory", + "description": "", + "input_schema": {"type": "object"}, + }, + { + "name": "search_messages", + "description": "", + "input_schema": {"type": "object"}, + }, + ], + tool_choice="auto", + tool_executor=tracker.execute, + max_tool_iterations=5, + response_model=None, + json_mode=False, + temperature=None, + stop_seqs=None, + verbosity=None, + enable_retry=False, + retry_attempts=1, + max_input_tokens=1000, + get_attempt_plan=_make_plan, + before_retry_callback=lambda _r: None, + stream_final=False, + telemetry=None, + ) + + +@pytest.mark.asyncio +async def test_tool_calls_in_one_iteration_run_concurrently(): + """Two tools requested in one turn overlap rather than queueing.""" + tracker = _Tracker() + + started = asyncio.get_running_loop().time() + await _run(tracker) + elapsed = asyncio.get_running_loop().time() - started + + assert tracker.peak == 2, f"tools ran sequentially (peak in-flight {tracker.peak})" + # Sequential would be >= 2 * delay; concurrent stays near one delay. + assert elapsed < TOOL_DELAY_SECONDS * 1.8, ( + f"elapsed {elapsed:.3f}s suggests serial execution" + ) + + +@pytest.mark.asyncio +async def test_results_keep_request_order(): + """Concurrency must not reorder results relative to the calls.""" + tracker = _Tracker() + + result = await _run(tracker) + + names = [call["tool_name"] for call in result.tool_calls_made] + assert names == ["search_memory", "search_messages"] + + +@pytest.mark.asyncio +async def test_fan_out_is_capped(): + """A turn asking for many tools does not stampede the database. + + Production has produced 18 tool calls in a single iteration; firing all of + them at once would mean that many concurrent embedding + pgvector queries + on one instance. + """ + tracker = _Tracker() + + await _run(tracker, n_calls=18) + + assert tracker.peak <= MAX_CONCURRENT_TOOL_CALLS, ( + f"fan-out reached {tracker.peak}, above the {MAX_CONCURRENT_TOOL_CALLS} cap" + ) + assert tracker.peak > 1, "cap must not serialise execution entirely"