diff --git a/python/packages/core/agent_framework/__init__.py b/python/packages/core/agent_framework/__init__.py index c80b5025eb..691c35e911 100644 --- a/python/packages/core/agent_framework/__init__.py +++ b/python/packages/core/agent_framework/__init__.py @@ -374,7 +374,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": ( @@ -646,6 +646,7 @@ "WorkflowEventType", "WorkflowException", "WorkflowExecutor", + "WorkflowInvocationKwargs", "WorkflowMessage", "WorkflowRunResult", "WorkflowRunState", diff --git a/python/packages/core/agent_framework/__init__.pyi b/python/packages/core/agent_framework/__init__.pyi index a18848f56f..b303e0d98e 100644 --- a/python/packages/core/agent_framework/__init__.pyi +++ b/python/packages/core/agent_framework/__init__.pyi @@ -328,7 +328,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 @@ -603,6 +603,7 @@ __all__ = [ "WorkflowExecutor", "WorkflowMessage", "WorkflowRunResult", + "WorkflowInvocationKwargs", "WorkflowRunState", "WorkflowRunnerException", "WorkflowValidationError", diff --git a/python/packages/core/agent_framework/_workflows/_agent.py b/python/packages/core/agent_framework/_workflows/_agent.py index 0df47d2b34..8198ad297e 100644 --- a/python/packages/core/agent_framework/_workflows/_agent.py +++ b/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__) @@ -157,9 +157,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, - ) -> 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( @@ -171,9 +174,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]]: ... + 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]: ... 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. diff --git a/python/packages/core/agent_framework/_workflows/_agent_executor.py b/python/packages/core/agent_framework/_workflows/_agent_executor.py index 512706e36f..1c3a7203f4 100644 --- a/python/packages/core/agent_framework/_workflows/_agent_executor.py +++ b/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 {})} diff --git a/python/packages/core/agent_framework/_workflows/_const.py b/python/packages/core/agent_framework/_workflows/_const.py index 27b9c24961..f01b45d840 100644 --- a/python/packages/core/agent_framework/_workflows/_const.py +++ b/python/packages/core/agent_framework/_workflows/_const.py @@ -17,6 +17,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__" diff --git a/python/packages/core/agent_framework/_workflows/_workflow.py b/python/packages/core/agent_framework/_workflows/_workflow.py index ab89df1e5b..6e9f256b9b 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow.py +++ b/python/packages/core/agent_framework/_workflows/_workflow.py @@ -21,7 +21,14 @@ from ..exceptions import WorkflowException from ..observability import OtelAttr, capture_exception, create_workflow_span from ._checkpoint import CheckpointStorage -from ._const import DEFAULT_MAX_ITERATIONS, GLOBAL_KWARGS_KEY, INTERNAL_SOURCE_ID, WORKFLOW_RUN_KWARGS_KEY +from ._const import ( + DEFAULT_MAX_ITERATIONS, + GLOBAL_KWARGS_KEY, + INTERNAL_SOURCE_ID, + RAW_CLIENT_KWARGS_KEY, + RAW_FUNCTION_INVOCATION_KWARGS_KEY, + WORKFLOW_RUN_KWARGS_KEY, +) from ._edge import ( EdgeGroup, FanOutEdgeGroup, @@ -206,6 +213,20 @@ def classify(self, executor_id: str) -> Literal["output", "intermediate"] | None return None +@dataclass(frozen=True) +class WorkflowInvocationKwargs: + """Explicit global and executor-specific kwargs for a workflow run. + + Use this wrapper when shared kwargs should be combined with executor-specific + overrides. Plain mappings retain their existing global or per-executor behavior. + """ + + global_kwargs: Mapping[str, Any] = field(default_factory=lambda: dict[str, Any]()) + executor_kwargs: Mapping[str, Mapping[str, Any]] = field( + default_factory=lambda: dict[str, Mapping[str, Any]]() + ) + + class Workflow(DictConvertible): """A graph-based execution engine that orchestrates connected executors. @@ -482,8 +503,11 @@ async def _run_workflow_with_tracing( is_continuation: bool = False, streaming: bool = False, 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]: """Private method to run workflow with proper tracing. @@ -558,10 +582,18 @@ async def _run_workflow_with_tracing( combined_kwargs["function_invocation_kwargs"] = self._resolve_invocation_kwargs( function_invocation_kwargs, "function_invocation_kwargs" ) + if isinstance(function_invocation_kwargs, WorkflowInvocationKwargs) or any( + isinstance(value, Mapping) for value in function_invocation_kwargs.values() + ): + combined_kwargs[RAW_FUNCTION_INVOCATION_KWARGS_KEY] = function_invocation_kwargs if client_kwargs is not None: combined_kwargs["client_kwargs"] = self._resolve_invocation_kwargs( client_kwargs, "client_kwargs" ) + if isinstance(client_kwargs, WorkflowInvocationKwargs) or any( + isinstance(value, Mapping) for value in client_kwargs.values() + ): + combined_kwargs[RAW_CLIENT_KWARGS_KEY] = client_kwargs self._runner.state.set(WORKFLOW_RUN_KWARGS_KEY, combined_kwargs) elif not is_continuation: self._runner.state.set(WORKFLOW_RUN_KWARGS_KEY, {}) @@ -691,8 +723,14 @@ 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, + 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[WorkflowEvent, WorkflowRunResult]: ... @overload @@ -706,8 +744,8 @@ def run( checkpoint_storage: CheckpointStorage | None = None, include_status_events: bool = False, 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, + function_invocation_kwargs: WorkflowInvocationKwargs | Mapping[str, Any] | None = None, + client_kwargs: WorkflowInvocationKwargs | Mapping[str, Any] | None = None, ) -> Awaitable[WorkflowRunResult]: ... def run( @@ -720,8 +758,11 @@ def run( checkpoint_storage: CheckpointStorage | None = None, include_status_events: bool = False, 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, ) -> ResponseStream[WorkflowEvent, WorkflowRunResult] | Awaitable[WorkflowRunResult]: """Run the workflow, optionally streaming events. @@ -746,10 +787,14 @@ def run( tools: Runtime tools available to agent executors. function_invocation_kwargs: Keyword arguments forwarded to tool invocations in subagents. Either a mapping for agent name or agent executor id to kwargs, - or a flat mapping of kwargs for all tool invocations. + a flat mapping of kwargs for all tool invocations, or a + ``WorkflowInvocationKwargs`` instance to combine global and executor-specific + kwargs. client_kwargs: Keyword arguments forwarded to chat client calls in subagents. Either a mapping for agent name or agent executor id to kwargs, - or a flat mapping of kwargs for all chat client calls. + a flat mapping of kwargs for all chat client calls, or a + ``WorkflowInvocationKwargs`` instance to combine global and executor-specific + kwargs. Returns: When stream=True: A ResponseStream[WorkflowEvent, WorkflowRunResult] for @@ -791,7 +836,7 @@ def run( checkpoint_id=checkpoint_id, checkpoint_storage=checkpoint_storage, streaming=stream, - tools=runtime_tools, + tools=tools, function_invocation_kwargs=function_invocation_kwargs, client_kwargs=client_kwargs, ), @@ -812,8 +857,11 @@ async def _run_core( checkpoint_storage: CheckpointStorage | None = None, streaming: bool = False, 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]: """Single core execution path for both streaming and non-streaming modes. @@ -1079,7 +1127,7 @@ def _get_executor_by_id(self, executor_id: str) -> Executor: def _resolve_invocation_kwargs( self, - kwargs: Mapping[str, Any], + kwargs: WorkflowInvocationKwargs | Mapping[str, Any], param_name: str, ) -> dict[str, Any]: """Resolve invocation kwargs into a normalized per-executor or global format. @@ -1087,17 +1135,24 @@ def _resolve_invocation_kwargs( Detects whether the provided kwargs dict uses per-executor targeting by checking if any top-level key matches a known executor ID in the workflow. If at least one key matches, all entries are treated as per-executor. Otherwise the dict is treated - as global kwargs that apply to every executor. + as global kwargs that apply to every executor. The ``"__global__"`` key can be used + explicitly to combine global kwargs with per-executor overrides. Args: kwargs: The raw invocation kwargs from the caller. param_name: The parameter name (for logging), e.g. ``"function_invocation_kwargs"``. Returns: - A dict with either: - - ``{"__global__": }`` for global kwargs, or - - The original dict unchanged for per-executor kwargs. + A dict containing normalized global or per-executor mappings. """ + if isinstance(kwargs, WorkflowInvocationKwargs): + resolved = {GLOBAL_KWARGS_KEY: dict(kwargs.global_kwargs)} + resolved.update({ + executor_id: dict(executor_kwargs) for executor_id, executor_kwargs in kwargs.executor_kwargs.items() + }) + logger.info("Explicit global %s provided with executor-specific overrides.", param_name) + return resolved + executor_ids = set(self.executors.keys()) matched_ids = kwargs.keys() & executor_ids if matched_ids: @@ -1221,29 +1276,10 @@ async def cancel_pending_requests( 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, Any] | None = None, + client_kwargs: WorkflowInvocationKwargs | Mapping[str, Any] | None = None, ) -> WorkflowRunResult: - """Cancel pending external requests and release their owning executor state. - - Cancellation follows requests through nested workflows and clears any executor-owned - correlation without synthesizing a response. If cancellation drains an executor's pending - set after sibling responses were already accepted, the executor resumes through its normal - continuation path. Unknown or already-handled request IDs are ignored. - - Args: - request_ids: Request identifiers to cancel. - - Keyword Args: - checkpoint_id: Checkpoint to restore before applying cancellation. - checkpoint_storage: Runtime checkpoint storage for the cancellation continuation. - tools: Request-scoped tools available while cancellation resumes executors. - function_invocation_kwargs: Keyword arguments forwarded to resumed tool invocations. - client_kwargs: Keyword arguments forwarded to resumed chat client calls. - - Returns: - Events produced while applying cancellation and any resulting continuation. - """ + """Cancel pending external requests and continue the workflow.""" selected_ids = set(request_ids) if not all(isinstance(request_id, str) and request_id for request_id in selected_ids): raise ValueError("Pending workflow request IDs must be non-empty strings.") @@ -1260,10 +1296,7 @@ async def apply_cancellations() -> None: state=self._runner.state, runner_context=self._runner.context, ) - await executor._cancel_pending_request( # pyright: ignore[reportPrivateUsage] - request_id, - context, - ) + await executor._cancel_pending_request(request_id, context) # pyright: ignore[reportPrivateUsage] if checkpoint_storage is not None: self._runner.context.set_runtime_checkpoint_storage(checkpoint_storage) diff --git a/python/packages/core/agent_framework/_workflows/_workflow_executor.py b/python/packages/core/agent_framework/_workflows/_workflow_executor.py index fce77a540c..749283d2d7 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow_executor.py +++ b/python/packages/core/agent_framework/_workflows/_workflow_executor.py @@ -3,13 +3,19 @@ import logging import sys import types +from collections.abc import Mapping from dataclasses import dataclass -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, cast if TYPE_CHECKING: from ._workflow import Workflow -from ._const import GLOBAL_KWARGS_KEY, WORKFLOW_RUN_KWARGS_KEY +from ._const import ( + GLOBAL_KWARGS_KEY, + RAW_CLIENT_KWARGS_KEY, + RAW_FUNCTION_INVOCATION_KWARGS_KEY, + WORKFLOW_RUN_KWARGS_KEY, +) from ._edge_runner import gather_cancelling_siblings_on_error from ._events import ( WorkflowEvent, @@ -20,7 +26,7 @@ from ._request_info_mixin import response_handler from ._runner_context import WorkflowMessage from ._typing_utils import is_instance_of -from ._workflow import WorkflowRunResult +from ._workflow import WorkflowInvocationKwargs, WorkflowRunResult from ._workflow_context import WorkflowContext if sys.version_info >= (3, 12): @@ -375,29 +381,36 @@ async def process_workflow(self, input_data: object, ctx: WorkflowContext[Any, A # Get kwargs from parent workflow's State to propagate to subworkflow parent_kwargs: dict[str, Any] = ctx.get_state(WORKFLOW_RUN_KWARGS_KEY, {}) - # Extract invocation kwargs recognised by Workflow.run() - # The state stores resolved format (with __global__ wrapper for global kwargs). - # Unwrap __global__ before passing to the subworkflow so it gets re-resolved - # against the subworkflow's own executor IDs. - fi_kwargs: dict[str, Any] | None = None - ci_kwargs: dict[str, Any] | None = None - tools = ctx.get_runtime_tools() + # Use the caller's raw kwargs so legacy per-executor mappings are resolved + # against the child workflow's executor IDs rather than the parent's. + fi_kwargs: WorkflowInvocationKwargs | Mapping[str, Any] | None = None + ci_kwargs: WorkflowInvocationKwargs | Mapping[str, Any] | None = None for key in ("function_invocation_kwargs", "client_kwargs"): - resolved = parent_kwargs.get(key) - if isinstance(resolved, dict): - # Unwrap global sentinel; pass per-executor dicts as-is - unwrapped: dict[str, Any] = resolved.get(GLOBAL_KWARGS_KEY, resolved) # type: ignore + raw_key = ( + RAW_FUNCTION_INVOCATION_KWARGS_KEY if key == "function_invocation_kwargs" else RAW_CLIENT_KWARGS_KEY + ) + raw_value = parent_kwargs.get(raw_key) + if raw_value is not None: + resolved = cast(WorkflowInvocationKwargs | Mapping[str, Any], raw_value) + else: + normalized: Any = parent_kwargs.get(key) + if isinstance(normalized, dict): + normalized_dict = cast(dict[str, Any], normalized) + if len(normalized_dict) == 1 and GLOBAL_KWARGS_KEY in normalized_dict: + normalized = normalized_dict[GLOBAL_KWARGS_KEY] + resolved = cast(WorkflowInvocationKwargs | Mapping[str, Any] | None, normalized) + if resolved is not None: if key == "function_invocation_kwargs": - fi_kwargs = unwrapped # type: ignore + fi_kwargs = resolved else: - ci_kwargs = unwrapped # type: ignore + ci_kwargs = resolved # Run the sub-workflow and collect all events, passing parent kwargs result = await self.workflow.run( input_data, - tools=tools, - function_invocation_kwargs=fi_kwargs, # type: ignore - client_kwargs=ci_kwargs, # type: ignore + tools=ctx.get_runtime_tools(), + function_invocation_kwargs=fi_kwargs, + client_kwargs=ci_kwargs, ) logger.debug(f"WorkflowExecutor {self.id} sub-workflow {self.workflow.id} completed with {len(result)} events") @@ -621,5 +634,5 @@ async def _handle_response( # Forward the response to the sub-workflow, which resumes and validates it against its own # pending requests, then process whatever the sub-workflow produces. - result = await self.workflow.run(responses={request_id: response}, tools=ctx.get_runtime_tools()) + result = await self.workflow.run(responses={request_id: response}) await self._process_workflow_result(result, ctx) diff --git a/python/packages/core/tests/workflow/test_agent_executor.py b/python/packages/core/tests/workflow/test_agent_executor.py index 2cc2ed2ce6..9d7e4978a0 100644 --- a/python/packages/core/tests/workflow/test_agent_executor.py +++ b/python/packages/core/tests/workflow/test_agent_executor.py @@ -1,7 +1,6 @@ # Copyright (c) Microsoft. All rights reserved. import pickle - from collections.abc import AsyncIterable, Awaitable from typing import Any, Literal, overload @@ -338,9 +337,6 @@ class _NonCopyableRaw: def __deepcopy__(self, memo: dict) -> Any: raise TypeError("Cannot deepcopy this object") - def __reduce__(self) -> Any: - raise TypeError("Cannot pickle this object") - class _AgentWithRawRepr(BaseAgent): """Agent that returns responses with a non-copyable raw_representation.""" @@ -392,20 +388,6 @@ async def test_agent_executor_workflow_with_non_copyable_raw_representation() -> assert agent_responses[0].raw_representation is raw -def test_serialization_mixin_omits_non_pickleable_raw_representation() -> None: - """Pickling framework objects should not include runtime-only raw representations.""" - raw = _NonCopyableRaw() - response = AgentResponse( - messages=[Message("assistant", [Content.from_text(text="reply", raw_representation=raw)])], - raw_representation=raw, - ) - - restored = pickle.loads(pickle.dumps(response)) - - assert restored.raw_representation is None - assert restored.messages[0].contents[0].raw_representation is None - - # --------------------------------------------------------------------------- # Context mode tests # --------------------------------------------------------------------------- @@ -641,15 +623,15 @@ async def test_resolve_executor_kwargs_returns_none_for_none_input() -> None: assert result is None -async def test_resolve_executor_kwargs_prefers_executor_id_over_global() -> None: - """_resolve_executor_kwargs prefers executor-specific entry over __global__.""" +async def test_resolve_executor_kwargs_merges_executor_id_over_global() -> None: + """_resolve_executor_kwargs merges executor-specific entries over __global__.""" agent = _CountingAgent(id="a", name="A") executor = AgentExecutor(agent, id="exec_a") # Dict has both a per-executor entry and a global entry resolved = {"exec_a": {"specific": True}, GLOBAL_KWARGS_KEY: {"global": True}} result = executor._resolve_executor_kwargs(resolved) # pyright: ignore[reportPrivateUsage] - assert result == {"specific": True} + assert result == {"global": True, "specific": True} async def test_prepare_agent_run_args_extracts_function_invocation_kwargs() -> None: @@ -708,16 +690,15 @@ async def test_prepare_agent_run_args_per_executor_no_match() -> None: assert fi_kwargs is None -async def test_resolve_executor_kwargs_empty_per_executor_does_not_fallback_to_global() -> None: - """An explicit empty per-executor dict should not fall through to global kwargs.""" +async def test_resolve_executor_kwargs_empty_per_executor_keeps_global_kwargs() -> None: + """An explicit empty per-executor dict keeps the global kwargs.""" agent = _CountingAgent(id="a", name="A") executor = AgentExecutor(agent, id="exec_a") - # Per-executor entry for exec_a is empty, but global has values. - # The empty dict should be honoured (no fallback to global). + # Per-executor entry for exec_a is empty, so only global values apply. resolved = {"exec_a": {}, GLOBAL_KWARGS_KEY: {"global_key": "global_val"}} # type: ignore[var-annotated] result = executor._resolve_executor_kwargs(resolved) # pyright: ignore[reportPrivateUsage] - assert result == {} + assert result == {"global_key": "global_val"} # region Tool approval emission diff --git a/python/packages/core/tests/workflow/test_workflow_kwargs.py b/python/packages/core/tests/workflow/test_workflow_kwargs.py index c21a25a5d7..d91d224c54 100644 --- a/python/packages/core/tests/workflow/test_workflow_kwargs.py +++ b/python/packages/core/tests/workflow/test_workflow_kwargs.py @@ -17,6 +17,7 @@ FunctionTool, Message, ResponseStream, + WorkflowInvocationKwargs, WorkflowRunState, ) from agent_framework._workflows._const import WORKFLOW_RUN_KWARGS_KEY @@ -1070,6 +1071,35 @@ async def test_nested_subworkflow_kwargs_propagation() -> None: ) +async def test_mixed_kwargs_route_through_subworkflow() -> None: + """Mixed kwargs preserve global values and child executor-specific routing.""" + from agent_framework._workflows._workflow_executor import WorkflowExecutor + + inner_agent1 = _KwargsCapturingAgent(name="inner_agent1") + inner_agent2 = _KwargsCapturingAgent(name="inner_agent2") + inner_workflow = SequentialBuilder(participants=[inner_agent1, inner_agent2]).build() + subworkflow_executor = WorkflowExecutor(workflow=inner_workflow, id="subworkflow") + outer_workflow = SequentialBuilder(participants=[subworkflow_executor]).build() + + fi_kwargs = WorkflowInvocationKwargs( + global_kwargs={"shared": "value", "overridden": "global"}, + executor_kwargs={"inner_agent2": {"overridden": "inner_agent2"}}, + ) + + async for event in outer_workflow.run("test", stream=True, function_invocation_kwargs=fi_kwargs): + if event.type == "status" and event.state == WorkflowRunState.IDLE: + break + + assert inner_agent1.captured_kwargs[0].get("function_invocation_kwargs") == { + "shared": "value", + "overridden": "global", + } + assert inner_agent2.captured_kwargs[0].get("function_invocation_kwargs") == { + "shared": "value", + "overridden": "inner_agent2", + } + + # endregion @@ -1151,6 +1181,35 @@ async def test_per_executor_function_invocation_kwargs_routes_to_correct_agent() assert agent2.captured_kwargs[0].get("function_invocation_kwargs") == {"tool_param": "value_for_agent2"} +async def test_global_and_per_executor_function_invocation_kwargs_are_merged() -> None: + """Global function kwargs are merged with executor-specific overrides.""" + agent1 = _KwargsCapturingAgent(name="agent1") + agent2 = _KwargsCapturingAgent(name="agent2") + workflow = SequentialBuilder(participants=[agent1, agent2]).build() + + fi_kwargs = WorkflowInvocationKwargs( + global_kwargs={"shared": "value", "overridden": "global"}, + executor_kwargs={ + "agent1": {"overridden": "agent1"}, + "agent2": {"agent_only": True}, + }, + ) + + async for event in workflow.run("test", stream=True, function_invocation_kwargs=fi_kwargs): + if event.type == "status" and event.state == WorkflowRunState.IDLE: + break + + assert agent1.captured_kwargs[0].get("function_invocation_kwargs") == { + "shared": "value", + "overridden": "agent1", + } + assert agent2.captured_kwargs[0].get("function_invocation_kwargs") == { + "shared": "value", + "overridden": "global", + "agent_only": True, + } + + async def test_per_executor_kwargs_unmatched_agent_gets_none() -> None: """An agent not targeted in per-executor kwargs should receive None for that kwarg.""" agent1 = _KwargsCapturingAgent(name="agent1") @@ -1221,6 +1280,35 @@ async def test_per_executor_client_kwargs_routes_correctly() -> None: assert agent2.captured_kwargs[0].get("client_kwargs") == {"temperature": 0.9} +async def test_global_and_per_executor_client_kwargs_are_merged() -> None: + """Global client kwargs are merged with executor-specific overrides.""" + agent1 = _KwargsCapturingAgent(name="agent1") + agent2 = _KwargsCapturingAgent(name="agent2") + workflow = SequentialBuilder(participants=[agent1, agent2]).build() + + ci_kwargs = WorkflowInvocationKwargs( + global_kwargs={"shared": "value", "overridden": "global"}, + executor_kwargs={ + "agent1": {"overridden": "agent1"}, + "agent2": {"agent_only": True}, + }, + ) + + async for event in workflow.run("test", stream=True, client_kwargs=ci_kwargs): + if event.type == "status" and event.state == WorkflowRunState.IDLE: + break + + assert agent1.captured_kwargs[0].get("client_kwargs") == { + "shared": "value", + "overridden": "agent1", + } + assert agent2.captured_kwargs[0].get("client_kwargs") == { + "shared": "value", + "overridden": "global", + "agent_only": True, + } + + async def test_resolve_invocation_kwargs_logs_per_executor(caplog: "LogCaptureFixture") -> None: """Workflow._resolve_invocation_kwargs logs info when per-executor format is detected.""" import logging