mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
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:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user