diff --git a/src/dialectic/core.py b/src/dialectic/core.py index 57964c87..6842fbe1 100644 --- a/src/dialectic/core.py +++ b/src/dialectic/core.py @@ -35,6 +35,7 @@ from src.telemetry.prometheus.metrics import DialecticComponents, TokenTypes from src.utils.agent_tools import ( DIALECTIC_TOOLS, DIALECTIC_TOOLS_MINIMAL, + MUTATING_TOOL_NAMES, create_tool_executor, search_memory, ) @@ -508,6 +509,11 @@ class DialecticAgent: tool_choice=level_settings.TOOL_CHOICE, tool_executor=tool_executor, max_tool_iterations=level_settings.MAX_TOOL_ITERATIONS, + # Dialectic's loadout is reads today, but say so explicitly rather + # than relying on that: the loop keeps any state-changing call + # in the order the model asked for, so adding a write tool here + # later stays correct without anyone remembering this. + mutating_tools=MUTATING_TOOL_NAMES, messages=self.messages, max_input_tokens=settings.DIALECTIC.MAX_INPUT_TOKENS, trace_name="dialectic_chat", @@ -584,6 +590,11 @@ class DialecticAgent: tool_choice=level_settings.TOOL_CHOICE, tool_executor=tool_executor, max_tool_iterations=level_settings.MAX_TOOL_ITERATIONS, + # Dialectic's loadout is reads today, but say so explicitly rather + # than relying on that: the loop keeps any state-changing call + # in the order the model asked for, so adding a write tool here + # later stays correct without anyone remembering this. + mutating_tools=MUTATING_TOOL_NAMES, messages=self.messages, max_input_tokens=settings.DIALECTIC.MAX_INPUT_TOKENS, trace_name="dialectic_chat", diff --git a/src/llm/api.py b/src/llm/api.py index 13c3a5e5..ca9795d6 100644 --- a/src/llm/api.py +++ b/src/llm/api.py @@ -68,6 +68,7 @@ async def honcho_llm_call( tool_choice: str | dict[str, Any] | None = None, tool_executor: Callable[[str, dict[str, Any]], Any] | None = None, max_tool_iterations: int = 10, + mutating_tools: frozenset[str] | None = None, messages: list[dict[str, Any]] | None = None, max_input_tokens: int | None = None, trace_name: str | None = None, @@ -97,6 +98,7 @@ async def honcho_llm_call( tool_choice: str | dict[str, Any] | None = None, tool_executor: Callable[[str, dict[str, Any]], Any] | None = None, max_tool_iterations: int = 10, + mutating_tools: frozenset[str] | None = None, messages: list[dict[str, Any]] | None = None, max_input_tokens: int | None = None, trace_name: str | None = None, @@ -126,6 +128,7 @@ async def honcho_llm_call( tool_choice: str | dict[str, Any] | None = None, tool_executor: Callable[[str, dict[str, Any]], Any] | None = None, max_tool_iterations: int = 10, + mutating_tools: frozenset[str] | None = None, messages: list[dict[str, Any]] | None = None, max_input_tokens: int | None = None, trace_name: str | None = None, @@ -154,6 +157,7 @@ async def honcho_llm_call( tool_choice: str | dict[str, Any] | None = None, tool_executor: Callable[[str, dict[str, Any]], Any] | None = None, max_tool_iterations: int = 10, + mutating_tools: frozenset[str] | None = None, messages: list[dict[str, Any]] | None = None, max_input_tokens: int | None = None, trace_name: str | None = None, @@ -450,6 +454,7 @@ async def honcho_llm_call( tool_choice=tool_choice, tool_executor=tool_executor, max_tool_iterations=max_tool_iterations, + mutating_tools=mutating_tools, response_model=response_model, json_mode=json_mode, temperature=temperature, diff --git a/src/llm/tool_loop.py b/src/llm/tool_loop.py index 783d4965..8a70291d 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 @@ -345,6 +353,7 @@ async def execute_tool_loop( tool_choice: str | dict[str, Any] | None, tool_executor: Callable[[str, dict[str, Any]], Any], max_tool_iterations: int, + mutating_tools: frozenset[str] | None = None, response_model: type[BaseModel] | None, json_mode: bool, temperature: float | None, @@ -391,6 +400,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 +594,97 @@ 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, ) + # Reads may overlap freely; state-changing calls may not, because + # `ctx.db_lock` decides that they take turns but not which turn each + # one gets. Anything named in `mutating_tools` therefore runs in the + # order the model asked for, while the rest run concurrently + # alongside it. `mutating_tools=None` means "assume every tool + # mutates", so a caller that has not opted in keeps the fully + # sequential behaviour it had before. + calls = list(enumerate(response.tool_calls_made)) + + def _must_stay_ordered(tool_call: dict[str, Any]) -> bool: + return mutating_tools is None or tool_call["name"] in mutating_tools + + # create_task even for the ordered ones: each still needs its own + # context copy for the telemetry ContextVars, and awaiting them one + # at a time is what preserves the order. + concurrent = { + seq: asyncio.create_task(run_tool_call(seq, tool_call)) + for seq, tool_call in calls + if not _must_stay_ordered(tool_call) + } + + outcomes: dict[int, tuple[dict[str, Any], dict[str, Any] | None]] = {} + for seq, tool_call in calls: + if _must_stay_ordered(tool_call): + outcomes[seq] = await run_tool_call(seq, tool_call) + for seq, task in concurrent.items(): + outcomes[seq] = await task + + tool_results: list[dict[str, Any]] = [] + for seq, _ in calls: + result_entry, call_entry = outcomes[seq] + 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/src/utils/agent_tools.py b/src/utils/agent_tools.py index de07e38f..86c2fdd3 100644 --- a/src/utils/agent_tools.py +++ b/src/utils/agent_tools.py @@ -2619,6 +2619,23 @@ _TOOL_HANDLERS: dict[str, Callable[[ToolContext, dict[str, Any]], Any]] = { "get_reasoning_chain": _handle_get_reasoning_chain, } +# Tools that change state, and so must not be reordered relative to one another +# within a single assistant turn. Each of these routes to a handler that takes +# `ctx.db_lock` (`_handle_create_observations_impl`, `_handle_update_peer_card`, +# `_handle_delete_observations`); the lock stops them interleaving, but it does +# not decide who gets it first, so the tool loop keeps them in the order the +# model asked for. `extract_preferences` and `finish_consolidation` are not +# here: both only return text telling the model what to call next. +MUTATING_TOOL_NAMES: frozenset[str] = frozenset( + { + "create_observations", + "create_observations_deductive", + "create_observations_inductive", + "update_peer_card", + "delete_observations", + } +) + async def create_tool_executor( workspace_name: str, diff --git a/tests/llm/test_tool_loop_parallel.py b/tests/llm/test_tool_loop_parallel.py new file mode 100644 index 00000000..f49e66f2 --- /dev/null +++ b/tests/llm/test_tool_loop_parallel.py @@ -0,0 +1,219 @@ +# 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, + mutating_tools: frozenset[str] | None = frozenset(), + names: list[str] | None = None, +) -> HonchoLLMCallResponse[Any]: + names = names or ["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, + mutating_tools=mutating_tools, + 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" + + +@pytest.mark.asyncio +async def test_mutating_tools_are_not_reordered(): + """State-changing calls keep the order the model asked for. + + `ctx.db_lock` makes the mutating handlers take turns but does not decide + whose turn comes first, so running them concurrently could apply a turn's + writes in an order the model did not request. + """ + order: list[str] = [] + + class _Recorder(_Tracker): + async def execute(self, name: str, _input: dict[str, Any]) -> str: + self.in_flight += 1 + self.peak = max(self.peak, self.in_flight) + # Later calls sleep less, so anything concurrent finishes reversed. + await asyncio.sleep(TOOL_DELAY_SECONDS / (len(order) + 1)) + order.append(name) + self.in_flight -= 1 + return "ok" + + await _run( + _Recorder(), + n_calls=3, + mutating_tools=frozenset({"update_peer_card"}), + names=["update_peer_card"], + ) + + assert order == ["update_peer_card"] * 3 + + +@pytest.mark.asyncio +async def test_default_is_fully_sequential(): + """A caller that does not opt in keeps the old one-at-a-time behaviour.""" + tracker = _Tracker() + + await _run(tracker, n_calls=4, mutating_tools=None) + + assert tracker.peak == 1