Python: support mixed workflow invocation kwargs

python/packages/core/agent_framework/__init__.py

@@ -362,7 +362,7 @@
         "validate_workflow_graph",
     ),
     "._workflows._viz": ("WorkflowViz",),
-    "._workflows._workflow": ("Workflow", "WorkflowRunResult"),
+    "._workflows._workflow": ("Workflow", "WorkflowInvocationKwargs", "WorkflowRunResult"),
     "._workflows._workflow_builder": ("WorkflowBuilder",),
     "._workflows._workflow_context": ("WorkflowContext",),
     "._workflows._workflow_executor": (
@@ -626,6 +626,7 @@
     "WorkflowEventType",
     "WorkflowException",
     "WorkflowExecutor",
+    "WorkflowInvocationKwargs",
     "WorkflowMessage",
     "WorkflowRunResult",
     "WorkflowRunState",

python/packages/core/agent_framework/__init__.pyi

@@ -325,7 +325,7 @@ from ._workflows._validation import (
     validate_workflow_graph,
 )
 from ._workflows._viz import WorkflowViz
-from ._workflows._workflow import Workflow, WorkflowRunResult
+from ._workflows._workflow import Workflow, WorkflowInvocationKwargs, WorkflowRunResult
 from ._workflows._workflow_builder import WorkflowBuilder
 from ._workflows._workflow_context import WorkflowContext
 from ._workflows._workflow_executor import SubWorkflowRequestMessage, SubWorkflowResponseMessage, WorkflowExecutor
@@ -592,6 +592,7 @@ __all__ = [
     "WorkflowExecutor",
     "WorkflowMessage",
     "WorkflowRunResult",
+    "WorkflowInvocationKwargs",
     "WorkflowRunState",
     "WorkflowRunnerException",
     "WorkflowValidationError",

python/packages/core/agent_framework/_workflows/_agent.py

@@ -45,7 +45,7 @@
     from typing_extensions import TypedDict  # pragma: no cover
 
 if TYPE_CHECKING:
-    from ._workflow import Workflow
+    from ._workflow import Workflow, WorkflowInvocationKwargs
 
 logger = logging.getLogger(__name__)
 
@@ -152,28 +152,34 @@ def run(
         self,
         messages: AgentRunInputs | None = None,
         *,
-        stream: Literal[False] = ...,
+        stream: Literal[True],
         session: AgentSession | None = None,
         checkpoint_id: str | None = None,
         checkpoint_storage: CheckpointStorage | None = None,
         tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None,
-        function_invocation_kwargs: Mapping[str, Any] | None = None,
-        client_kwargs: Mapping[str, Any] | None = None,
-    ) -> Awaitable[AgentResponse[Any]]: ...
+        function_invocation_kwargs: WorkflowInvocationKwargs
+        | Mapping[str, Mapping[str, Any]]
+        | Mapping[str, Any]
+        | None = None,
+        client_kwargs: WorkflowInvocationKwargs | Mapping[str, Mapping[str, Any]] | Mapping[str, Any] | None = None,
+    ) -> ResponseStream[AgentResponseUpdate, AgentResponse]: ...
 
     @overload
-    def run(
+    async def run(
         self,
         messages: AgentRunInputs | None = None,
         *,
-        stream: Literal[True],
+        stream: Literal[False] = ...,
         session: AgentSession | None = None,
         checkpoint_id: str | None = None,
         checkpoint_storage: CheckpointStorage | None = None,
         tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None,
-        function_invocation_kwargs: Mapping[str, Any] | None = None,
-        client_kwargs: Mapping[str, Any] | None = None,
-    ) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
+        function_invocation_kwargs: WorkflowInvocationKwargs
+        | Mapping[str, Mapping[str, Any]]
+        | Mapping[str, Any]
+        | None = None,
+        client_kwargs: WorkflowInvocationKwargs | Mapping[str, Mapping[str, Any]] | Mapping[str, Any] | None = None,
+    ) -> AgentResponse: ...
 
     def run(
         self,
@@ -184,9 +190,12 @@ def run(
         checkpoint_id: str | None = None,
         checkpoint_storage: CheckpointStorage | None = None,
         tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None,
-        function_invocation_kwargs: Mapping[str, Any] | None = None,
-        client_kwargs: Mapping[str, Any] | None = None,
-    ) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]] | Awaitable[AgentResponse[Any]]:
+        function_invocation_kwargs: WorkflowInvocationKwargs
+        | Mapping[str, Mapping[str, Any]]
+        | Mapping[str, Any]
+        | None = None,
+        client_kwargs: WorkflowInvocationKwargs | Mapping[str, Mapping[str, Any]] | Mapping[str, Any] | None = None,
+    ) -> ResponseStream[AgentResponseUpdate, AgentResponse] | Awaitable[AgentResponse]:
         """Get a response from the workflow agent.
 
         Args:
@@ -254,8 +263,11 @@ async def _run_impl(
         checkpoint_id: str | None = None,
         checkpoint_storage: CheckpointStorage | None = None,
         tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None,
-        function_invocation_kwargs: Mapping[str, Mapping[str, Any]] | Mapping[str, Any] | None = None,
-        client_kwargs: Mapping[str, Mapping[str, Any]] | Mapping[str, Any] | None = None,
+        function_invocation_kwargs: WorkflowInvocationKwargs
+        | Mapping[str, Mapping[str, Any]]
+        | Mapping[str, Any]
+        | None = None,
+        client_kwargs: WorkflowInvocationKwargs | Mapping[str, Mapping[str, Any]] | Mapping[str, Any] | None = None,
     ) -> AgentResponse:
         """Internal implementation of non-streaming execution.
 
@@ -337,8 +349,11 @@ async def _run_stream_impl(
         checkpoint_id: str | None = None,
         checkpoint_storage: CheckpointStorage | None = None,
         tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None,
-        function_invocation_kwargs: Mapping[str, Mapping[str, Any]] | Mapping[str, Any] | None = None,
-        client_kwargs: Mapping[str, Mapping[str, Any]] | Mapping[str, Any] | None = None,
+        function_invocation_kwargs: WorkflowInvocationKwargs
+        | Mapping[str, Mapping[str, Any]]
+        | Mapping[str, Any]
+        | None = None,
+        client_kwargs: WorkflowInvocationKwargs | Mapping[str, Mapping[str, Any]] | Mapping[str, Any] | None = None,
     ) -> AsyncIterable[AgentResponseUpdate]:
         """Internal implementation of streaming execution.
 
@@ -419,8 +434,11 @@ async def _run_core(
         checkpoint_storage: CheckpointStorage | None,
         streaming: bool,
         tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None,
-        function_invocation_kwargs: Mapping[str, Mapping[str, Any]] | Mapping[str, Any] | None = None,
-        client_kwargs: Mapping[str, Mapping[str, Any]] | Mapping[str, Any] | None = None,
+        function_invocation_kwargs: WorkflowInvocationKwargs
+        | Mapping[str, Mapping[str, Any]]
+        | Mapping[str, Any]
+        | None = None,
+        client_kwargs: WorkflowInvocationKwargs | Mapping[str, Mapping[str, Any]] | Mapping[str, Any] | None = None,
     ) -> AsyncIterable[WorkflowEvent]:
         """Core implementation that yields workflow events for both streaming and non-streaming modes.
 
@@ -470,8 +488,7 @@ async def _run_core(
             # NOTE: It is possible that some pending requests are not fulfilled,
             # and we will let the workflow to handle this -- the agent does not
             # have an opinion on this.
-            pending_requests = await self.workflow._runner_context.get_pending_request_info_events()  # pyright: ignore[reportPrivateUsage]
-            function_responses = self._extract_function_responses(input_messages, pending_requests)
+            function_responses = self._extract_function_responses(input_messages)
             if streaming:
                 async for event in self.workflow.run(
                     responses=function_responses,
@@ -748,51 +765,22 @@ def _process_request_info_event(
             arguments=args,
         )
 
-    def _extract_function_responses(
-        self,
-        input_messages: Sequence[Message],
-        pending_requests: Mapping[str, WorkflowEvent[Any]] | None = None,
-    ) -> dict[str, Any]:
+    def _extract_function_responses(self, input_messages: Sequence[Message]) -> dict[str, Any]:
         """Extract function responses from input messages.
 
         The responses are for pending requests that the workflow is waiting on, and
         will be passed to the workflow. The pending requests are processed to either
         `function_approval_request` or `function_call` content by `_process_request_info_event`.
         """
-        pending_requests = pending_requests or {}
         function_responses: dict[str, Any] = {}
         for message in input_messages:
             for content in message.contents:
                 if content.type == "function_approval_response":
-                    request_id = content.id
-                    if request_id is None:
-                        raise AgentInvalidResponseException("Function approval response is missing its request ID.")
+                    request_id: str = content.id  # type: ignore[assignment]
                     function_responses[request_id] = content
                 elif content.type == "function_result":
-                    request_id = content.call_id
-                    if request_id is None:
-                        raise AgentInvalidResponseException("Function result is missing its call ID.")
-                    response_request_id = request_id
-                    pending_request = pending_requests.get(response_request_id)
-                    if pending_request is None:
-                        matching_requests = [
-                            (pending_id, pending_event)
-                            for pending_id, pending_event in pending_requests.items()
-                            if isinstance(pending_event.data, Content)
-                            and pending_event.data.type == "function_call"
-                            and pending_event.data.call_id == request_id
-                        ]
-                        if len(matching_requests) == 1:
-                            response_request_id, pending_request = matching_requests[0]
-                    response_data = (
-                        content
-                        if pending_request is not None
-                        and pending_request.response_type is Content
-                        and isinstance(pending_request.data, Content)
-                        and pending_request.data.type == "function_call"
-                        else content.result
-                    )
-                    function_responses[response_request_id] = response_data
+                    response_data = content.result if hasattr(content, "result") else str(content)
+                    function_responses[content.call_id] = response_data  # type: ignore
                 else:
                     raise AgentInvalidResponseException(
                         "Unexpected content type while awaiting request info responses."

python/packages/core/agent_framework/_workflows/_agent_executor.py

@@ -614,22 +614,26 @@ def _resolve_executor_kwargs(self, resolved: dict[str, Any] | None) -> dict[str,
         """
         if not isinstance(resolved, dict):
             return None
-        # Use explicit key-presence checks so that an empty per-executor dict is
-        # honoured (e.g. to clear kwargs) instead of falling through to global.
-        if self.id in resolved:
-            executor_kwargs = resolved[self.id]
-        elif GLOBAL_KWARGS_KEY in resolved:
-            executor_kwargs = resolved[GLOBAL_KWARGS_KEY]
-        else:
+        global_kwargs: Any = resolved.get(GLOBAL_KWARGS_KEY)
+        executor_kwargs: Any = resolved.get(self.id)
+        if global_kwargs is None and executor_kwargs is None:
             return None
 
-        if not isinstance(executor_kwargs, dict):
+        if global_kwargs is not None and not isinstance(global_kwargs, dict):
             logger.warning(
-                "Executor %s expected a dict for its kwargs, but got %s. Ignoring.",
+                "Executor %s expected a dict for global kwargs, but got %s. Ignoring.",
                 self.id,
-                type(executor_kwargs),  # type: ignore
+                cast(type[Any], type(global_kwargs)),
             )
+            return None
 
+        if executor_kwargs is not None and not isinstance(executor_kwargs, dict):
+            logger.warning(
+                "Executor %s expected a dict for its kwargs, but got %s. Ignoring.",
+                self.id,
+                cast(type[Any], type(executor_kwargs)),
+            )
             return None
 
-        return executor_kwargs  # type: ignore
+        # Specific values override global values for the same function argument.
+        return {**(global_kwargs or {}), **(executor_kwargs or {})}

python/packages/core/agent_framework/_workflows/_const.py

@@ -14,6 +14,10 @@
 # to pass kwargs from workflow.run() through to agent.run() and @tool functions.
 WORKFLOW_RUN_KWARGS_KEY = "_workflow_run_kwargs"
 
+# State keys used to preserve caller-provided kwargs for nested workflow routing.
+RAW_FUNCTION_INVOCATION_KWARGS_KEY = "_raw_function_invocation_kwargs"
+RAW_CLIENT_KWARGS_KEY = "_raw_client_kwargs"
+
 # Sentinel key used in resolved invocation kwargs dicts to denote global kwargs
 # that apply to all executors (as opposed to per-executor keyed entries).
 GLOBAL_KWARGS_KEY = "__global__"