Python: [BREAKING]: Introducing Options as TypedDict and Generic

.github/workflows/python-merge-tests.yml

@@ -97,7 +97,7 @@ jobs:
         id: azure-functions-setup
       - name: Test with pytest
         timeout-minutes: 10
-        run: uv run poe all-tests -n logical --dist loadfile --dist worksteal --timeout 300 --retries 3 --retry-delay 10
+        run: uv run poe all-tests -n logical --dist loadfile --dist worksteal --timeout 600 --retries 3 --retry-delay 10
         working-directory: ./python
       - name: Test core samples
         timeout-minutes: 10

docs/decisions/0011-python-typeddict-options.md

@@ -0,0 +1,129 @@
+---
+# These are optional elements. Feel free to remove any of them.
+status: proposed
+contact: eavanvalkenburg
+date: 2026-01-08
+deciders: eavanvalkenburg, markwallace-microsoft,  sphenry, alliscode, johanst, brettcannon
+consulted: taochenosu, moonbox3, dmytrostruk, giles17
+---
+
+# Leveraging TypedDict and Generic Options in Python Chat Clients
+
+## Context and Problem Statement
+
+The Agent Framework Python SDK provides multiple chat client implementations for different providers (OpenAI, Anthropic, Azure AI, Bedrock, Ollama, etc.). Each provider has unique configuration options beyond the common parameters defined in `ChatOptions`. Currently, developers using these clients lack type safety and IDE autocompletion for provider-specific options, leading to runtime errors and a poor developer experience.
+
+How can we provide type-safe, discoverable options for each chat client while maintaining a consistent API across all implementations?
+
+## Decision Drivers
+
+- **Type Safety**: Developers should get compile-time/static analysis errors when using invalid options
+- **IDE Support**: Full autocompletion and inline documentation for all available options
+- **Extensibility**: Users should be able to define custom options that extend provider-specific options
+- **Consistency**: All chat clients should follow the same pattern for options handling
+- **Provider Flexibility**: Each provider can expose its unique options without affecting the common interface
+
+## Considered Options
+
+- **Option 1: Status Quo - Class `ChatOptions` with `**kwargs`**
+- **Option 2: TypedDict with Generic Type Parameters**
+
+### Option 1: Status Quo - Class `ChatOptions` with `**kwargs`
+
+The current approach uses a base `ChatOptions` Class with common parameters, and provider-specific options are passed via `**kwargs` or loosely typed dictionaries.
+
+```python
+# Current usage - no type safety for provider-specific options
+response = await client.get_response(
+    messages=messages,
+    temperature=0.7,
+    top_k=40,
+    random=42, # No validation
+)
+```
+
+**Pros:**
+- Simple implementation
+- Maximum flexibility
+
+**Cons:**
+- No type checking for provider-specific options
+- No IDE autocompletion for available options
+- Runtime errors for typos or invalid options
+- Documentation must be consulted for each provider
+
+### Option 2: TypedDict with Generic Type Parameters (Chosen)
+
+Each chat client is parameterized with a TypeVar bound to a provider-specific `TypedDict` that extends `ChatOptions`. This enables full type safety and IDE support.
+
+```python
+# Provider-specific TypedDict
+class AnthropicChatOptions(ChatOptions, total=False):
+    """Anthropic-specific chat options."""
+    top_k: int
+    thinking: ThinkingConfig
+    # ... other Anthropic-specific options
+
+# Generic chat client
+class AnthropicChatClient(ChatClientBase[TAnthropicChatOptions]):
+    ...
+
+client = AnthropicChatClient(...)
+
+# Usage with full type safety
+response = await client.get_response(
+    messages=messages,
+    options={
+        "temperature": 0.7,
+        "top_k": 40,
+        "random": 42,  # fails type checking and IDE would flag this
+    }
+)
+
+# Users can extend for custom options
+class MyAnthropicOptions(AnthropicChatOptions, total=False):
+    custom_field: str
+
+
+client = AnthropicChatClient[MyAnthropicOptions](...)
+
+# Usage of custom options with full type safety
+response = await client.get_response(
+    messages=messages,
+    options={
+        "temperature": 0.7,
+        "top_k": 40,
+        "custom_field": "value",
+    }
+)
+
+```
+
+**Pros:**
+- Full type safety with static analysis
+- IDE autocompletion for all options
+- Compile-time error detection
+- Self-documenting through type hints
+- Users can extend options for their specific needs or advances in models
+
+**Cons:**
+- More complex implementation
+- Some type: ignore comments needed for TypedDict field overrides
+- Minor: Requires TypeVar with default (Python 3.13+ or typing_extensions)
+
+> [NOTE!]
+> In .NET this is already achieved through overloads on the `GetResponseAsync` method for each provider-specific options class, e.g., `AnthropicChatOptions`, `OpenAIChatOptions`, etc. So this does not apply to .NET.
+
+### Implementation Details
+
+1. **Base Protocol**: `ChatClientProtocol[TOptions]` is generic over options type, with default set to `ChatOptions` (the new TypedDict)
+2. **Provider TypedDicts**: Each provider defines its options extending `ChatOptions`
+    They can even override fields with type=None to indicate they are not supported.
+3. **TypeVar Pattern**: `TProviderOptions = TypeVar("TProviderOptions", bound=TypedDict, default=ProviderChatOptions, contravariant=True)`
+4. **Option Translation**: Common options are kept in place,and explicitly documented in the Options class how they are used. (e.g., `user` → `metadata.user_id`) in `_prepare_options` (for Anthropic) to preserve easy use of common options.
+
+## Decision Outcome
+
+Chosen option: **"Option 2: TypedDict with Generic Type Parameters"**, because it provides full type safety, excellent IDE support with autocompletion, and allows users to extend provider-specific options for their use cases. Extended this Generic to ChatAgents in order to also properly type the options used in agent construction and run methods.
+
+See [typed_options.py](../../python/samples/getting_started/chat_client/typed_options.py) for a complete example demonstrating the usage of typed options with custom extensions.

python/.cspell.json

@@ -25,6 +25,7 @@
     "words": [
         "aeiou",
         "aiplatform",
+        "agui",
         "azuredocindex",
         "azuredocs",
         "azurefunctions",

python/packages/a2a/agent_framework_a2a/_agent.py

@@ -189,7 +189,7 @@ async def __aexit__(
 
     async def run(
         self,
-        messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
+        messages: str | ChatMessage | Sequence[str | ChatMessage] | None = None,
         *,
         thread: AgentThread | None = None,
         **kwargs: Any,
@@ -216,7 +216,7 @@ async def run(
 
     async def run_stream(
         self,
-        messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
+        messages: str | ChatMessage | Sequence[str | ChatMessage] | None = None,
         *,
         thread: AgentThread | None = None,
         **kwargs: Any,

python/packages/ag-ui/agent_framework_ag_ui/__init__.py

@@ -16,7 +16,7 @@
 from ._endpoint import add_agent_framework_fastapi_endpoint
 from ._event_converters import AGUIEventConverter
 from ._http_service import AGUIHttpService
-from ._types import AGUIRequest
+from ._types import AgentState, AGUIChatOptions, AGUIRequest, PredictStateConfig, RunMetadata
 
 try:
     __version__ = importlib.metadata.version(__name__)
@@ -30,11 +30,15 @@
     "AgentFrameworkAgent",
     "add_agent_framework_fastapi_endpoint",
     "AGUIChatClient",
+    "AGUIChatOptions",
     "AGUIEventConverter",
     "AGUIHttpService",
     "AGUIRequest",
+    "AgentState",
     "ConfirmationStrategy",
     "DefaultConfirmationStrategy",
+    "PredictStateConfig",
+    "RunMetadata",
     "TaskPlannerConfirmationStrategy",
     "RecipeConfirmationStrategy",
     "DocumentWriterConfirmationStrategy",

python/packages/ag-ui/agent_framework_ag_ui/_client.py

@@ -4,17 +4,17 @@
 
 import json
 import logging
+import sys
 import uuid
 from collections.abc import AsyncIterable, MutableSequence
 from functools import wraps
-from typing import Any, TypeVar, cast
+from typing import TYPE_CHECKING, Any, Generic, cast
 
 import httpx
 from agent_framework import (
     AIFunction,
     BaseChatClient,
     ChatMessage,
-    ChatOptions,
     ChatResponse,
     ChatResponseUpdate,
     DataContent,
@@ -30,6 +30,26 @@
 from ._message_adapters import agent_framework_messages_to_agui
 from ._utils import convert_tools_to_agui_format
 
+if TYPE_CHECKING:
+    from ._types import AGUIChatOptions
+
+from typing import TypedDict
+
+if sys.version_info >= (3, 13):
+    from typing import TypeVar
+else:
+    from typing_extensions import TypeVar
+
+if sys.version_info >= (3, 12):
+    from typing import override  # type: ignore # pragma: no cover
+else:
+    from typing_extensions import override  # type: ignore[import] # pragma: no cover
+
+if sys.version_info >= (3, 11):
+    from typing import Self  # pragma: no cover
+else:
+    from typing_extensions import Self  # pragma: no cover
+
 logger: logging.Logger = logging.getLogger(__name__)
 
 
@@ -55,7 +75,14 @@ def _unwrap_server_function_call_contents(contents: MutableSequence[Contents | d
             contents[idx] = content.function_call_content  # type: ignore[assignment]
 
 
-TBaseChatClient = TypeVar("TBaseChatClient", bound=type[BaseChatClient])
+TBaseChatClient = TypeVar("TBaseChatClient", bound=type[BaseChatClient[Any]])
+
+TAGUIChatOptions = TypeVar(
+    "TAGUIChatOptions",
+    bound=TypedDict,  # type: ignore[valid-type]
+    default="AGUIChatOptions",
+    covariant=True,
+)
 
 
 def _apply_server_function_call_unwrap(chat_client: TBaseChatClient) -> TBaseChatClient:
@@ -91,7 +118,7 @@ async def response_wrapper(self, *args: Any, **kwargs: Any) -> ChatResponse:
 @use_function_invocation
 @use_instrumentation
 @use_chat_middleware
-class AGUIChatClient(BaseChatClient):
+class AGUIChatClient(BaseChatClient[TAGUIChatOptions], Generic[TAGUIChatOptions]):
     """Chat client for communicating with AG-UI compliant servers.
 
     This client implements the BaseChatClient interface and automatically handles:
@@ -168,6 +195,19 @@ class AGUIChatClient(BaseChatClient):
             async with AGUIChatClient(endpoint="http://localhost:8888/") as client:
                 response = await client.get_response("Hello!")
                 print(response.messages[0].text)
+
+        Using custom ChatOptions with type safety:
+
+        .. code-block:: python
+
+            from typing import TypedDict
+            from agent_framework_ag_ui import AGUIChatClient, AGUIChatOptions
+
+            class MyOptions(AGUIChatOptions, total=False):
+                my_custom_option: str
+
+            client: AGUIChatClient[MyOptions] = AGUIChatClient(endpoint="http://localhost:8888/")
+            response = await client.get_response("Hello", options={"my_custom_option": "value"})
     """
 
     OTEL_PROVIDER_NAME = "agui"
@@ -201,7 +241,7 @@ async def close(self) -> None:
         """Close the HTTP client."""
         await self._http_service.close()
 
-    async def __aenter__(self) -> "AGUIChatClient":
+    async def __aenter__(self) -> Self:
         """Enter async context manager."""
         return self
 
@@ -280,36 +320,38 @@ def _convert_messages_to_agui_format(self, messages: list[ChatMessage]) -> list[
         """
         return agent_framework_messages_to_agui(messages)
 
-    def _get_thread_id(self, chat_options: ChatOptions) -> str:
+    def _get_thread_id(self, options: dict[str, Any]) -> str:
         """Get or generate thread ID from chat options.
 
         Args:
-            chat_options: Chat options containing metadata
+            options: Chat options containing metadata
 
         Returns:
             Thread ID string
         """
         thread_id = None
-        if chat_options.metadata:
-            thread_id = chat_options.metadata.get("thread_id")
+        metadata = options.get("metadata")
+        if metadata:
+            thread_id = metadata.get("thread_id")
 
         if not thread_id:
             thread_id = f"thread_{uuid.uuid4().hex}"
 
         return thread_id
 
+    @override
     async def _inner_get_response(
         self,
         *,
         messages: MutableSequence[ChatMessage],
-        chat_options: ChatOptions,
+        options: dict[str, Any],
         **kwargs: Any,
     ) -> ChatResponse:
         """Internal method to get non-streaming response.
 
         Keyword Args:
             messages: List of chat messages
-            chat_options: Chat options for the request
+            options: Chat options for the request
             **kwargs: Additional keyword arguments
 
         Returns:
@@ -318,44 +360,46 @@ async def _inner_get_response(
         return await ChatResponse.from_chat_response_generator(
             self._inner_get_streaming_response(
                 messages=messages,
-                chat_options=chat_options,
+                options=options,
                 **kwargs,
             )
         )
 
+    @override
     async def _inner_get_streaming_response(
         self,
         *,
         messages: MutableSequence[ChatMessage],
-        chat_options: ChatOptions,
+        options: dict[str, Any],
         **kwargs: Any,
     ) -> AsyncIterable[ChatResponseUpdate]:
         """Internal method to get streaming response.
 
         Keyword Args:
             messages: List of chat messages
-            chat_options: Chat options for the request
+            options: Chat options for the request
             **kwargs: Additional keyword arguments
 
         Yields:
             ChatResponseUpdate objects
         """
         messages_to_send, state = self._extract_state_from_messages(messages)
 
-        thread_id = self._get_thread_id(chat_options)
+        thread_id = self._get_thread_id(options)
         run_id = f"run_{uuid.uuid4().hex}"
 
         agui_messages = self._convert_messages_to_agui_format(messages_to_send)
 
         # Send client tools to server so LLM knows about them
         # Client tools execute via ChatAgent's @use_function_invocation wrapper
-        agui_tools = convert_tools_to_agui_format(chat_options.tools)
+        agui_tools = convert_tools_to_agui_format(options.get("tools"))
 
         # Build set of client tool names (matches .NET clientToolSet)
         # Used to distinguish client vs server tools in response stream
         client_tool_set: set[str] = set()
-        if chat_options.tools:
-            for tool in chat_options.tools:
+        tools = options.get("tools")
+        if tools:
+            for tool in tools:
                 if hasattr(tool, "name"):
                     client_tool_set.add(tool.name)  # type: ignore[arg-type]
         self._last_client_tool_set = client_tool_set  # type: ignore[attr-defined]