Python: Flow custom kwargs to agents via Workflow SharedState (#2894)

* Flow custom kwargs to agents via SharedState * Address Copilot feedback * Improve sample typing * Fix test
2026-06-16 21:04:09 +08:00 · 2025-12-17 09:04:00 +09:00
parent 8fca71e5ad
commit 6adcac2e97
8 changed files with 921 additions and 17 deletions
@@ -149,6 +149,8 @@ to configure which agents can route to which others with a fluent, type-safe API
 | Sample | File | Concepts |
 |---|---|---|
 | Shared States | [state-management/shared_states_with_agents.py](./state-management/shared_states_with_agents.py) | Store in shared state once and later reuse across agents |
+| Workflow Kwargs (Custom Context) | [state-management/workflow_kwargs.py](./state-management/workflow_kwargs.py) | Pass custom context (data, user tokens) via kwargs to `@ai_function` tools |
+

 ### visualization

@@ -0,0 +1,132 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import asyncio
+import json
+from typing import Annotated, Any
+
+from agent_framework import ChatMessage, SequentialBuilder, WorkflowOutputEvent, ai_function
+from agent_framework.openai import OpenAIChatClient
+from pydantic import Field
+
+"""
+Sample: Workflow kwargs Flow to @ai_function Tools
+
+This sample demonstrates how to flow custom context (skill data, user tokens, etc.)
+through any workflow pattern to @ai_function tools using the **kwargs pattern.
+
+Key Concepts:
+- Pass custom context as kwargs when invoking workflow.run_stream() or workflow.run()
+- kwargs are stored in SharedState and passed to all agent invocations
+- @ai_function tools receive kwargs via **kwargs parameter
+- Works with Sequential, Concurrent, GroupChat, Handoff, and Magentic patterns
+
+Prerequisites:
+- OpenAI environment variables configured
+"""
+
+
+# Define tools that accept custom context via **kwargs
+@ai_function
+def get_user_data(
+    query: Annotated[str, Field(description="What user data to retrieve")],
+    **kwargs: Any,
+) -> str:
+    """Retrieve user-specific data based on the authenticated context."""
+    user_token = kwargs.get("user_token", {})
+    user_name = user_token.get("user_name", "anonymous")
+    access_level = user_token.get("access_level", "none")
+
+    print(f"\n[get_user_data] Received kwargs keys: {list(kwargs.keys())}")
+    print(f"[get_user_data] User: {user_name}")
+    print(f"[get_user_data] Access level: {access_level}")
+
+    return f"Retrieved data for user {user_name} with {access_level} access: {query}"
+
+
+@ai_function
+def call_api(
+    endpoint_name: Annotated[str, Field(description="Name of the API endpoint to call")],
+    **kwargs: Any,
+) -> str:
+    """Call an API using the configured endpoints from custom_data."""
+    custom_data = kwargs.get("custom_data", {})
+    api_config = custom_data.get("api_config", {})
+
+    base_url = api_config.get("base_url", "unknown")
+    endpoints = api_config.get("endpoints", {})
+
+    print(f"\n[call_api] Received kwargs keys: {list(kwargs.keys())}")
+    print(f"[call_api] Base URL: {base_url}")
+    print(f"[call_api] Available endpoints: {list(endpoints.keys())}")
+
+    if endpoint_name in endpoints:
+        return f"Called {base_url}{endpoints[endpoint_name]} successfully"
+    return f"Endpoint '{endpoint_name}' not found in configuration"
+
+
+async def main() -> None:
+    print("=" * 70)
+    print("Workflow kwargs Flow Demo (SequentialBuilder)")
+    print("=" * 70)
+
+    # Create chat client
+    chat_client = OpenAIChatClient()
+
+    # Create agent with tools that use kwargs
+    agent = chat_client.create_agent(
+        name="assistant",
+        instructions=(
+            "You are a helpful assistant. Use the available tools to help users. "
+            "When asked about user data, use get_user_data. "
+            "When asked to call an API, use call_api."
+        ),
+        tools=[get_user_data, call_api],
+    )
+
+    # Build a simple sequential workflow
+    workflow = SequentialBuilder().participants([agent]).build()
+
+    # Define custom context that will flow to ai_functions via kwargs
+    custom_data = {
+        "api_config": {
+            "base_url": "https://api.example.com",
+            "endpoints": {
+                "users": "/v1/users",
+                "orders": "/v1/orders",
+                "products": "/v1/products",
+            },
+        },
+    }
+
+    user_token = {
+        "user_name": "bob@contoso.com",
+        "access_level": "admin",
+    }
+
+    print("\nCustom Data being passed:")
+    print(json.dumps(custom_data, indent=2))
+    print(f"\nUser: {user_token['user_name']}")
+    print("\n" + "-" * 70)
+    print("Workflow Execution (watch for [tool_name] logs showing kwargs received):")
+    print("-" * 70)
+
+    # Run workflow with kwargs - these will flow through to ai_functions
+    async for event in workflow.run_stream(
+        "Please get my user data and then call the users API endpoint.",
+        custom_data=custom_data,
+        user_token=user_token,
+    ):
+        if isinstance(event, WorkflowOutputEvent):
+            output_data = event.data
+            if isinstance(output_data, list):
+                for item in output_data:
+                    if isinstance(item, ChatMessage) and item.text:
+                        print(f"\n[Final Answer]: {item.text}")
+
+    print("\n" + "=" * 70)
+    print("Sample Complete")
+    print("=" * 70)
+
+
+if __name__ == "__main__":
+    asyncio.run(main())