Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
228 changes: 176 additions & 52 deletions python/packages/core/agent_framework/_tools.py
Original file line number Diff line number Diff line change
Expand Up @@ -321,6 +321,7 @@ def __init__(
func: Callable[..., Any] | None = None,
input_model: type[BaseModel] | Mapping[str, Any] | None = None,
result_parser: Callable[[Any], str | list[Content]] | _SkipParsingSentinel | None = None,
concurrency_group: str | None = None,
**kwargs: Any,
) -> None:
"""Initialize the FunctionTool.
Expand All @@ -335,10 +336,11 @@ def __init__(
max_invocations: The maximum number of times this function can be invoked
across the **lifetime of this tool instance**. If None (default),
there is no limit. Should be at least 1. If the tool is called multiple
times in one iteration, those will execute, after that it will stop working. For example,
if max_invocations is 3 and the tool is called 5 times in a single iteration,
these will complete, but any subsequent calls to the tool (in the same or future iterations)
will raise a ToolException.
times in one iteration, those will execute, after that it will stop
working. For example, if max_invocations is 3 and the tool is called 5
times in a single iteration, these will complete, but any subsequent
calls to the tool (in the same or future iterations) will raise a
ToolException.

.. note::
This counter lives on the tool instance and is never automatically
Expand All @@ -349,30 +351,37 @@ def __init__(
``FunctionInvocationConfiguration["max_function_calls"]``
for per-request limits instead.

max_invocation_exceptions: The maximum number of exceptions allowed during invocations.
If None, there is no limit. Should be at least 1.
max_invocation_exceptions: The maximum number of exceptions allowed
during invocations. If None, there is no limit. Should be at least 1.
additional_properties: Additional properties to set on the function.
func: The function to wrap. When ``None``, creates a declaration-only tool
that has no implementation. Declaration-only tools are useful when you want
the agent to reason about tool usage without executing them, or when the
actual implementation exists elsewhere (e.g., client-side rendering).
input_model: The Pydantic model that defines the input parameters for the function.
This can also be a JSON schema dictionary.
If not provided and ``func`` is not ``None``, it will be inferred from
the function signature. When ``func`` is ``None`` and ``input_model`` is
not provided, the tool will use an empty input model (no parameters) in
its JSON schema. For declaration-only tools that should declare
parameters, explicitly provide ``input_model`` (either a Pydantic
``BaseModel`` or a JSON schema dictionary) so the model can reason about
the expected arguments.
result_parser: An optional callable with signature ``Callable[[Any], str]`` that
overrides the default result parsing behavior. When provided, this callable
is used to convert the raw function return value to a string instead of the
built-in :meth:`parse_result` logic. Pass the :data:`SKIP_PARSING` sentinel
instead of a callable to opt out of parsing entirely; in that case
:meth:`invoke` returns the wrapped function's raw return value. Depending
on your function, it may be easiest to just do the serialization directly
in the function body rather than providing a custom ``result_parser``.
func: The function to wrap. When ``None``, creates a declaration-only
tool that has no implementation. Declaration-only tools are useful
when you want the agent to reason about tool usage without executing
them, or when the actual implementation exists elsewhere (e.g.,
client-side rendering).
input_model: The Pydantic model that defines the input parameters for the
function. This can also be a JSON schema dictionary.
If not provided and ``func`` is not ``None``, it will be inferred
from the function signature. When ``func`` is ``None`` and
``input_model`` is not provided, the tool will use an empty input
model (no parameters) in its JSON schema. For declaration-only tools
that should declare parameters, explicitly provide ``input_model``
(either a Pydantic ``BaseModel`` or a JSON schema dictionary) so the
model can reason about the expected arguments.
result_parser: An optional callable with signature ``Callable[[Any], str]``
that overrides the default result parsing behavior. When provided,
this callable is used to convert the raw function return value to a
string instead of the built-in :meth:`parse_result` logic. Pass the
:data:`SKIP_PARSING` sentinel instead of a callable to opt out of
parsing entirely; in that case :meth:`invoke` returns the wrapped
function's raw return value. Depending on your function, it may be
easiest to just do the serialization directly in the function body
rather than providing a custom ``result_parser``.
concurrency_group: If provided, tool calls with the same
concurrency_group will execute sequentially in the order they were
invoked by the model. Tools without a group, or with different
groups, will execute concurrently. Useful for stateful tools with
write->read dependencies to prevent race conditions.
**kwargs: Additional keyword arguments.
"""
# Core attributes (formerly from BaseTool)
Expand Down Expand Up @@ -417,6 +426,7 @@ def __init__(
self._invocation_duration_histogram = _default_histogram()
self.type: Literal["function_tool"] = "function_tool"
self.result_parser = result_parser
self.concurrency_group = concurrency_group

def _discover_injected_parameters(self) -> None:
"""Inspect the wrapped function for runtime injection parameters."""
Expand Down Expand Up @@ -905,9 +915,13 @@ def to_json_schema_spec(self) -> dict[str, Any]:
@override
def to_dict(self, *, exclude: set[str] | None = None, exclude_none: bool = True) -> dict[str, Any]:
as_dict = super().to_dict(exclude=exclude, exclude_none=exclude_none)
if (not exclude or "concurrency_group" not in exclude) and (
not exclude_none or self.concurrency_group is not None
):
as_dict["concurrency_group"] = self.concurrency_group
if (exclude and "input_model" in exclude) or not self.input_model:
return as_dict
as_dict["input_model"] = self.parameters() # Use cached parameters()
as_dict["input_model"] = self.parameters()
return as_dict


Expand Down Expand Up @@ -1144,6 +1158,7 @@ def tool(
max_invocations: int | None = None,
max_invocation_exceptions: int | None = None,
additional_properties: dict[str, Any] | None = None,
concurrency_group: str | None = None,
result_parser: Callable[[Any], str | list[Content]] | _SkipParsingSentinel | None = None,
) -> FunctionTool: ...

Expand All @@ -1160,6 +1175,7 @@ def tool(
max_invocations: int | None = None,
max_invocation_exceptions: int | None = None,
additional_properties: dict[str, Any] | None = None,
concurrency_group: str | None = None,
result_parser: Callable[[Any], str | list[Content]] | _SkipParsingSentinel | None = None,
) -> Callable[[Callable[..., Any]], FunctionTool]: ...

Expand All @@ -1175,6 +1191,7 @@ def tool(
max_invocations: int | None = None,
max_invocation_exceptions: int | None = None,
additional_properties: dict[str, Any] | None = None,
concurrency_group: str | None = None,
result_parser: Callable[[Any], str | list[Content]] | _SkipParsingSentinel | None = None,
) -> FunctionTool | Callable[[Callable[..., Any]], FunctionTool]:
"""Decorate a function to turn it into a FunctionTool that can be passed to models and executed automatically.
Expand Down Expand Up @@ -1219,6 +1236,11 @@ def tool(
max_invocation_exceptions: The maximum number of exceptions allowed during invocations.
If None, there is no limit, should be at least 1.
additional_properties: Additional properties to set on the function.
concurrency_group: If provided, tool calls with the same
concurrency_group will execute sequentially in the order they were
invoked by the model. Tools without a group, or with different
groups, will execute concurrently. Useful for stateful tools with
write->read dependencies to prevent race conditions.
result_parser: An optional callable with signature ``Callable[[Any], str]`` that
overrides the default result parsing. When provided, this callable converts the
raw function return value to a string instead of using the built-in
Expand Down Expand Up @@ -1319,6 +1341,7 @@ def wrapper(f: Callable[..., Any]) -> FunctionTool:
func=f,
input_model=schema,
result_parser=result_parser,
concurrency_group=concurrency_group,
)

return wrapper(func)
Expand Down Expand Up @@ -1384,6 +1407,7 @@ class FunctionInvocationConfiguration(TypedDict, total=False):
terminate_on_unknown_calls: bool
additional_tools: Sequence[FunctionTool]
include_detailed_errors: bool
tool_execution_order: Literal["parallel", "sequential"]


def normalize_function_invocation_configuration(
Expand All @@ -1397,6 +1421,7 @@ def normalize_function_invocation_configuration(
"terminate_on_unknown_calls": False,
"additional_tools": [],
"include_detailed_errors": False,
"tool_execution_order": "parallel",
}
if config:
normalized.update(config)
Expand Down Expand Up @@ -1777,7 +1802,8 @@ async def _try_execute_function_call_groups(
has_declaration_only_call = False
# A user-input pause takes precedence over unknown-call termination in mixed batches.
for function_call in actionable_calls:
function_name = function_call.name
function_name = _underlying_function_call(function_call).name

logger.debug(
"Checking function call: type=%s, name=%s, in approval_tools=%s",
function_call.type,
Expand Down Expand Up @@ -1845,23 +1871,48 @@ async def _try_execute_function_call_groups(
# Only a fully executable batch reaches this point; run calls concurrently but retain per-call result groups.
# Create each task inside a copied context so the active agent span is
# preserved for every parallel tool invocation.
execution_tasks = [
contextvars.copy_context().run(
asyncio.create_task,
_execute_single_function_call(
function_call,
custom_args=custom_args,
config=config,
tool_map=tool_map,
invocation_session=invocation_session,
middleware_pipeline=middleware_pipeline,
live_tools=live_tools,
),
)
for function_call in function_calls
]
execution_order = config.get("tool_execution_order", "parallel")

groups: dict[str, list[int]] = {}
for idx, function_call in enumerate(function_calls):
group_key: str | None = None
if execution_order == "parallel":
tool_name = _underlying_function_call(function_call).name
if tool_name is not None:
tool = tool_map.get(tool_name)
if tool is not None:
group_key = getattr(tool, "concurrency_group", None)
if group_key is None:
group_key = "__sequential_all__" if execution_order == "sequential" else f"__ungrouped_{idx}"

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Synthetic keys share the same string namespace as user-provided concurrency_group values. For example, a named group __ungrouped_1 collides with the ungrouped call at index 1, unexpectedly serializing them and potentially deadlocking when the first waits for the second. Please use structurally distinct internal keys or object sentinels so user group names cannot collide with scheduler keys.

if group_key not in groups:
groups[group_key] = []
groups[group_key].append(idx)

ordered_results: list[tuple[list[Content], bool] | None] = [None] * len(function_calls)

async def _execute_group(indices: list[int]) -> None:
for idx in indices:

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This assumes function_calls remains in model-request order, but approval resume collects explicit responses before appending the previously hidden safe siblings. A batch emitted as [write, approval-required read] can therefore resume as [read, write], and this serial loop executes the dependency backward. Please preserve each call's original batch ordinal through approval storage and merge resumed calls by that ordinal before grouping.

call = function_calls[idx]
ctx = contextvars.copy_context()
task = ctx.run(
asyncio.create_task,
_execute_single_function_call(
call,
custom_args=custom_args,
config=config,
tool_map=tool_map,
invocation_session=invocation_session,
middleware_pipeline=middleware_pipeline,
live_tools=live_tools,
),
)
res = await task

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

When one group raises MiddlewareFailure, another group that has just completed a call can advance to its next queued tool before the outer gather() observes the failure and cancels the group tasks. That lets a side-effecting call start after a fail-closed middleware abort, contrary to the function-loop contract that no further tool call starts. Please propagate a batch-wide abort signal at individual-call completion and check it before dequeuing each subsequent call.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

_execute_single_function_call returns should_terminate=True for MiddlewareTermination, but this loop records the tuple and immediately starts the next call in the same group. With sequential execution, a policy middleware can request approval or termination and a later side-effecting tool still runs before the flag is checked after the entire batch. Please stop dequeuing calls as soon as a result requests termination and prevent other groups from starting additional queued work.

ordered_results[idx] = res

execution_tasks = [asyncio.create_task(_execute_group(indices)) for indices in groups.values()]

try:
execution_results = await asyncio.gather(*execution_tasks)
await asyncio.gather(*execution_tasks)
except BaseException:
# A loud escape from one call (e.g. MiddlewareFailure aborting the run
# fail-closed) fails the whole batch: cancel in-flight siblings and wait for
Expand All @@ -1875,8 +1926,61 @@ async def _try_execute_function_call_groups(
await asyncio.gather(*execution_tasks, return_exceptions=True)
raise

should_terminate = any(terminate for _, terminate in execution_results)
return [result_contents for result_contents, _ in execution_results], should_terminate
if any(result is None for result in ordered_results):
raise RuntimeError("Internal error: missing tool execution result(s).")

completed_results = cast(list[tuple[list[Content], bool]], ordered_results)
should_terminate = any(terminate for _, terminate in completed_results)
return [result_contents for result_contents, _ in completed_results], should_terminate

groups: dict[str, list[int]] = {}
for idx, function_call in enumerate(function_calls):
group_key: str | None = None

if execution_order == "parallel":
tool_name = _underlying_function_call(function_call).name
if tool_name is not None:
tool = tool_map.get(tool_name)
if tool is not None:
group_key = getattr(tool, "concurrency_group", None)

if group_key is None:
group_key = "__sequential_all__" if execution_order == "sequential" else f"__ungrouped_{idx}"

if group_key not in groups:
groups[group_key] = []
groups[group_key].append(idx)

ordered_results: list[tuple[list[Content], bool] | None] = [None] * len(function_calls)

async def _execute_group(indices: list[int]) -> None:
for idx in indices:
call = function_calls[idx]
ctx = contextvars.copy_context()
task = ctx.run(
asyncio.create_task,
_execute_single_function_call(
call,
custom_args=custom_args,
config=config,
tool_map=tool_map,
invocation_session=invocation_session,
middleware_pipeline=middleware_pipeline,
live_tools=live_tools,
),
)
res = await task
ordered_results[idx] = res

execution_tasks = [asyncio.create_task(_execute_group(indices)) for indices in groups.values()]

await asyncio.gather(*execution_tasks)

if any(result is None for result in ordered_results):
raise RuntimeError("Internal error: missing tool execution result(s).")
completed_results = cast(list[tuple[list[Content], bool]], ordered_results)
should_terminate = any(terminate for _, terminate in completed_results)
return [result_contents for result_contents, _ in completed_results], should_terminate


@dataclass
Expand Down Expand Up @@ -1911,6 +2015,11 @@ async def _execute_function_calls(
invocation_session: AgentSession | None = None,
middleware_pipeline: FunctionMiddlewarePipeline | None = None,
) -> _FunctionExecutionBatch:

run_config = cast("FunctionInvocationConfiguration", dict(config) if config else {})
if custom_args and "tool_execution_order" in custom_args:

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This assignment runs after the invocation layer has applied the explicit per-run option, so function_invocation_kwargs["tool_execution_order"] silently overwrites options["tool_execution_order"] despite the preceding log saying the explicit option wins. It also lets an ordinary injected tool argument change framework scheduling. Please consume this reserved key when constructing run_config, apply it only when no explicit option was supplied, and avoid forwarding it as a tool runtime argument.

run_config["tool_execution_order"] = custom_args["tool_execution_order"]

tools = _extract_tools(options)
if not tools:
return _FunctionExecutionBatch(result_groups=[])
Expand All @@ -1920,7 +2029,7 @@ async def _execute_function_calls(
tools=tools,
invocation_session=invocation_session,
middleware_pipeline=middleware_pipeline,
config=config,
config=run_config,
)
return _FunctionExecutionBatch(
result_groups=result_groups,
Expand Down Expand Up @@ -3625,17 +3734,32 @@ def get_response(
invocation_session = raw_session if isinstance(raw_session, _AgentSession) else None

# Bind one executor with the run's custom arguments, middleware, configuration, and session.
options = dict(options) if options else {}
run_config = cast(
"FunctionInvocationConfiguration",
dict(self.function_invocation_configuration) if self.function_invocation_configuration else {},
)

if not isinstance(options, dict): # pragma: no cover
options = {}
if tool_exec_order := options.pop("tool_execution_order", None):
run_config["tool_execution_order"] = tool_exec_order
if additional_function_arguments and "tool_execution_order" in additional_function_arguments:
logger.debug(
"overriding tool_execution_order from function_invocation_kwargs with explicit run option: %s",
tool_exec_order,
)

mutable_options: dict[str, Any] = dict(options)

execute_function_calls = partial(
_execute_function_calls,
custom_args=additional_function_arguments,
config=self.function_invocation_configuration,
config=run_config,
invocation_session=invocation_session,
middleware_pipeline=function_middleware_pipeline,
)

# Give the loop private mutable options and one shared run-local tool list for progressive tool changes.
# Make options mutable so we can update conversation_id during function invocation loop
mutable_options: dict[str, Any] = dict(options) if options else {}
# Remove additional_function_arguments from options passed to underlying chat client
# It's for tool invocation only and not recognized by chat service APIs
mutable_options.pop("additional_function_arguments", None)
Expand Down
4 changes: 4 additions & 0 deletions python/packages/core/agent_framework/_types.py
Original file line number Diff line number Diff line change
Expand Up @@ -3708,6 +3708,10 @@ class _ChatOptionsBase(TypedDict, total=False):
tool_choice: ToolMode | Literal["auto", "required", "none"]
allow_multiple_tool_calls: bool

# Dictates whether multiple tool calls in a single message batch
# are executed concurrently (parallel) or one-by-one (sequential).
tool_execution_order: Literal["parallel", "sequential"]

# Response configuration
response_format: type[BaseModel] | Mapping[str, Any] | None

Expand Down
Loading
Loading