Coverage for haystack/components/agents/tool_calling.py: 99%
252 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-12 09:42 +0000
1# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
2#
3# SPDX-License-Identifier: Apache-2.0
5import asyncio
6import contextvars
7import json
8from collections.abc import Callable
9from concurrent.futures import ThreadPoolExecutor
10from typing import Any
12from haystack import logging, tracing
13from haystack.components.agents.state.state import State
14from haystack.core.component.sockets import Sockets
15from haystack.core.type_utils import _resolve_parameter_types
16from haystack.dataclasses import ChatMessage, ToolCall
17from haystack.dataclasses.streaming_chunk import StreamingCallbackT, StreamingChunk, _invoke_streaming_callback
18from haystack.tools import ComponentTool, Tool, ToolsType, _check_duplicate_tool_names, flatten_tools_or_toolsets
19from haystack.tools.errors import ToolInvocationError
20from haystack.tools.parameters_schema_utils import _is_state_annotation
21from haystack.tracing.utils import _serializable_value
23logger = logging.getLogger(__name__)
26class _AllStateKeys:
27 """Sentinel representing "every state key", used for tools that receive the full State object."""
29 def __repr__(self) -> str:
30 return "<all state keys>"
33_ALL_STATE_KEYS = _AllStateKeys()
35# A set of state keys, or the _ALL_STATE_KEYS sentinel meaning "every key".
36_StateKeys = set[str] | _AllStateKeys
39class ToolNotFoundException(Exception):
40 """Exception raised when a tool is not found in the list of available tools."""
42 def __init__(self, tool_name: str, available_tools: list[str]) -> None:
43 message = f"Tool '{tool_name}' not found. Available tools: {', '.join(available_tools)}"
44 super().__init__(message)
47def _validate_and_prepare_tools(tools: ToolsType) -> dict[str, Tool]:
48 """
49 Flatten, deduplicate-check, and index tools by name.
51 :raises ValueError: If no tools are provided or if duplicate tool names are found.
52 """
53 if not tools:
54 raise ValueError("Tool execution requires at least one tool.")
56 available_tools = flatten_tools_or_toolsets(tools)
57 _check_duplicate_tool_names(available_tools)
58 tool_names = [tool.name for tool in available_tools]
60 return dict(zip(tool_names, available_tools, strict=True))
63def _merge_tool_outputs_into_state(tool: Tool, result: Any, state: State) -> None:
64 """
65 Write tool outputs into State according to the tool's `outputs_to_state` mapping.
67 :raises RuntimeError: If writing an output value into the state fails.
68 """
69 if not isinstance(result, dict):
70 return
72 for state_key, config in (tool.outputs_to_state or {}).items():
73 source_key = config.get("source", None)
74 if source_key and source_key not in result:
75 continue
76 output_value = result.get(source_key) if source_key else result
77 try:
78 state.set(state_key, output_value, handler_override=config.get("handler"))
79 except Exception as e:
80 raise RuntimeError(f"Tool '{tool.name}': failed to merge outputs into state. {e}") from e
83def _result_to_string(result: Any) -> str:
84 """
85 Convert a tool result to a string.
87 Strings are returned as-is; all other types are passed through a JSON serialization step to produce more readable
88 output, with a fallback to plain str() conversion if serialization fails.
90 :param result: The tool result to convert.
91 :returns: A string representation of the tool result.
92 """
93 if isinstance(result, str):
94 return result
95 serializable = _serializable_value(value=result, use_placeholders=False)
96 try:
97 return json.dumps(serializable, ensure_ascii=False)
98 except Exception as error:
99 logger.warning(
100 "Tool result is not JSON serializable. Falling back to str conversion. Result: {result}\nError: {err}",
101 result=result,
102 err=error,
103 )
104 return str(result)
107def _process_tool_output(config: dict[str, Any], result: Any, tool_call: ToolCall, *, raise_on_failure: bool) -> Any:
108 """
109 Extract and convert a single tool output according to `config`.
111 `config` may contain `source` (key to extract from result dict), `handler` (conversion callable), and
112 `raw_result` (return the value without string conversion).
114 If a configured `handler` raises, the exception is re-raised when `raise_on_failure` is True; otherwise
115 a warning is logged and the value is converted via `_result_to_string`.
116 """
117 source_key = config.get("source")
118 value = result.get(source_key) if source_key is not None and isinstance(result, dict) else result
120 handler = config.get("handler")
121 raw_result = config.get("raw_result", False)
123 if handler is None:
124 # raw result is mostly used to allow ImageContent or TextContent blocks to be directly returned and consumed
125 # by ChatMessage.from_tool without string conversion.
126 if raw_result:
127 return value
128 return _result_to_string(value)
130 try:
131 return handler(value)
132 except Exception as e:
133 if raise_on_failure:
134 raise
135 logger.warning(
136 "Output handler '{handler}' for tool '{tool}' failed, falling back to string conversion. Error: {err}",
137 handler=handler.__name__,
138 tool=tool_call.tool_name,
139 err=e,
140 )
141 return _result_to_string(value)
144def _build_tool_result_message(result: Any, tool_call: ToolCall, tool: Tool, *, raise_on_failure: bool) -> ChatMessage:
145 """Convert a raw tool result into a ChatMessage, applying `outputs_to_string` config if present."""
146 outputs_config = tool.outputs_to_string or {}
148 # Single-output config (or no config): keys are at the root level
149 if not outputs_config or any(k in outputs_config for k in ("source", "handler", "raw_result")):
150 tool_result = _process_tool_output(outputs_config, result, tool_call, raise_on_failure=raise_on_failure)
151 return ChatMessage.from_tool(tool_result=tool_result, origin=tool_call)
153 # Multi-output config: each key maps to its own sub-config — stringify each value, then stringify the whole dict
154 tool_result_dict = {
155 output_key: _process_tool_output(
156 {**cfg, "raw_result": False}, result, tool_call, raise_on_failure=raise_on_failure
157 )
158 for output_key, cfg in outputs_config.items()
159 }
160 return ChatMessage.from_tool(tool_result=_result_to_string(tool_result_dict), origin=tool_call)
163def _create_tool_result_streaming_chunk(tool_message: ChatMessage, tool_call: ToolCall, index: int) -> StreamingChunk:
164 """
165 Create a streaming chunk that carries a tool result.
167 :param tool_message: The tool result message to stream.
168 :param tool_call: The ToolCall object that triggered the tool invocation.
169 :param index: The position of this tool result in the stream (in execution order).
170 :returns: A StreamingChunk containing the tool result and metadata about the tool call.
171 """
172 return StreamingChunk(
173 content="",
174 index=index,
175 tool_call_result=tool_message.tool_call_results[0],
176 start=True,
177 meta={"tool_result": tool_message.tool_call_results[0].result, "tool_call": tool_call},
178 )
181def _create_tool_span(tool: Tool, tool_call: ToolCall, parent_span: tracing.Span | None) -> Any:
182 """
183 Create one tracing span for a single tool call, nested under `parent_span`.
185 The established standard for tracing agents is one span per tool call, so each call gets its own span rather than
186 grouping all of a step's calls together. The OpenTelemetry GenAI semantic conventions codify this with a dedicated
187 "execute tool" span per invocation
188 (https://github.com/open-telemetry/semantic-conventions-genai/blob/main/docs/gen-ai/gen-ai-agent-spans.md#execute-tool-span),
189 and tracing backends such as Langfuse follow the same model. The span is tagged with the tool's identity; the caller
190 adds the call arguments and result as content tags, and sets `haystack.tool.error` to `True` when the call fails.
192 The parent is passed in rather than read from the tracer here: the calls of a step run concurrently, so a tracer
193 that tracks the active span globally would report a sibling tool span as current and the spans would end up chained
194 instead of side by side.
195 """
196 return tracing.tracer.trace(
197 "haystack.agent.step.tool",
198 tags={"haystack.tool.name": tool_call.tool_name, "haystack.tool.description": tool.description},
199 parent_span=parent_span,
200 )
203def _make_context_bound_invoke(
204 tool: Tool, args: dict[str, Any], tool_call: ToolCall, parent_span: tracing.Span | None
205) -> Callable[[], Any]:
206 """
207 Return a zero-arg callable that runs `tool.invoke(**args)` under the current contextvars snapshot.
209 This preserves tracing spans and other context-local state across thread-pool boundaries, so spans opened by the
210 tool itself stay attached to the calling trace. The callable returns a ToolInvocationError instead of raising so
211 that parallel executions can collect failures without aborting the whole batch.
212 """
213 ctx = contextvars.copy_context()
215 def _invoke() -> Any:
216 with _create_tool_span(tool, tool_call, parent_span) as span:
217 span.set_content_tag("haystack.agent.step.tool.input", tool_call.arguments)
218 try:
219 result = tool.invoke(**args)
220 except ToolInvocationError as e:
221 span.set_tag("haystack.tool.error", True)
222 span.set_content_tag("haystack.agent.step.tool.output", {"error": str(e)})
223 return e
224 span.set_content_tag("haystack.agent.step.tool.output", result)
225 return result
227 def _runner() -> Any:
228 return ctx.run(_invoke)
230 return _runner
233def _make_bounded_invoke_async(
234 tool: Tool,
235 args: dict[str, Any],
236 semaphore: asyncio.Semaphore,
237 tool_call: ToolCall,
238 parent_span: tracing.Span | None,
239) -> Callable[[], Any]:
240 """
241 Return a zero-arg async callable that awaits `tool.invoke_async(**args)` while holding `semaphore`.
243 Concurrency is bounded uniformly across native-async tools and sync-fallback tools (which dispatch
244 to a worker thread inside `Tool.invoke_async`). ContextVars naturally inherit into child tasks for
245 the native-async branch, and `asyncio.to_thread` propagates them for the fallback branch, so spans opened by
246 the tool itself stay attached to the calling trace.
248 Returns a `ToolInvocationError` instead of raising so that gathered executions can collect failures
249 without aborting the whole batch.
250 """
252 async def _runner() -> Any:
253 async with semaphore:
254 with _create_tool_span(tool, tool_call, parent_span) as span:
255 span.set_content_tag("haystack.agent.step.tool.input", tool_call.arguments)
256 try:
257 result = await tool.invoke_async(**args)
258 except ToolInvocationError as e:
259 span.set_tag("haystack.tool.error", True)
260 span.set_content_tag("haystack.agent.step.tool.output", {"error": str(e)})
261 return e
262 span.set_content_tag("haystack.agent.step.tool.output", result)
263 return result
265 return _runner
268def _get_func_params(tool: Tool) -> dict[str, Any]:
269 """
270 Return parameter names → annotations for a tool's invocation function.
272 - For ComponentTool, this is the annotated input schema defined on the underlying component.
273 - For regular Tools, this is the function signature of the `function` callable, falling back to `async_function`
274 for async-only tools.
276 :param tool: The tool to inspect.
277 :returns: A dict mapping parameter names to their type annotations.
278 """
279 if isinstance(tool, ComponentTool):
280 assert hasattr(tool._component, "__haystack_input__") and isinstance(
281 tool._component.__haystack_input__, Sockets
282 )
283 return {name: socket.type for name, socket in tool._component.__haystack_input__._sockets_dict.items()}
284 # Tool.__post_init__ guarantees that at least one of `function` / `async_function` is set.
285 target = tool.function if tool.function is not None else tool.async_function
286 return _resolve_parameter_types(target) # type: ignore[arg-type]
289def _inject_state_args(tool: Tool, llm_args: dict[str, Any], state: State) -> dict[str, Any]:
290 """
291 Merge LLM-provided arguments with state-sourced arguments.
293 LLM args take precedence. State values are pulled in only for the keys a tool explicitly declares via its
294 `inputs_from_state` mapping, then the live State object is injected for any param annotated as State.
296 :param tool: The tool being invoked, used to determine parameter mappings and State injection.
297 :param llm_args: The arguments provided by the LLM, which take precedence over state values.
298 :param state: The current runtime state, used to source additional arguments as needed.
299 :returns: A dict of arguments to invoke the tool with, combining LLM and state values according to the rules
300 described above.
301 """
302 final_args = dict(llm_args)
303 func_params = _get_func_params(tool)
305 # A tool reads from State by name only via an explicit `inputs_from_state` mapping
306 for state_key, param_name in (tool.inputs_from_state or {}).items():
307 if param_name not in final_args and state.has(state_key):
308 final_args[param_name] = state.get(state_key)
310 # We also inject the full State object for any parameter annotated as State
311 for param_name, param_type in func_params.items():
312 if _is_state_annotation(param_type):
313 final_args[param_name] = state
315 return final_args
318def _prepare_tool_args(
319 *,
320 tool: Tool,
321 tool_call_arguments: dict[str, Any],
322 state: State,
323 streaming_callback: StreamingCallbackT | None = None,
324 enable_streaming_passthrough: bool = False,
325) -> dict[str, Any]:
326 """
327 Prepare the final arguments for a tool by injecting state inputs and optionally a streaming callback.
329 :param tool:
330 The tool instance to prepare arguments for.
331 :param tool_call_arguments:
332 The initial arguments provided for the tool call.
333 :param state:
334 The current state containing inputs to be injected into the tool arguments.
335 :param streaming_callback:
336 Optional streaming callback to be injected if enabled and applicable.
337 :param enable_streaming_passthrough:
338 Flag indicating whether to inject the streaming callback into the tool arguments.
340 :returns:
341 A dictionary of final arguments ready for tool invocation.
342 """
343 # Combine user + state inputs
344 final_args = _inject_state_args(tool, tool_call_arguments.copy(), state)
345 # Check whether to inject streaming_callback
346 if (
347 enable_streaming_passthrough
348 and streaming_callback is not None
349 and "streaming_callback" not in final_args
350 and "streaming_callback" in _get_func_params(tool)
351 ):
352 final_args["streaming_callback"] = streaming_callback
353 return final_args
356def _resolve_tool_calls(
357 messages_with_tool_calls: list[ChatMessage], tools_with_names: dict[str, Tool], *, raise_on_failure: bool
358) -> tuple[list[ToolCall], list[Tool], list[ChatMessage]]:
359 """
360 Walk all tool calls in `messages_with_tool_calls` and resolve each to its Tool.
362 Argument preparation is deliberately *not* done here: args are prepared per execution batch (see
363 `_schedule_tool_calls`) so that a tool reading from State observes writes made by tools that ran earlier in the same
364 step.
366 :returns: (tool_calls, resolved_tools, error_messages)
367 - tool_calls: ToolCall objects for each valid call, in call order
368 - resolved_tools: the resolved Tool for each entry in `tool_calls` (parallel list)
369 - error_messages: ChatMessages for tool-not-found errors (when raise_on_failure is False)
370 """
371 tool_calls: list[ToolCall] = []
372 resolved_tools: list[Tool] = []
373 error_messages: list[ChatMessage] = []
375 for message in messages_with_tool_calls:
376 for tool_call in message.tool_calls:
377 tool_name = tool_call.tool_name
379 if tool_name not in tools_with_names:
380 error = ToolNotFoundException(tool_name, list(tools_with_names.keys()))
381 if raise_on_failure:
382 raise error
383 logger.error("{error_exception}", error_exception=error)
384 error_messages.append(ChatMessage.from_tool(tool_result=str(error), origin=tool_call, error=True))
385 continue
387 tool_calls.append(tool_call)
388 resolved_tools.append(tools_with_names[tool_name])
390 return tool_calls, resolved_tools, error_messages
393def _keys_intersect(a: _StateKeys, b: _StateKeys) -> bool:
394 """
395 Return whether two State-key sets share at least one key, treating `_ALL_STATE_KEYS` as a wildcard.
397 Used to detect read-after-write dependencies between tool calls: the reader's read set is tested against the
398 writer's write set.
400 :param a: A set of state keys, or the `_ALL_STATE_KEYS` wildcard meaning "every key".
401 :param b: A set of state keys, or the `_ALL_STATE_KEYS` wildcard meaning "every key".
402 :returns: True if the sets overlap (a wildcard overlaps any non-empty set, and two wildcards always overlap).
403 """
404 if a is _ALL_STATE_KEYS:
405 # `a` covers every key, so it overlaps `b` as long as `b` touches any key. Two wildcards always overlap;
406 # otherwise `bool(b)` is True iff the concrete set `b` is non-empty.
407 return b is _ALL_STATE_KEYS or bool(b)
408 if b is _ALL_STATE_KEYS:
409 # Symmetric case: wildcard `b` overlaps `a` iff the concrete set `a` is non-empty (`bool(set)` == non-empty).
410 return bool(a)
411 # Both are concrete sets: they overlap iff their set intersection is non-empty.
412 return bool(a & b) # type: ignore[operator]
415def _state_io_for_call(tool: Tool, llm_args: dict[str, Any]) -> tuple[_StateKeys, _StateKeys]:
416 """
417 Compute the State keys a tool call reads from and writes to.
419 Mirrors the resolution logic in `_inject_state_args`:
420 - A tool with a `State`-annotated parameter can read/write any key, so both sets are the `_ALL_STATE_KEYS` wildcard.
421 - Otherwise reads come from the tool's explicit `inputs_from_state` mapping, excluding any parameter the LLM already
422 supplied (LLM args take precedence and short-circuit the state lookup).
423 - Writes are the keys in `outputs_to_state`.
425 :returns: A `(reads, writes)` tuple of state-key sets (or the `_ALL_STATE_KEYS` wildcard).
426 """
427 func_params = _get_func_params(tool)
428 # Check if State is in func_params
429 if any(_is_state_annotation(param_type) for param_type in func_params.values()):
430 return _ALL_STATE_KEYS, _ALL_STATE_KEYS
432 # Calculate reads
433 param_mappings = tool.inputs_from_state or {}
434 reads = {state_key for state_key, param_name in param_mappings.items() if param_name not in llm_args}
435 # Calculate writes
436 writes = set((tool.outputs_to_state or {}).keys())
438 return reads, writes
441def _schedule_tool_calls(tool_calls: list[ToolCall], tools: list[Tool]) -> list[list[int]]:
442 """
443 Group tool calls into ordered execution batches based on their State read/write sets.
445 Calls within a batch are mutually independent and run in parallel; batches run sequentially. The schedule guarantees
446 that a call reading a State key always runs in a later batch than any call (in the same step) that writes that
447 key — so read-after-write dependencies are honored regardless of the order the LLM requested the calls in.
449 This is a layered topological sort: each round, every call whose dependencies have all been scheduled forms the
450 next parallel batch. Dependency cycles — e.g. a tool that both reads and writes the same key, requested more than
451 once — cannot be ordered by the read-after-write rule alone, so they are broken deterministically by call order
452 (the lowest-index remaining call runs next, on its own).
454 Write-write overlaps alone create no dependency. Outputs are merged into State after each batch, in the
455 LLM's requested order within that batch, so subsequent batches see the updated values. This keeps merging
456 deterministic, but does not guarantee that writes follow the LLM's requested order across batches.
458 :param tool_calls: The tool calls to schedule, in call order.
459 :param tools: The resolved Tool for each entry in `tool_calls` (parallel list).
460 :returns: A list of batches, each a list of indices into `tool_calls`.
461 """
462 # Per-call (reads, writes) State-key sets, in call order.
463 io_list = [_state_io_for_call(tool, tc.arguments) for tc, tool in zip(tool_calls, tools, strict=True)]
464 n = len(io_list)
466 # deps[j] = indices that must run before j because j reads a key they write (read-after-write).
467 deps: list[set[int]] = [set() for _ in range(n)]
468 for j in range(n):
469 reads_j, _ = io_list[j]
470 for i in range(n):
471 if i == j:
472 continue
473 _, writes_i = io_list[i]
474 if _keys_intersect(reads_j, writes_i):
475 deps[j].add(i)
477 scheduled = [False] * n
478 done: set[int] = set()
479 batches: list[list[int]] = []
481 while len(done) < n:
482 # A call is ready once every writer it depends on has already been scheduled (`deps[k] <= done`, i.e. its
483 # dependency set is a subset of the already-done set). All ready calls have no dependency on each other —
484 # if one read a key another writes, it would still be waiting — so the whole `ready` list runs in parallel.
485 ready = [k for k in range(n) if not scheduled[k] and deps[k] <= done]
486 if not ready:
487 # A dependency cycle remains: break it deterministically by running the lowest-index call next.
488 ready = [next(k for k in range(n) if not scheduled[k])]
489 for k in ready:
490 scheduled[k] = True
491 done.update(ready)
492 batches.append(ready)
494 return batches
497def _finalize_tool_result(
498 result: Any, tool_call: ToolCall, tool: Tool, state: State, *, raise_on_failure: bool
499) -> ChatMessage:
500 """
501 Turn a single tool invocation result into a tool-result ChatMessage, merging outputs into State.
503 On a `ToolInvocationError`, either re-raise (when `raise_on_failure`) or return an error message. Otherwise
504 merge the tool's outputs into State and build the result message.
505 """
506 if isinstance(result, ToolInvocationError):
507 if raise_on_failure:
508 raise result
509 logger.error("{error_exception}", error_exception=result)
510 return ChatMessage.from_tool(tool_result=str(result), origin=tool_call, error=True)
512 _merge_tool_outputs_into_state(tool, result, state)
513 return _build_tool_result_message(result, tool_call, tool, raise_on_failure=raise_on_failure)
516def _run_tool(
517 *,
518 messages: list[ChatMessage],
519 state: State,
520 tools: ToolsType,
521 streaming_callback: StreamingCallbackT | None = None,
522 raise_on_failure: bool = True,
523 enable_streaming_callback_passthrough: bool = False,
524 max_workers: int = 4,
525) -> tuple[list[ChatMessage], State]:
526 """
527 Invoke all tools referenced by tool calls in `messages`.
529 :param messages: ChatMessage objects that may contain tool calls.
530 :param state: Runtime state passed to and updated by tools.
531 :param tools: The tools available for invocation.
532 :param streaming_callback: Called once per tool result as it becomes available.
533 :param raise_on_failure: If True, raise on tool invocation failure; otherwise return an error message.
534 :param enable_streaming_callback_passthrough: If True, pass the streaming callback to tools that accept it.
535 :param max_workers: Maximum number of parallel tool invocations.
536 :returns: (tool_messages, updated_state)
537 """
538 tools_with_names = _validate_and_prepare_tools(tools)
540 messages_with_tool_calls = [m for m in messages if m.tool_calls]
541 if not messages_with_tool_calls:
542 return [], state
544 tool_calls, resolved_tools, error_messages = _resolve_tool_calls(
545 messages_with_tool_calls, tools_with_names, raise_on_failure=raise_on_failure
546 )
547 if not tool_calls:
548 return error_messages, state
550 # Group the calls into batches that honor read-after-write dependencies on State (see `_schedule_tool_calls`).
551 batches = _schedule_tool_calls(tool_calls, resolved_tools)
553 # Results are indexed by call position so the returned messages stay in call order, even though batches may
554 # execute the calls in a different order.
555 results: list[ChatMessage | None] = [None] * len(tool_calls)
556 stream_index = 0
558 # Resolved once, before any tool runs, so all per-call spans of this step become siblings under the same parent.
559 parent_span = tracing.tracer.current_span()
561 with ThreadPoolExecutor(max_workers=max_workers) as executor:
562 for batch in batches:
563 # Prepare args at the start of each batch so tools that read from State observe writes merged by earlier
564 # batches.
565 futures = {}
566 for idx in batch:
567 args = _prepare_tool_args(
568 tool=resolved_tools[idx],
569 tool_call_arguments=tool_calls[idx].arguments,
570 state=state,
571 streaming_callback=streaming_callback,
572 enable_streaming_passthrough=enable_streaming_callback_passthrough,
573 )
574 futures[idx] = executor.submit(
575 _make_context_bound_invoke(
576 tool=resolved_tools[idx], args=args, tool_call=tool_calls[idx], parent_span=parent_span
577 )
578 )
580 # Merge results in call order within the batch so write-write merges stay deterministic.
581 for idx in batch:
582 message = _finalize_tool_result(
583 futures[idx].result(),
584 tool_calls[idx],
585 resolved_tools[idx],
586 state,
587 raise_on_failure=raise_on_failure,
588 )
589 results[idx] = message
590 if streaming_callback is not None:
591 streaming_callback(_create_tool_result_streaming_chunk(message, tool_calls[idx], stream_index))
592 stream_index += 1
594 tool_messages = error_messages + [m for m in results if m is not None]
596 # We emit a final empty chunk with finish_reason "tool_call_results" to signal the end of the tool results stream.
597 if tool_messages and streaming_callback is not None:
598 streaming_callback(
599 StreamingChunk(content="", finish_reason="tool_call_results", meta={"finish_reason": "tool_call_results"})
600 )
602 return tool_messages, state
605async def _run_tool_async(
606 *,
607 messages: list[ChatMessage],
608 state: State,
609 tools: ToolsType,
610 streaming_callback: StreamingCallbackT | None = None,
611 raise_on_failure: bool = True,
612 enable_streaming_callback_passthrough: bool = False,
613 max_workers: int = 4,
614) -> tuple[list[ChatMessage], State]:
615 """
616 Asynchronous variant of `run_tool`. Tool calls execute concurrently via a thread pool.
618 :param messages: ChatMessage objects that may contain tool calls.
619 :param state: Runtime state passed to and updated by tools.
620 :param tools: The tools available for invocation.
621 :param streaming_callback: Async callback called once per tool result.
622 :param raise_on_failure: If True, raise on tool invocation failure; otherwise return an error message.
623 :param enable_streaming_callback_passthrough: If True, pass the streaming callback to tools that accept it.
624 :param max_workers: Maximum number of parallel tool invocations.
625 :returns: (tool_messages, updated_state)
626 """
627 tools_with_names = _validate_and_prepare_tools(tools)
629 messages_with_tool_calls = [m for m in messages if m.tool_calls]
630 if not messages_with_tool_calls:
631 return [], state
633 tool_calls, resolved_tools, error_messages = _resolve_tool_calls(
634 messages_with_tool_calls, tools_with_names, raise_on_failure=raise_on_failure
635 )
636 if not tool_calls:
637 return error_messages, state
639 # Group the calls into batches that honor read-after-write dependencies on State (see `_schedule_tool_calls`).
640 batches = _schedule_tool_calls(tool_calls, resolved_tools)
642 # Results are indexed by call position so the returned messages stay in call order, even though batches may
643 # execute the calls in a different order.
644 results: list[ChatMessage | None] = [None] * len(tool_calls)
645 stream_index = 0
647 # `max_workers` + Semaphore bounds concurrency for both sync and async tool calls async tools are awaited directly,
648 # and sync tools are dispatched to a worker thread inside `Tool.invoke_async`.
649 semaphore = asyncio.Semaphore(max_workers)
651 # Resolved once, before any tool runs, so all per-call spans of this step become siblings under the same parent.
652 parent_span = tracing.tracer.current_span()
654 for batch in batches:
655 # Prepare args at the start of each batch so readers observe writes merged by earlier batches.
656 tasks = {}
657 for idx in batch:
658 args = _prepare_tool_args(
659 tool=resolved_tools[idx],
660 tool_call_arguments=tool_calls[idx].arguments,
661 state=state,
662 streaming_callback=streaming_callback,
663 enable_streaming_passthrough=enable_streaming_callback_passthrough,
664 )
665 tasks[idx] = _make_bounded_invoke_async(
666 tool=resolved_tools[idx],
667 args=args,
668 semaphore=semaphore,
669 tool_call=tool_calls[idx],
670 parent_span=parent_span,
671 )()
672 batch_results = await asyncio.gather(*tasks.values())
674 # Merge results in call order within the batch so write-write merges stay deterministic.
675 for idx, result in zip(tasks.keys(), batch_results, strict=True):
676 message = _finalize_tool_result(
677 result, tool_calls[idx], resolved_tools[idx], state, raise_on_failure=raise_on_failure
678 )
679 results[idx] = message
680 if streaming_callback is not None:
681 await _invoke_streaming_callback(
682 streaming_callback, _create_tool_result_streaming_chunk(message, tool_calls[idx], stream_index)
683 )
684 stream_index += 1
686 tool_messages = error_messages + [m for m in results if m is not None]
688 # We emit a final empty chunk with finish_reason "tool_call_results" to signal the end of the tool results stream.
689 if tool_messages and streaming_callback is not None:
690 await _invoke_streaming_callback(
691 streaming_callback,
692 StreamingChunk(content="", finish_reason="tool_call_results", meta={"finish_reason": "tool_call_results"}),
693 )
695 return tool_messages, state