Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 9 additions & 9 deletions python/packages/ag-ui/agent_framework_ag_ui/_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -59,17 +59,17 @@ def _unwrap_server_function_call_contents(contents: MutableSequence[Content | di
contents[idx] = content.function_call # type: ignore[assignment, union-attr]


TBaseChatClient = TypeVar("TBaseChatClient", bound=type[BaseChatClient[Any]])
BaseChatClientT = TypeVar("BaseChatClientT", bound=type[BaseChatClient[Any]])

TAGUIChatOptions = TypeVar(
"TAGUIChatOptions",
AGUIChatOptionsT = TypeVar(
"AGUIChatOptionsT",
bound=TypedDict, # type: ignore[valid-type]
default="AGUIChatOptions",
covariant=True,
)


def _apply_server_function_call_unwrap(chat_client: TBaseChatClient) -> TBaseChatClient:
def _apply_server_function_call_unwrap(chat_client: BaseChatClientT) -> BaseChatClientT:
"""Class decorator that unwraps server-side function calls after tool handling."""

original_get_response = chat_client.get_response
Expand Down Expand Up @@ -111,11 +111,11 @@ def _map_update(update: ChatResponseUpdate) -> ChatResponseUpdate:

@_apply_server_function_call_unwrap
class AGUIChatClient(
ChatMiddlewareLayer[TAGUIChatOptions],
FunctionInvocationLayer[TAGUIChatOptions],
ChatTelemetryLayer[TAGUIChatOptions],
BaseChatClient[TAGUIChatOptions],
Generic[TAGUIChatOptions],
ChatMiddlewareLayer[AGUIChatOptionsT],
FunctionInvocationLayer[AGUIChatOptionsT],
ChatTelemetryLayer[AGUIChatOptionsT],
BaseChatClient[AGUIChatOptionsT],
Generic[AGUIChatOptionsT],
):
"""Chat client for communicating with AG-UI compliant servers.

Expand Down
6 changes: 3 additions & 3 deletions python/packages/ag-ui/agent_framework_ag_ui/_types.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,8 +18,8 @@
from typing_extensions import TypedDict # type: ignore # pragma: no cover


TAGUIChatOptions = TypeVar("TAGUIChatOptions", bound=TypedDict, default="AGUIChatOptions", covariant=True) # type: ignore[valid-type]
TResponseModel = TypeVar("TResponseModel", bound=BaseModel | None, default=None)
AGUIChatOptionsT = TypeVar("AGUIChatOptionsT", bound=TypedDict, default="AGUIChatOptions", covariant=True) # type: ignore[valid-type]
ResponseModelT = TypeVar("ResponseModelT", bound=BaseModel | None, default=None)


class PredictStateConfig(TypedDict):
Expand Down Expand Up @@ -84,7 +84,7 @@ class AGUIRequest(BaseModel):
# region AG-UI Chat Options TypedDict


class AGUIChatOptions(ChatOptions[TResponseModel], Generic[TResponseModel], total=False):
class AGUIChatOptions(ChatOptions[ResponseModelT], Generic[ResponseModelT], total=False):
"""AG-UI protocol-specific chat options dict.

Extends base ChatOptions for the AG-UI (Agent-UI) protocol.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -165,10 +165,10 @@
For other requests, use the appropriate tool (create_chart, display_timeline, show_comparison_table).
"""

TOptions = TypeVar("TOptions", bound=TypedDict, default="ChatOptions") # type: ignore[valid-type]
OptionsT = TypeVar("OptionsT", bound=TypedDict, default="ChatOptions") # type: ignore[valid-type]


def ui_generator_agent(chat_client: ChatClientProtocol[TOptions]) -> AgentFrameworkAgent:
def ui_generator_agent(chat_client: ChatClientProtocol[OptionsT]) -> AgentFrameworkAgent:
"""Create a UI generator agent with custom React component rendering.

Args:
Expand Down
18 changes: 9 additions & 9 deletions python/packages/ag-ui/tests/ag_ui/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@
Content,
SupportsAgentRun,
)
from agent_framework._clients import TOptions_co
from agent_framework._clients import OptionsCoT
from agent_framework._middleware import ChatMiddlewareLayer
from agent_framework._tools import FunctionInvocationLayer
from agent_framework._types import ResponseStream
Expand All @@ -37,11 +37,11 @@


class StreamingChatClientStub(
ChatMiddlewareLayer[TOptions_co],
FunctionInvocationLayer[TOptions_co],
ChatTelemetryLayer[TOptions_co],
BaseChatClient[TOptions_co],
Generic[TOptions_co],
ChatMiddlewareLayer[OptionsCoT],
FunctionInvocationLayer[OptionsCoT],
ChatTelemetryLayer[OptionsCoT],
BaseChatClient[OptionsCoT],
Generic[OptionsCoT],
):
"""Typed streaming stub that satisfies ChatClientProtocol."""

Expand All @@ -68,7 +68,7 @@ def get_response(
messages: str | ChatMessage | Sequence[str | ChatMessage],
*,
stream: Literal[False] = ...,
options: TOptions_co | ChatOptions[None] | None = ...,
options: OptionsCoT | ChatOptions[None] | None = ...,
**kwargs: Any,
) -> Awaitable[ChatResponse[Any]]: ...

Expand All @@ -78,7 +78,7 @@ def get_response(
messages: str | ChatMessage | Sequence[str | ChatMessage],
*,
stream: Literal[True],
options: TOptions_co | ChatOptions[Any] | None = ...,
options: OptionsCoT | ChatOptions[Any] | None = ...,
**kwargs: Any,
) -> ResponseStream[ChatResponseUpdate, ChatResponse[Any]]: ...

Expand All @@ -87,7 +87,7 @@ def get_response(
messages: str | ChatMessage | Sequence[str | ChatMessage],
*,
stream: bool = False,
options: TOptions_co | ChatOptions[Any] | None = None,
options: OptionsCoT | ChatOptions[Any] | None = None,
**kwargs: Any,
) -> Awaitable[ChatResponse[Any]] | ResponseStream[ChatResponseUpdate, ChatResponse[Any]]:
self.last_thread = kwargs.get("thread")
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -78,7 +78,7 @@
BETA_FLAGS: Final[list[str]] = ["mcp-client-2025-04-04", "code-execution-2025-08-25"]
STRUCTURED_OUTPUTS_BETA_FLAG: Final[str] = "structured-outputs-2025-11-13"

TResponseModel = TypeVar("TResponseModel", bound=BaseModel | None, default=None)
ResponseModelT = TypeVar("ResponseModelT", bound=BaseModel | None, default=None)


# region Anthropic Chat Options TypedDict
Expand All @@ -102,7 +102,7 @@ class ThinkingConfig(TypedDict, total=False):
budget_tokens: int


class AnthropicChatOptions(ChatOptions[TResponseModel], Generic[TResponseModel], total=False):
class AnthropicChatOptions(ChatOptions[ResponseModelT], Generic[ResponseModelT], total=False):
"""Anthropic-specific chat options.

Extends ChatOptions with options specific to Anthropic's Messages API.
Expand Down Expand Up @@ -160,8 +160,8 @@ class AnthropicChatOptions(ChatOptions[TResponseModel], Generic[TResponseModel],
conversation_id: None # type: ignore[misc]


TAnthropicOptions = TypeVar(
"TAnthropicOptions",
AnthropicOptionsT = TypeVar(
"AnthropicOptionsT",
bound=TypedDict, # type: ignore[valid-type]
default="AnthropicChatOptions",
covariant=True,
Expand Down Expand Up @@ -232,11 +232,11 @@ class AnthropicSettings(AFBaseSettings):


class AnthropicClient(
ChatMiddlewareLayer[TAnthropicOptions],
FunctionInvocationLayer[TAnthropicOptions],
ChatTelemetryLayer[TAnthropicOptions],
BaseChatClient[TAnthropicOptions],
Generic[TAnthropicOptions],
ChatMiddlewareLayer[AnthropicOptionsT],
FunctionInvocationLayer[AnthropicOptionsT],
ChatTelemetryLayer[AnthropicOptionsT],
BaseChatClient[AnthropicOptionsT],
Generic[AnthropicOptionsT],
):
"""Anthropic Chat client with middleware, telemetry, and function invocation support."""

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -38,17 +38,17 @@
from typing_extensions import TypedDict # type: ignore # pragma: no cover


# Type variable for options - allows typed ChatAgent[TOptions] returns
# Type variable for options - allows typed ChatAgent[OptionsCoT] returns
# Default matches AzureAIAgentClient's default options type
TOptions_co = TypeVar(
"TOptions_co",
OptionsCoT = TypeVar(
"OptionsCoT",
bound=TypedDict, # type: ignore[valid-type]
default="AzureAIAgentOptions",
covariant=True,
)


class AzureAIAgentsProvider(Generic[TOptions_co]):
class AzureAIAgentsProvider(Generic[OptionsCoT]):
"""Provider for Azure AI Agent Service V1 (Persistent Agents API).

This provider enables creating, retrieving, and wrapping Azure AI agents as ChatAgent
Expand Down Expand Up @@ -176,10 +176,10 @@ async def create_agent(
| MutableMapping[str, Any]
| Sequence[ToolProtocol | Callable[..., Any] | MutableMapping[str, Any]]
| None = None,
default_options: TOptions_co | None = None,
default_options: OptionsCoT | None = None,
middleware: Sequence[MiddlewareTypes] | None = None,
context_provider: ContextProvider | None = None,
) -> ChatAgent[TOptions_co]:
) -> ChatAgent[OptionsCoT]:
"""Create a new agent on the Azure AI service and return a ChatAgent.

This method creates a persistent agent on the Azure AI service with the specified
Expand Down Expand Up @@ -273,10 +273,10 @@ async def get_agent(
| MutableMapping[str, Any]
| Sequence[ToolProtocol | Callable[..., Any] | MutableMapping[str, Any]]
| None = None,
default_options: TOptions_co | None = None,
default_options: OptionsCoT | None = None,
middleware: Sequence[MiddlewareTypes] | None = None,
context_provider: ContextProvider | None = None,
) -> ChatAgent[TOptions_co]:
) -> ChatAgent[OptionsCoT]:
"""Retrieve an existing agent from the service and return a ChatAgent.

This method fetches an agent by ID from the Azure AI service
Expand Down Expand Up @@ -329,10 +329,10 @@ def as_agent(
| MutableMapping[str, Any]
| Sequence[ToolProtocol | Callable[..., Any] | MutableMapping[str, Any]]
| None = None,
default_options: TOptions_co | None = None,
default_options: OptionsCoT | None = None,
middleware: Sequence[MiddlewareTypes] | None = None,
context_provider: ContextProvider | None = None,
) -> ChatAgent[TOptions_co]:
) -> ChatAgent[OptionsCoT]:
"""Wrap an existing Agent SDK object as a ChatAgent without making HTTP calls.

Use this method when you already have an Agent object from a previous
Expand Down Expand Up @@ -382,10 +382,10 @@ def _to_chat_agent_from_agent(
self,
agent: Agent,
provided_tools: Sequence[ToolProtocol | MutableMapping[str, Any]] | None = None,
default_options: TOptions_co | None = None,
default_options: OptionsCoT | None = None,
middleware: Sequence[MiddlewareTypes] | None = None,
context_provider: ContextProvider | None = None,
) -> ChatAgent[TOptions_co]:
) -> ChatAgent[OptionsCoT]:
"""Create a ChatAgent from an Agent SDK object.

Args:
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -193,8 +193,8 @@ class AzureAIAgentOptions(ChatOptions, total=False):
}
"""Maps ChatOptions keys to Azure AI Agents API parameter names."""

TAzureAIAgentOptions = TypeVar(
"TAzureAIAgentOptions",
AzureAIAgentOptionsT = TypeVar(
"AzureAIAgentOptionsT",
bound=TypedDict, # type: ignore[valid-type]
default="AzureAIAgentOptions",
covariant=True,
Expand All @@ -205,11 +205,11 @@ class AzureAIAgentOptions(ChatOptions, total=False):


class AzureAIAgentClient(
ChatMiddlewareLayer[TAzureAIAgentOptions],
FunctionInvocationLayer[TAzureAIAgentOptions],
ChatTelemetryLayer[TAzureAIAgentOptions],
BaseChatClient[TAzureAIAgentOptions],
Generic[TAzureAIAgentOptions],
ChatMiddlewareLayer[AzureAIAgentOptionsT],
FunctionInvocationLayer[AzureAIAgentOptionsT],
ChatTelemetryLayer[AzureAIAgentOptionsT],
BaseChatClient[AzureAIAgentOptionsT],
Generic[AzureAIAgentOptionsT],
):
"""Azure AI Agent Chat client with middleware, telemetry, and function invocation support."""

Expand Down Expand Up @@ -1296,12 +1296,12 @@ def as_agent(
| MutableMapping[str, Any]
| Sequence[ToolProtocol | Callable[..., Any] | MutableMapping[str, Any]]
| None = None,
default_options: TAzureAIAgentOptions | Mapping[str, Any] | None = None,
default_options: AzureAIAgentOptionsT | Mapping[str, Any] | None = None,
chat_message_store_factory: Callable[[], ChatMessageStoreProtocol] | None = None,
context_provider: ContextProvider | None = None,
middleware: Sequence[MiddlewareTypes] | None = None,
**kwargs: Any,
) -> ChatAgent[TAzureAIAgentOptions]:
) -> ChatAgent[AzureAIAgentOptionsT]:
"""Convert this chat client to a ChatAgent.

This method creates a ChatAgent instance with this client pre-configured.
Expand Down
20 changes: 10 additions & 10 deletions python/packages/azure-ai/agent_framework_azure_ai/_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -60,15 +60,15 @@ class AzureAIProjectAgentOptions(OpenAIResponsesOptions, total=False):
"""Configuration for enabling reasoning capabilities (requires azure.ai.projects.models.Reasoning)."""


TAzureAIClientOptions = TypeVar(
"TAzureAIClientOptions",
AzureAIClientOptionsT = TypeVar(
"AzureAIClientOptionsT",
bound=TypedDict, # type: ignore[valid-type]
default="AzureAIProjectAgentOptions",
covariant=True,
)


class RawAzureAIClient(RawOpenAIResponsesClient[TAzureAIClientOptions], Generic[TAzureAIClientOptions]):
class RawAzureAIClient(RawOpenAIResponsesClient[AzureAIClientOptionsT], Generic[AzureAIClientOptionsT]):
"""Raw Azure AI client without middleware, telemetry, or function invocation layers.

Warning:
Expand Down Expand Up @@ -570,12 +570,12 @@ def as_agent(
| MutableMapping[str, Any]
| Sequence[ToolProtocol | Callable[..., Any] | MutableMapping[str, Any]]
| None = None,
default_options: TAzureAIClientOptions | Mapping[str, Any] | None = None,
default_options: AzureAIClientOptionsT | Mapping[str, Any] | None = None,
chat_message_store_factory: Callable[[], ChatMessageStoreProtocol] | None = None,
context_provider: ContextProvider | None = None,
middleware: Sequence[MiddlewareTypes] | None = None,
**kwargs: Any,
) -> ChatAgent[TAzureAIClientOptions]:
) -> ChatAgent[AzureAIClientOptionsT]:
"""Convert this chat client to a ChatAgent.

This method creates a ChatAgent instance with this client pre-configured.
Expand Down Expand Up @@ -615,11 +615,11 @@ def as_agent(


class AzureAIClient(
ChatMiddlewareLayer[TAzureAIClientOptions],
FunctionInvocationLayer[TAzureAIClientOptions],
ChatTelemetryLayer[TAzureAIClientOptions],
RawAzureAIClient[TAzureAIClientOptions],
Generic[TAzureAIClientOptions],
ChatMiddlewareLayer[AzureAIClientOptionsT],
FunctionInvocationLayer[AzureAIClientOptionsT],
ChatTelemetryLayer[AzureAIClientOptionsT],
RawAzureAIClient[AzureAIClientOptionsT],
Generic[AzureAIClientOptionsT],
):
"""Azure AI client with middleware, telemetry, and function invocation support.

Expand Down
Loading
Loading