-
Notifications
You must be signed in to change notification settings - Fork 2.2k
Python: feat(core): add tool concurrency groups and sequential execution order #7881
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
e3efb1e
b8b1f03
7783ec2
fb4cb4f
e554130
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -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. | ||
|
|
@@ -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 | ||
|
|
@@ -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) | ||
|
|
@@ -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.""" | ||
|
|
@@ -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 | ||
|
|
||
|
|
||
|
|
@@ -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: ... | ||
|
|
||
|
|
@@ -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]: ... | ||
|
|
||
|
|
@@ -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. | ||
|
|
@@ -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 | ||
|
|
@@ -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) | ||
|
|
@@ -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( | ||
|
|
@@ -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) | ||
|
|
@@ -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, | ||
|
|
@@ -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}" | ||
| 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: | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. This assumes |
||
| 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 | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. When one group raises
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
|
||
| 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 | ||
|
|
@@ -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 | ||
|
|
@@ -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: | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe 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 |
||
| run_config["tool_execution_order"] = custom_args["tool_execution_order"] | ||
|
|
||
| tools = _extract_tools(options) | ||
| if not tools: | ||
| return _FunctionExecutionBatch(result_groups=[]) | ||
|
|
@@ -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, | ||
|
|
@@ -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) | ||
|
|
||
There was a problem hiding this comment.
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_groupvalues. For example, a named group__ungrouped_1collides 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.