Coverage for haystack/components/agents/tool_calling.py: 98%
250 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-21 13:53 +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 _unwrap_optional
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.
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_content_tag("haystack.agent.step.tool.output", {"error": str(e)})
222 return e
223 span.set_content_tag("haystack.agent.step.tool.output", result)
224 return result
226 def _runner() -> Any:
227 return ctx.run(_invoke)
229 return _runner
232def _make_bounded_invoke_async(
233 tool: Tool,
234 args: dict[str, Any],
235 semaphore: asyncio.Semaphore,
236 tool_call: ToolCall,
237 parent_span: tracing.Span | None,
238) -> Callable[[], Any]:
239 """
240 Return a zero-arg async callable that awaits `tool.invoke_async(**args)` while holding `semaphore`.
242 Concurrency is bounded uniformly across native-async tools and sync-fallback tools (which dispatch
243 to a worker thread inside `Tool.invoke_async`). ContextVars naturally inherit into child tasks for
244 the native-async branch, and `asyncio.to_thread` propagates them for the fallback branch, so spans opened by
245 the tool itself stay attached to the calling trace.
247 Returns a `ToolInvocationError` instead of raising so that gathered executions can collect failures
248 without aborting the whole batch.
249 """
251 async def _runner() -> Any:
252 async with semaphore:
253 with _create_tool_span(tool, tool_call, parent_span) as span:
254 span.set_content_tag("haystack.agent.step.tool.input", tool_call.arguments)
255 try:
256 result = await tool.invoke_async(**args)
257 except ToolInvocationError as e:
258 span.set_content_tag("haystack.agent.step.tool.output", {"error": str(e)})
259 return e
260 span.set_content_tag("haystack.agent.step.tool.output", result)
261 return result
263 return _runner
266def _get_func_params(tool: Tool) -> dict[str, Any]:
267 """
268 Return parameter names → annotations for a tool's invocation function.
270 - For ComponentTool, this is the annotated input schema defined on the underlying component.
271 - For regular Tools, this is the function signature of the `function` callable, falling back to `async_function`
272 for async-only tools.
274 :param tool: The tool to inspect.
275 :returns: A dict mapping parameter names to their type annotations.
276 """
277 if isinstance(tool, ComponentTool):
278 assert hasattr(tool._component, "__haystack_input__") and isinstance(
279 tool._component.__haystack_input__, Sockets
280 )
281 return {name: socket.type for name, socket in tool._component.__haystack_input__._sockets_dict.items()}
282 # Tool.__post_init__ guarantees that at least one of `function` / `async_function` is set.
283 target = tool.function if tool.function is not None else tool.async_function
284 return _resolve_parameter_types(target) # type: ignore[arg-type]
287def _inject_state_args(tool: Tool, llm_args: dict[str, Any], state: State) -> dict[str, Any]:
288 """
289 Merge LLM-provided arguments with state-sourced arguments.
291 LLM args take precedence. State values are pulled in only for the keys a tool explicitly declares via its
292 `inputs_from_state` mapping, then the live State object is injected for any param annotated as State.
294 :param tool: The tool being invoked, used to determine parameter mappings and State injection.
295 :param llm_args: The arguments provided by the LLM, which take precedence over state values.
296 :param state: The current runtime state, used to source additional arguments as needed.
297 :returns: A dict of arguments to invoke the tool with, combining LLM and state values according to the rules
298 described above.
299 """
300 final_args = dict(llm_args)
301 func_params = _get_func_params(tool)
303 # A tool reads from State by name only via an explicit `inputs_from_state` mapping
304 for state_key, param_name in (tool.inputs_from_state or {}).items():
305 if param_name not in final_args and state.has(state_key):
306 final_args[param_name] = state.get(state_key)
308 # We also inject the full State object for any parameter annotated as State
309 for param_name, param_type in func_params.items():
310 if _unwrap_optional(param_type) is State:
311 final_args[param_name] = state
313 return final_args
316def _prepare_tool_args(
317 *,
318 tool: Tool,
319 tool_call_arguments: dict[str, Any],
320 state: State,
321 streaming_callback: StreamingCallbackT | None = None,
322 enable_streaming_passthrough: bool = False,
323) -> dict[str, Any]:
324 """
325 Prepare the final arguments for a tool by injecting state inputs and optionally a streaming callback.
327 :param tool:
328 The tool instance to prepare arguments for.
329 :param tool_call_arguments:
330 The initial arguments provided for the tool call.
331 :param state:
332 The current state containing inputs to be injected into the tool arguments.
333 :param streaming_callback:
334 Optional streaming callback to be injected if enabled and applicable.
335 :param enable_streaming_passthrough:
336 Flag indicating whether to inject the streaming callback into the tool arguments.
338 :returns:
339 A dictionary of final arguments ready for tool invocation.
340 """
341 # Combine user + state inputs
342 final_args = _inject_state_args(tool, tool_call_arguments.copy(), state)
343 # Check whether to inject streaming_callback
344 if (
345 enable_streaming_passthrough
346 and streaming_callback is not None
347 and "streaming_callback" not in final_args
348 and "streaming_callback" in _get_func_params(tool)
349 ):
350 final_args["streaming_callback"] = streaming_callback
351 return final_args
354def _resolve_tool_calls(
355 messages_with_tool_calls: list[ChatMessage], tools_with_names: dict[str, Tool], *, raise_on_failure: bool
356) -> tuple[list[ToolCall], list[Tool], list[ChatMessage]]:
357 """
358 Walk all tool calls in `messages_with_tool_calls` and resolve each to its Tool.
360 Argument preparation is deliberately *not* done here: args are prepared per execution batch (see
361 `_schedule_tool_calls`) so that a tool reading from State observes writes made by tools that ran earlier in the same
362 step.
364 :returns: (tool_calls, resolved_tools, error_messages)
365 - tool_calls: ToolCall objects for each valid call, in call order
366 - resolved_tools: the resolved Tool for each entry in `tool_calls` (parallel list)
367 - error_messages: ChatMessages for tool-not-found errors (when raise_on_failure is False)
368 """
369 tool_calls: list[ToolCall] = []
370 resolved_tools: list[Tool] = []
371 error_messages: list[ChatMessage] = []
373 for message in messages_with_tool_calls:
374 for tool_call in message.tool_calls:
375 tool_name = tool_call.tool_name
377 if tool_name not in tools_with_names:
378 error = ToolNotFoundException(tool_name, list(tools_with_names.keys()))
379 if raise_on_failure:
380 raise error
381 logger.error("{error_exception}", error_exception=error)
382 error_messages.append(ChatMessage.from_tool(tool_result=str(error), origin=tool_call, error=True))
383 continue
385 tool_calls.append(tool_call)
386 resolved_tools.append(tools_with_names[tool_name])
388 return tool_calls, resolved_tools, error_messages
391def _keys_intersect(a: _StateKeys, b: _StateKeys) -> bool:
392 """
393 Return whether two State-key sets share at least one key, treating `_ALL_STATE_KEYS` as a wildcard.
395 Used to detect read-after-write dependencies between tool calls: the reader's read set is tested against the
396 writer's write set.
398 :param a: A set of state keys, or the `_ALL_STATE_KEYS` wildcard meaning "every key".
399 :param b: A set of state keys, or the `_ALL_STATE_KEYS` wildcard meaning "every key".
400 :returns: True if the sets overlap (a wildcard overlaps any non-empty set, and two wildcards always overlap).
401 """
402 if a is _ALL_STATE_KEYS:
403 # `a` covers every key, so it overlaps `b` as long as `b` touches any key. Two wildcards always overlap;
404 # otherwise `bool(b)` is True iff the concrete set `b` is non-empty.
405 return b is _ALL_STATE_KEYS or bool(b)
406 if b is _ALL_STATE_KEYS:
407 # Symmetric case: wildcard `b` overlaps `a` iff the concrete set `a` is non-empty (`bool(set)` == non-empty).
408 return bool(a)
409 # Both are concrete sets: they overlap iff their set intersection is non-empty.
410 return bool(a & b) # type: ignore[operator]
413def _state_io_for_call(tool: Tool, llm_args: dict[str, Any]) -> tuple[_StateKeys, _StateKeys]:
414 """
415 Compute the State keys a tool call reads from and writes to.
417 Mirrors the resolution logic in `_inject_state_args`:
418 - A tool with a `State`-annotated parameter can read/write any key, so both sets are the `_ALL_STATE_KEYS` wildcard.
419 - Otherwise reads come from the tool's explicit `inputs_from_state` mapping, excluding any parameter the LLM already
420 supplied (LLM args take precedence and short-circuit the state lookup).
421 - Writes are the keys in `outputs_to_state`.
423 :returns: A `(reads, writes)` tuple of state-key sets (or the `_ALL_STATE_KEYS` wildcard).
424 """
425 func_params = _get_func_params(tool)
426 # Check if State is in func_params
427 if any(_unwrap_optional(param_type) is State for param_type in func_params.values()):
428 return _ALL_STATE_KEYS, _ALL_STATE_KEYS
430 # Calculate reads
431 param_mappings = tool.inputs_from_state or {}
432 reads = {state_key for state_key, param_name in param_mappings.items() if param_name not in llm_args}
433 # Calculate writes
434 writes = set((tool.outputs_to_state or {}).keys())
436 return reads, writes
439def _schedule_tool_calls(tool_calls: list[ToolCall], tools: list[Tool]) -> list[list[int]]:
440 """
441 Group tool calls into ordered execution batches based on their State read/write sets.
443 Calls within a batch are mutually independent and run in parallel; batches run sequentially. The schedule guarantees
444 that a call reading a State key always runs in a later batch than any call (in the same step) that writes that
445 key — so read-after-write dependencies are honored regardless of the order the LLM requested the calls in.
447 This is a layered topological sort: each round, every call whose dependencies have all been scheduled forms the
448 next parallel batch. Dependency cycles — e.g. a tool that both reads and writes the same key, requested more than
449 once — cannot be ordered by the read-after-write rule alone, so they are broken deterministically by call order
450 (the lowest-index remaining call runs next, on its own).
452 Pure write-write overlaps create no dependency: nobody reads the contended key, and outputs are merged into State
453 sequentially in call order afterward, so the result stays deterministic without serializing execution.
455 :param tool_calls: The tool calls to schedule, in call order.
456 :param tools: The resolved Tool for each entry in `tool_calls` (parallel list).
457 :returns: A list of batches, each a list of indices into `tool_calls`.
458 """
459 # Per-call (reads, writes) State-key sets, in call order.
460 io_list = [_state_io_for_call(tool, tc.arguments) for tc, tool in zip(tool_calls, tools, strict=True)]
461 n = len(io_list)
463 # deps[j] = indices that must run before j because j reads a key they write (read-after-write).
464 deps: list[set[int]] = [set() for _ in range(n)]
465 for j in range(n):
466 reads_j, _ = io_list[j]
467 for i in range(n):
468 if i == j:
469 continue
470 _, writes_i = io_list[i]
471 if _keys_intersect(reads_j, writes_i):
472 deps[j].add(i)
474 scheduled = [False] * n
475 done: set[int] = set()
476 batches: list[list[int]] = []
478 while len(done) < n:
479 # A call is ready once every writer it depends on has already been scheduled (`deps[k] <= done`, i.e. its
480 # dependency set is a subset of the already-done set). All ready calls have no dependency on each other —
481 # if one read a key another writes, it would still be waiting — so the whole `ready` list runs in parallel.
482 ready = [k for k in range(n) if not scheduled[k] and deps[k] <= done]
483 if not ready:
484 # A dependency cycle remains: break it deterministically by running the lowest-index call next.
485 ready = [next(k for k in range(n) if not scheduled[k])]
486 for k in ready:
487 scheduled[k] = True
488 done.update(ready)
489 batches.append(ready)
491 return batches
494def _finalize_tool_result(
495 result: Any, tool_call: ToolCall, tool: Tool, state: State, *, raise_on_failure: bool
496) -> ChatMessage:
497 """
498 Turn a single tool invocation result into a tool-result ChatMessage, merging outputs into State.
500 On a `ToolInvocationError`, either re-raise (when `raise_on_failure`) or return an error message. Otherwise
501 merge the tool's outputs into State (in call order, so write-write merges stay deterministic) and build the
502 result message.
503 """
504 if isinstance(result, ToolInvocationError):
505 if raise_on_failure:
506 raise result
507 logger.error("{error_exception}", error_exception=result)
508 return ChatMessage.from_tool(tool_result=str(result), origin=tool_call, error=True)
510 _merge_tool_outputs_into_state(tool, result, state)
511 return _build_tool_result_message(result, tool_call, tool, raise_on_failure=raise_on_failure)
514def _run_tool(
515 *,
516 messages: list[ChatMessage],
517 state: State,
518 tools: ToolsType,
519 streaming_callback: StreamingCallbackT | None = None,
520 raise_on_failure: bool = True,
521 enable_streaming_callback_passthrough: bool = False,
522 max_workers: int = 4,
523) -> tuple[list[ChatMessage], State]:
524 """
525 Invoke all tools referenced by tool calls in `messages`.
527 :param messages: ChatMessage objects that may contain tool calls.
528 :param state: Runtime state passed to and updated by tools.
529 :param tools: The tools available for invocation.
530 :param streaming_callback: Called once per tool result as it becomes available.
531 :param raise_on_failure: If True, raise on tool invocation failure; otherwise return an error message.
532 :param enable_streaming_callback_passthrough: If True, pass the streaming callback to tools that accept it.
533 :param max_workers: Maximum number of parallel tool invocations.
534 :returns: (tool_messages, updated_state)
535 """
536 tools_with_names = _validate_and_prepare_tools(tools)
538 messages_with_tool_calls = [m for m in messages if m.tool_calls]
539 if not messages_with_tool_calls:
540 return [], state
542 tool_calls, resolved_tools, error_messages = _resolve_tool_calls(
543 messages_with_tool_calls, tools_with_names, raise_on_failure=raise_on_failure
544 )
545 if not tool_calls:
546 return error_messages, state
548 # Group the calls into batches that honor read-after-write dependencies on State (see `_schedule_tool_calls`).
549 batches = _schedule_tool_calls(tool_calls, resolved_tools)
551 # Results are indexed by call position so the returned messages stay in call order, even though batches may
552 # execute the calls in a different order.
553 results: list[ChatMessage | None] = [None] * len(tool_calls)
554 stream_index = 0
556 # Resolved once, before any tool runs, so all per-call spans of this step become siblings under the same parent.
557 parent_span = tracing.tracer.current_span()
559 with ThreadPoolExecutor(max_workers=max_workers) as executor:
560 for batch in batches:
561 # Prepare args at the start of each batch so tools that read from State observe writes merged by earlier
562 # batches.
563 futures = {}
564 for idx in batch:
565 args = _prepare_tool_args(
566 tool=resolved_tools[idx],
567 tool_call_arguments=tool_calls[idx].arguments,
568 state=state,
569 streaming_callback=streaming_callback,
570 enable_streaming_passthrough=enable_streaming_callback_passthrough,
571 )
572 futures[idx] = executor.submit(
573 _make_context_bound_invoke(
574 tool=resolved_tools[idx], args=args, tool_call=tool_calls[idx], parent_span=parent_span
575 )
576 )
578 # Merge results in call order within the batch so write-write merges stay deterministic.
579 for idx in batch:
580 message = _finalize_tool_result(
581 futures[idx].result(),
582 tool_calls[idx],
583 resolved_tools[idx],
584 state,
585 raise_on_failure=raise_on_failure,
586 )
587 results[idx] = message
588 if streaming_callback is not None:
589 streaming_callback(_create_tool_result_streaming_chunk(message, tool_calls[idx], stream_index))
590 stream_index += 1
592 tool_messages = error_messages + [m for m in results if m is not None]
594 # We emit a final empty chunk with finish_reason "tool_call_results" to signal the end of the tool results stream.
595 if tool_messages and streaming_callback is not None:
596 streaming_callback(
597 StreamingChunk(content="", finish_reason="tool_call_results", meta={"finish_reason": "tool_call_results"})
598 )
600 return tool_messages, state
603async def _run_tool_async(
604 *,
605 messages: list[ChatMessage],
606 state: State,
607 tools: ToolsType,
608 streaming_callback: StreamingCallbackT | None = None,
609 raise_on_failure: bool = True,
610 enable_streaming_callback_passthrough: bool = False,
611 max_workers: int = 4,
612) -> tuple[list[ChatMessage], State]:
613 """
614 Asynchronous variant of `run_tool`. Tool calls execute concurrently via a thread pool.
616 :param messages: ChatMessage objects that may contain tool calls.
617 :param state: Runtime state passed to and updated by tools.
618 :param tools: The tools available for invocation.
619 :param streaming_callback: Async callback called once per tool result.
620 :param raise_on_failure: If True, raise on tool invocation failure; otherwise return an error message.
621 :param enable_streaming_callback_passthrough: If True, pass the streaming callback to tools that accept it.
622 :param max_workers: Maximum number of parallel tool invocations.
623 :returns: (tool_messages, updated_state)
624 """
625 tools_with_names = _validate_and_prepare_tools(tools)
627 messages_with_tool_calls = [m for m in messages if m.tool_calls]
628 if not messages_with_tool_calls:
629 return [], state
631 tool_calls, resolved_tools, error_messages = _resolve_tool_calls(
632 messages_with_tool_calls, tools_with_names, raise_on_failure=raise_on_failure
633 )
634 if not tool_calls:
635 return error_messages, state
637 # Group the calls into batches that honor read-after-write dependencies on State (see `_schedule_tool_calls`).
638 batches = _schedule_tool_calls(tool_calls, resolved_tools)
640 # Results are indexed by call position so the returned messages stay in call order, even though batches may
641 # execute the calls in a different order.
642 results: list[ChatMessage | None] = [None] * len(tool_calls)
643 stream_index = 0
645 # `max_workers` + Semaphore bounds concurrency for both sync and async tool calls async tools are awaited directly,
646 # and sync tools are dispatched to a worker thread inside `Tool.invoke_async`.
647 semaphore = asyncio.Semaphore(max_workers)
649 # Resolved once, before any tool runs, so all per-call spans of this step become siblings under the same parent.
650 parent_span = tracing.tracer.current_span()
652 for batch in batches:
653 # Prepare args at the start of each batch so readers observe writes merged by earlier batches.
654 tasks = {}
655 for idx in batch:
656 args = _prepare_tool_args(
657 tool=resolved_tools[idx],
658 tool_call_arguments=tool_calls[idx].arguments,
659 state=state,
660 streaming_callback=streaming_callback,
661 enable_streaming_passthrough=enable_streaming_callback_passthrough,
662 )
663 tasks[idx] = _make_bounded_invoke_async(
664 tool=resolved_tools[idx],
665 args=args,
666 semaphore=semaphore,
667 tool_call=tool_calls[idx],
668 parent_span=parent_span,
669 )()
670 batch_results = await asyncio.gather(*tasks.values())
672 # Merge results in call order within the batch so write-write merges stay deterministic.
673 for idx, result in zip(tasks.keys(), batch_results, strict=True):
674 message = _finalize_tool_result(
675 result, tool_calls[idx], resolved_tools[idx], state, raise_on_failure=raise_on_failure
676 )
677 results[idx] = message
678 if streaming_callback is not None:
679 await _invoke_streaming_callback(
680 streaming_callback, _create_tool_result_streaming_chunk(message, tool_calls[idx], stream_index)
681 )
682 stream_index += 1
684 tool_messages = error_messages + [m for m in results if m is not None]
686 # We emit a final empty chunk with finish_reason "tool_call_results" to signal the end of the tool results stream.
687 if tool_messages and streaming_callback is not None:
688 await _invoke_streaming_callback(
689 streaming_callback,
690 StreamingChunk(content="", finish_reason="tool_call_results", meta={"finish_reason": "tool_call_results"}),
691 )
693 return tool_messages, state