diff --git a/python/packages/core/agent_framework/_workflows/_agent_executor.py b/python/packages/core/agent_framework/_workflows/_agent_executor.py index 1c3a7203f49..cd69504c71d 100644 --- a/python/packages/core/agent_framework/_workflows/_agent_executor.py +++ b/python/packages/core/agent_framework/_workflows/_agent_executor.py @@ -14,8 +14,8 @@ from .._agents import SupportsAgentRun from .._sessions import AgentSession from .._types import AgentResponse, AgentResponseUpdate, Message, ResponseStream -from ._agent_utils import resolve_agent_id -from ._const import GLOBAL_KWARGS_KEY, INTERNAL_SOURCE_ID, WORKFLOW_RUN_KWARGS_KEY +from ._agent_utils import prepare_agent_run_args, resolve_agent_id, resolve_executor_kwargs +from ._const import INTERNAL_SOURCE_ID, WORKFLOW_RUN_KWARGS_KEY from ._executor import Executor, handler from ._message_utils import normalize_messages_input from ._request_info_mixin import response_handler @@ -594,12 +594,7 @@ def _prepare_agent_run_args( Returns: A 2-tuple of (function_invocation_kwargs, client_kwargs). """ - fi_resolved = raw_run_kwargs.get("function_invocation_kwargs") - ci_resolved = raw_run_kwargs.get("client_kwargs") - function_invocation_kwargs = self._resolve_executor_kwargs(fi_resolved) - client_kwargs = self._resolve_executor_kwargs(ci_resolved) - - return function_invocation_kwargs, client_kwargs + return prepare_agent_run_args(self.id, raw_run_kwargs) def _resolve_executor_kwargs(self, resolved: dict[str, Any] | None) -> dict[str, Any] | None: """Extract this executor's kwargs from a resolved invocation kwargs dict. @@ -612,28 +607,4 @@ def _resolve_executor_kwargs(self, resolved: dict[str, Any] | None) -> dict[str, Returns: The kwargs for this executor, or ``None`` if not applicable. """ - if not isinstance(resolved, dict): - return None - 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 global_kwargs is not None and not isinstance(global_kwargs, dict): - logger.warning( - "Executor %s expected a dict for global kwargs, but got %s. Ignoring.", - self.id, - 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 - - # Specific values override global values for the same function argument. - return {**(global_kwargs or {}), **(executor_kwargs or {})} + return resolve_executor_kwargs(self.id, resolved) diff --git a/python/packages/core/agent_framework/_workflows/_agent_utils.py b/python/packages/core/agent_framework/_workflows/_agent_utils.py index f70d524ceeb..08034123328 100644 --- a/python/packages/core/agent_framework/_workflows/_agent_utils.py +++ b/python/packages/core/agent_framework/_workflows/_agent_utils.py @@ -1,6 +1,12 @@ # Copyright (c) Microsoft. All rights reserved. +import logging +from typing import Any, cast + from .._agents import SupportsAgentRun +from ._const import GLOBAL_KWARGS_KEY + +logger = logging.getLogger(__name__) def resolve_agent_id(agent: SupportsAgentRun) -> str: @@ -15,3 +21,67 @@ def resolve_agent_id(agent: SupportsAgentRun) -> str: The resolved unique identifier for the agent. """ return agent.name if agent.name else agent.id + + +def resolve_executor_kwargs(executor_id: str, resolved: dict[str, Any] | None) -> dict[str, Any] | None: + """Extract one executor's kwargs from a resolved invocation kwargs dict. + + Args: + executor_id: The id of the executor whose kwargs are wanted. + resolved: The resolved dict produced by ``Workflow._resolve_invocation_kwargs``, + containing either a ``__global__`` key (global kwargs) or executor-ID keys + (per-executor kwargs). May also be ``None``. + + Returns: + The kwargs for that executor, or ``None`` if not applicable. + """ + if not isinstance(resolved, dict): + return None + global_kwargs: Any = resolved.get(GLOBAL_KWARGS_KEY) + executor_kwargs: Any = resolved.get(executor_id) + if global_kwargs is None and executor_kwargs is None: + return None + + if global_kwargs is not None and not isinstance(global_kwargs, dict): + logger.warning( + "Executor %s expected a dict for global kwargs, but got %s. Ignoring.", + executor_id, + 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.", + executor_id, + cast(type[Any], type(executor_kwargs)), + ) + return None + + # Specific values override global values for the same function argument. + return {**(global_kwargs or {}), **(executor_kwargs or {})} + + +def prepare_agent_run_args( + executor_id: str, + raw_run_kwargs: dict[str, Any], +) -> tuple[dict[str, Any] | None, dict[str, Any] | None]: + """Prepare function_invocation_kwargs and client_kwargs for agent.run(). + + Extracts ``function_invocation_kwargs`` and ``client_kwargs`` from the workflow state + dict, resolving per-executor entries using ``executor_id``. The ``__global__`` sentinel + key (set by ``Workflow._resolve_invocation_kwargs``) denotes global kwargs that apply to + all executors. Per-executor dicts use executor IDs as keys; only the entry for + ``executor_id`` is extracted. + + Args: + executor_id: The id of the executor about to invoke the agent. + raw_run_kwargs: The workflow state dict stored under ``WORKFLOW_RUN_KWARGS_KEY``. + + Returns: + A 2-tuple of (function_invocation_kwargs, client_kwargs). + """ + function_invocation_kwargs = resolve_executor_kwargs(executor_id, raw_run_kwargs.get("function_invocation_kwargs")) + client_kwargs = resolve_executor_kwargs(executor_id, raw_run_kwargs.get("client_kwargs")) + + return function_invocation_kwargs, client_kwargs diff --git a/python/packages/orchestrations/agent_framework_orchestrations/_group_chat.py b/python/packages/orchestrations/agent_framework_orchestrations/_group_chat.py index bbb61edae28..879e09236cf 100644 --- a/python/packages/orchestrations/agent_framework_orchestrations/_group_chat.py +++ b/python/packages/orchestrations/agent_framework_orchestrations/_group_chat.py @@ -32,8 +32,9 @@ from agent_framework import Agent, AgentResponse, AgentResponseUpdate, AgentSession, Message, SupportsAgentRun from agent_framework._telemetry import mark_feature_used from agent_framework._workflows._agent_executor import AgentExecutor, AgentExecutorRequest, AgentExecutorResponse -from agent_framework._workflows._agent_utils import resolve_agent_id +from agent_framework._workflows._agent_utils import prepare_agent_run_args, resolve_agent_id from agent_framework._workflows._checkpoint import CheckpointStorage +from agent_framework._workflows._const import WORKFLOW_RUN_KWARGS_KEY from agent_framework._workflows._executor import Executor from agent_framework._workflows._workflow import Workflow from agent_framework._workflows._workflow_context import WorkflowContext @@ -353,7 +354,7 @@ async def _handle_messages( ): return - agent_orchestration_output = await self._invoke_agent() + agent_orchestration_output = await self._invoke_agent(cast(WorkflowContext[Any, Any], ctx)) if await self._check_agent_terminate_and_yield( agent_orchestration_output, cast(WorkflowContext[Never, AgentResponse | AgentResponseUpdate], ctx), @@ -393,7 +394,7 @@ async def _handle_response( ): return - agent_orchestration_output = await self._invoke_agent() + agent_orchestration_output = await self._invoke_agent(cast(WorkflowContext[Any, Any], ctx)) if await self._check_agent_terminate_and_yield( agent_orchestration_output, cast(WorkflowContext[Never, AgentResponse | AgentResponseUpdate], ctx), @@ -486,8 +487,16 @@ def _parse_agent_output(cls, agent_response: Any) -> AgentOrchestrationOutput: raise ValueError("Failed to parse agent orchestration output.") from last_error - async def _invoke_agent(self) -> AgentOrchestrationOutput: - """Invoke the orchestrator agent to determine the next speaker and termination.""" + async def _invoke_agent(self, ctx: WorkflowContext[Any, Any]) -> AgentOrchestrationOutput: + """Invoke the orchestrator agent to determine the next speaker and termination. + + Args: + ctx: The workflow context, read for the run kwargs stored by ``Workflow.run``. + The orchestrator agent runs outside ``AgentExecutor``, so it has to resolve + those itself or it is the only agent in the group chat that does not see them. + """ + raw_run_kwargs: dict[str, Any] = ctx.get_state(WORKFLOW_RUN_KWARGS_KEY, {}) + function_invocation_kwargs, client_kwargs = prepare_agent_run_args(self.id, raw_run_kwargs) async def _invoke_agent_helper(conversation: list[Message]) -> AgentOrchestrationOutput: # Run the agent in non-streaming mode for simplicity @@ -495,6 +504,8 @@ async def _invoke_agent_helper(conversation: list[Message]) -> AgentOrchestratio messages=conversation, session=self._session, options={"response_format": AgentOrchestrationOutput}, + function_invocation_kwargs=function_invocation_kwargs, + client_kwargs=client_kwargs, ) # Parse and validate the structured output agent_orchestration_output = self._parse_agent_output(agent_response) diff --git a/python/packages/orchestrations/tests/test_group_chat.py b/python/packages/orchestrations/tests/test_group_chat.py index 935111ccaad..92d600ba9ce 100644 --- a/python/packages/orchestrations/tests/test_group_chat.py +++ b/python/packages/orchestrations/tests/test_group_chat.py @@ -275,6 +275,78 @@ async def test_group_chat_as_agent_accepts_conversation() -> None: assert response.messages, "Expected agent conversation output" +class KwargsRecordingManagerAgent(StubManagerAgent): + """Manager agent that records the run kwargs it was invoked with.""" + + def __init__(self) -> None: + super().__init__() + self.seen_kwargs: list[dict[str, Any]] = [] + + async def run( # type: ignore[override] # ty: ignore[invalid-method-override] + self, + messages: str | Content | Message | Sequence[str | Content | Message] | None = None, + *, + session: AgentSession | None = None, + **kwargs: Any, + ) -> AgentResponse[Any]: + self.seen_kwargs.append(dict(kwargs)) + return await super().run(messages, session=session, **kwargs) + + +async def test_agent_manager_receives_workflow_run_kwargs() -> None: + """The orchestrator agent gets the same run kwargs the participants already get. + + ``workflow.run(function_invocation_kwargs=...)`` is stored in workflow state and + ``AgentExecutor`` forwards it to every participant agent. The orchestrator agent runs + outside ``AgentExecutor``, so it has to resolve the same state itself or hosts that put + request-scoped values there (a user id for tool ACL, say) silently get them on the + participants and not on the orchestrator. + """ + manager = KwargsRecordingManagerAgent() + worker = StubAgent("agent", "worker response") + + workflow = GroupChatBuilder( + participants=[worker], + orchestrator_agent=manager, + ).build() + + async for _ in workflow.run( + "coordinate task", + stream=True, + function_invocation_kwargs={"user_id": "user-123"}, + client_kwargs={"trace_id": "trace-abc"}, + ): + pass + + assert manager.seen_kwargs, "Expected the orchestrator agent to be invoked" + for call in manager.seen_kwargs: + assert call.get("function_invocation_kwargs") == {"user_id": "user-123"} + assert call.get("client_kwargs") == {"trace_id": "trace-abc"} + + +async def test_agent_manager_receives_no_run_kwargs_when_none_supplied() -> None: + """With nothing declared on the run, the orchestrator is invoked with both kwargs as None. + + This pins the shape rather than just the happy path: the resolution has to return None, + not an empty dict, so a client that distinguishes the two is not handed a stray {}. + """ + manager = KwargsRecordingManagerAgent() + worker = StubAgent("agent", "worker response") + + workflow = GroupChatBuilder( + participants=[worker], + orchestrator_agent=manager, + ).build() + + async for _ in workflow.run("coordinate task", stream=True): + pass + + assert manager.seen_kwargs, "Expected the orchestrator agent to be invoked" + for call in manager.seen_kwargs: + assert call.get("function_invocation_kwargs") is None + assert call.get("client_kwargs") is None + + async def test_agent_manager_handles_concatenated_json_output() -> None: manager = ConcatenatedJsonManagerAgent() worker = StubAgent("agent", "worker response")