Python: [BREAKING]: Introducing Options as TypedDict and Generic (#3140)

* WIP typeddict for options

* updated all clients and ChatAgents

* updated everything

* added ADR

* fix mypy

* proper typevar imports

* fixed import

* fixed other imports

* slight update in the sample

* updated from feedback

* fixes

* fixed missing covariants and test fixes

* fixed typing

* updated anthropic thinking config

* ruff fixes

* fixed int tests

* fix tests and mypy

* updated integration tests

* updated docstring and test fix

* improved options handling in obser

* mypy fix

* updated a host of integration tests

* fix tests

* bedrock fix
This commit is contained in:
Eduard van Valkenburg
2026-01-13 16:41:05 +00:00
committed by GitHub
parent 5faa2851bb
commit 3e97425245
111 changed files with 6141 additions and 4715 deletions
@@ -16,7 +16,7 @@ from ._confirmation_strategies import (
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 @@ __all__ = [
"AgentFrameworkAgent",
"add_agent_framework_fastapi_endpoint",
"AGUIChatClient",
"AGUIChatOptions",
"AGUIEventConverter",
"AGUIHttpService",
"AGUIRequest",
"AgentState",
"ConfirmationStrategy",
"DefaultConfirmationStrategy",
"PredictStateConfig",
"RunMetadata",
"TaskPlannerConfirmationStrategy",
"RecipeConfirmationStrategy",
"DocumentWriterConfirmationStrategy",
@@ -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 ._http_service import AGUIHttpService
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 @@ def _apply_server_function_call_unwrap(chat_client: TBaseChatClient) -> TBaseCha
@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 @@ class AGUIChatClient(BaseChatClient):
"""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 @@ class AGUIChatClient(BaseChatClient):
"""
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,23 +360,24 @@ class AGUIChatClient(BaseChatClient):
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:
@@ -342,20 +385,21 @@ class AGUIChatClient(BaseChatClient):
"""
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]
@@ -13,7 +13,7 @@ logger = logging.getLogger(__name__)
def collect_server_tools(agent: Any) -> list[Any]:
"""Collect server tools from ChatAgent or duck-typed agent."""
if isinstance(agent, ChatAgent):
tools_from_agent = agent.chat_options.tools
tools_from_agent = agent.default_options.get("tools")
server_tools = list(tools_from_agent) if tools_from_agent else []
logger.info(f"[TOOLS] Agent has {len(server_tools)} configured tools")
for tool in server_tools:
@@ -23,9 +23,11 @@ def collect_server_tools(agent: Any) -> list[Any]:
return server_tools
try:
chat_options_attr = getattr(agent, "chat_options", None)
if chat_options_attr is not None:
return getattr(chat_options_attr, "tools", None) or []
default_options_attr = getattr(agent, "default_options", None)
if default_options_attr is not None:
if isinstance(default_options_attr, dict):
return default_options_attr.get("tools") or []
return getattr(default_options_attr, "tools", None) or []
except AttributeError:
return []
return []
@@ -319,7 +319,7 @@ class DefaultOrchestrator(Orchestrator):
response_format = None
if isinstance(context.agent, ChatAgent):
response_format = context.agent.chat_options.response_format
response_format = context.agent.default_options.get("response_format")
skip_text_content = response_format is not None
client_tools = convert_agui_tools_to_agent_framework(context.input_data.get("tools"))
@@ -434,10 +434,10 @@ class DefaultOrchestrator(Orchestrator):
run_kwargs: dict[str, Any] = {
"thread": thread,
"tools": tools_param,
"metadata": safe_metadata,
"options": {"metadata": safe_metadata},
}
if safe_metadata:
run_kwargs["store"] = True
run_kwargs["options"]["store"] = True
async def _resolve_approval_responses(
messages: list[Any],
@@ -2,8 +2,23 @@
"""Type definitions for AG-UI integration."""
import sys
from typing import Any, TypedDict
from agent_framework import ChatOptions
if sys.version_info >= (3, 13):
from typing import TypeVar
else:
from typing_extensions import TypeVar
__all__ = [
"AGUIChatOptions",
"AgentState",
"PredictStateConfig",
"RunMetadata",
]
from pydantic import BaseModel, Field
@@ -48,3 +63,76 @@ class AGUIRequest(BaseModel):
None,
description="Optional shared state for agentic generative UI",
)
# region AG-UI Chat Options TypedDict
class AGUIChatOptions(ChatOptions, total=False):
"""AG-UI protocol-specific chat options dict.
Extends base ChatOptions for the AG-UI (Agent-UI) protocol.
AG-UI is a streaming protocol for connecting AI agents to user interfaces.
Options are forwarded to the remote AG-UI server.
See: https://github.com/ag-ui/ag-ui-protocol
Keys:
# Inherited from ChatOptions (forwarded to remote server):
model_id: The model identifier (forwarded as-is to server).
temperature: Sampling temperature.
top_p: Nucleus sampling parameter.
max_tokens: Maximum tokens to generate.
stop: Stop sequences.
tools: List of tools - sent to server so LLM knows about client tools.
Server executes its own tools; client tools execute locally via
@use_function_invocation middleware.
tool_choice: How the model should use tools.
metadata: Metadata dict containing thread_id for conversation continuity.
# Options with limited support (depends on remote server):
frequency_penalty: Forwarded if remote server supports it.
presence_penalty: Forwarded if remote server supports it.
seed: Forwarded if remote server supports it.
response_format: Forwarded if remote server supports it.
logit_bias: Forwarded if remote server supports it.
user: Forwarded if remote server supports it.
# Options not typically used in AG-UI:
store: Not applicable for AG-UI protocol.
allow_multiple_tool_calls: Handled by underlying server.
# AG-UI-specific options:
forward_props: Additional properties to forward to the AG-UI server.
Useful for passing custom parameters to specific server implementations.
context: Shared context/state to send to the server.
Note:
AG-UI is a protocol bridge - actual option support depends on the
remote server implementation. The client sends all options to the
server, which decides how to handle them.
Thread ID management:
- Pass ``thread_id`` in ``metadata`` to maintain conversation continuity
- If not provided, a new thread ID is auto-generated
"""
# AG-UI-specific options
forward_props: dict[str, Any]
"""Additional properties to forward to the AG-UI server."""
context: dict[str, Any]
"""Shared context/state to send to the server."""
# ChatOptions fields not applicable for AG-UI
store: None # type: ignore[misc]
"""Not applicable for AG-UI protocol."""
AGUI_OPTION_TRANSLATIONS: dict[str, str] = {}
"""Maps ChatOptions keys to AG-UI parameter names (protocol uses standard names)."""
TAGUIChatOptions = TypeVar("TAGUIChatOptions", bound=TypedDict, default="AGUIChatOptions", covariant=True) # type: ignore[valid-type]
# endregion