diff --git a/python/packages/core/agent_framework/_tools.py b/python/packages/core/agent_framework/_tools.py index ce9606d3de..688b4c5381 100644 --- a/python/packages/core/agent_framework/_tools.py +++ b/python/packages/core/agent_framework/_tools.py @@ -319,6 +319,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. @@ -333,10 +334,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 @@ -347,30 +349,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) @@ -415,6 +424,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.""" @@ -898,9 +908,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 @@ -1137,6 +1151,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: ... @@ -1153,6 +1168,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]: ... @@ -1168,6 +1184,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. @@ -1212,6 +1229,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 @@ -1312,6 +1334,7 @@ def wrapper(f: Callable[..., Any]) -> FunctionTool: func=f, input_model=schema, result_parser=result_parser, + concurrency_group=concurrency_group, ) return wrapper(func) @@ -1377,6 +1400,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( @@ -1390,6 +1414,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) @@ -1820,25 +1845,54 @@ 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_results = await asyncio.gather(*execution_tasks) + 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 = tool_map.get(function_call.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 - should_terminate = any(terminate for _, terminate in execution_results) - return [result_contents for result_contents, _ in execution_results], should_terminate + 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 @@ -3259,11 +3313,19 @@ def get_response( raw_session = request_kwargs.get("session") invocation_session = raw_session if isinstance(raw_session, _AgentSession) else None + # 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 {} + # Bind one executor with the run's custom arguments, middleware, configuration, and session. + request_config = dict(self.function_invocation_configuration) + if tool_exec_order := mutable_options.pop("tool_execution_order", None): + request_config["tool_execution_order"] = tool_exec_order + execute_function_calls = partial( _execute_function_calls, custom_args=additional_function_arguments, - config=self.function_invocation_configuration, + config=request_config, invocation_session=invocation_session, middleware_pipeline=function_middleware_pipeline, ) diff --git a/python/packages/core/agent_framework/_types.py b/python/packages/core/agent_framework/_types.py index 7a0a982ba8..a9fa592af3 100644 --- a/python/packages/core/agent_framework/_types.py +++ b/python/packages/core/agent_framework/_types.py @@ -3552,6 +3552,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 diff --git a/python/packages/core/tests/core/test_tools.py b/python/packages/core/tests/core/test_tools.py index a24e20cacd..bd45d7ee18 100644 --- a/python/packages/core/tests/core/test_tools.py +++ b/python/packages/core/tests/core/test_tools.py @@ -20,6 +20,7 @@ _auto_invoke_function, _parse_annotation, _parse_inputs, + _try_execute_function_call_groups, normalize_function_invocation_configuration, ) from agent_framework.observability import OtelAttr @@ -1547,3 +1548,106 @@ def test_skip_parsing_is_singleton() -> None: # endregion + + +def test_tool_decorator_accepts_concurrency_group(): + """Test that the @tool decorator accepts and stores the concurrency_group parameter.""" + + @tool(name="grouped_tool", concurrency_group="file_system") + def grouped_tool(x: int) -> int: + return x + + assert isinstance(grouped_tool, FunctionTool) + assert grouped_tool.concurrency_group == "file_system" + + +def test_function_invocation_configuration_accepts_execution_order(): + """Test that execution_order is accepted and defaults to 'parallel'.""" + config_seq = normalize_function_invocation_configuration({"tool_execution_order": "sequential"}) + assert config_seq["tool_execution_order"] == "sequential" + + config_default = normalize_function_invocation_configuration(None) + assert config_default["tool_execution_order"] == "parallel" + + +async def test_try_execute_function_call_groups_concurrency_group(): + """Tools in the same concurrency_group execute sequentially; ungrouped tools run concurrently.""" + execution_order = [] + + @tool(concurrency_group="files") + async def write_file(name: str): + execution_order.append("write_start") + await asyncio.sleep(0.05) + execution_order.append("write_end") + return f"wrote {name}" + + @tool(concurrency_group="files") + async def read_file(name: str): + execution_order.append("read_start") + await asyncio.sleep(0.01) + execution_order.append("read_end") + return f"read {name}" + + @tool() + async def ungrouped_tool(): + execution_order.append("ungrouped_start") + await asyncio.sleep(0.02) + execution_order.append("ungrouped_end") + return "ungrouped" + + # Create function call contents simulating a batch from the LLM + call_write = Content.from_function_call(call_id="1", name="write_file", arguments='{"name": "test"}') + call_read = Content.from_function_call(call_id="2", name="read_file", arguments='{"name": "test"}') + call_ungrouped = Content.from_function_call(call_id="3", name="ungrouped_tool", arguments="{}") + + config = normalize_function_invocation_configuration(None) + + results, should_terminate = await _try_execute_function_call_groups( + custom_args={}, + function_calls=[call_write, call_read, call_ungrouped], + tools=[write_file, read_file, ungrouped_tool], + config=config, + ) + + assert not should_terminate + assert len(results) == 3 + + assert execution_order.index("write_end") < execution_order.index("read_start") + assert execution_order.index("ungrouped_start") < execution_order.index("write_end") + + +async def test_try_execute_function_call_groups_sequential_config(): + """When execution_order is 'sequential', ALL tools run one-by-one regardless of groups.""" + execution_order = [] + + @tool() + async def tool_a(): + execution_order.append("a_start") + await asyncio.sleep(0.03) + execution_order.append("a_end") + return "a" + + @tool() + async def tool_b(): + execution_order.append("b_start") + await asyncio.sleep(0.01) + execution_order.append("b_end") + return "b" + + call_a = Content.from_function_call(call_id="1", name="tool_a", arguments="{}") + call_b = Content.from_function_call(call_id="2", name="tool_b", arguments="{}") + + config = normalize_function_invocation_configuration({"tool_execution_order": "sequential"}) + + results, should_terminate = await _try_execute_function_call_groups( + custom_args={}, + function_calls=[call_a, call_b], + tools=[tool_a, tool_b], + config=config, + ) + + assert not should_terminate + assert execution_order == ["a_start", "a_end", "b_start", "b_end"] + + +# endregion