agent/agent-framework
2
0
Fork 0
mirror of https://github.com/microsoft/agent-framework.git synced 2026-06-16 21:04:09 +08:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
* 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
Download Patch File Download Diff File Expand all files Collapse all files
python/packages/ag-ui/agent_framework_ag_ui/__init__.py
+5 -1
View File
@@ -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",
python/packages/ag-ui/agent_framework_ag_ui/_client.py
+62 -18
View File
@@ -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]
python/packages/ag-ui/agent_framework_ag_ui/_orchestration/_tooling.py
+6 -4
View File
@@ -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 []
python/packages/ag-ui/agent_framework_ag_ui/_orchestrators.py
+3 -3
View File
@@ -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],
python/packages/ag-ui/agent_framework_ag_ui/_types.py
+88
View File
@@ -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
python/packages/ag-ui/agent_framework_ag_ui_examples/agents/human_in_the_loop_agent.py
+2 -1
View File
@@ -3,6 +3,7 @@
"""Human-in-the-loop agent demonstrating step customization (Feature 5)."""
from enum import Enum
from typing import Any
from agent_framework import ChatAgent, ChatClientProtocol, ai_function
from pydantic import BaseModel, Field
@@ -42,7 +43,7 @@ def generate_task_steps(steps: list[TaskStep]) -> str:
return f"Generated {len(steps)} execution steps for the task."
def human_in_the_loop_agent(chat_client: ChatClientProtocol) -> ChatAgent:
def human_in_the_loop_agent(chat_client: ChatClientProtocol[Any]) -> ChatAgent[Any]:
"""Create a human-in-the-loop agent using tool-based approach for predictive state.
Args:
python/packages/ag-ui/agent_framework_ag_ui_examples/agents/recipe_agent.py
+2 -1
View File
@@ -3,6 +3,7 @@
"""Recipe agent example demonstrating shared state management (Feature 3)."""
from enum import Enum
from typing import Any
from agent_framework import ChatAgent, ChatClientProtocol, ai_function
from agent_framework.ag_ui import AgentFrameworkAgent, RecipeConfirmationStrategy
@@ -101,7 +102,7 @@ _RECIPE_INSTRUCTIONS = """You are a helpful recipe assistant that creates and mo
"""
def recipe_agent(chat_client: ChatClientProtocol) -> AgentFrameworkAgent:
def recipe_agent(chat_client: ChatClientProtocol[Any]) -> AgentFrameworkAgent:
"""Create a recipe agent with streaming state updates.
Args:
python/packages/ag-ui/agent_framework_ag_ui_examples/agents/research_assistant_agent.py
+3 -2
View File
@@ -3,6 +3,7 @@
"""Example agent demonstrating agentic generative UI with custom events during execution."""
import asyncio
from typing import Any
from agent_framework import ChatAgent, ChatClientProtocol, ai_function
from agent_framework.ag_ui import AgentFrameworkAgent
@@ -87,8 +88,8 @@ _RESEARCH_ASSISTANT_INSTRUCTIONS = (
)
def research_assistant_agent(chat_client: ChatClientProtocol) -> AgentFrameworkAgent:
"""Create a research assistant agent with progress events.
def research_assistant_agent(chat_client: ChatClientProtocol[Any]) -> AgentFrameworkAgent:
"""Create a research assistant agent.
Args:
chat_client: The chat client to use for the agent
python/packages/ag-ui/agent_framework_ag_ui_examples/agents/simple_agent.py
+4 -2
View File
@@ -2,10 +2,12 @@
"""Simple agentic chat example (Feature 1: Agentic Chat)."""
from typing import Any
from agent_framework import ChatAgent, ChatClientProtocol
def simple_agent(chat_client: ChatClientProtocol) -> ChatAgent:
def simple_agent(chat_client: ChatClientProtocol[Any]) -> ChatAgent[Any]:
"""Create a simple chat agent.
Args:
@@ -14,7 +16,7 @@ def simple_agent(chat_client: ChatClientProtocol) -> ChatAgent:
Returns:
A configured ChatAgent instance
"""
return ChatAgent(
return ChatAgent[Any](
name="simple_chat_agent",
instructions="You are a helpful assistant. Be concise and friendly.",
chat_client=chat_client,
python/packages/ag-ui/agent_framework_ag_ui_examples/agents/task_planner_agent.py
+3 -1
View File
@@ -2,6 +2,8 @@
"""Example agent demonstrating human-in-the-loop with function approvals."""
from typing import Any
from agent_framework import ChatAgent, ChatClientProtocol, ai_function
from agent_framework.ag_ui import AgentFrameworkAgent, TaskPlannerConfirmationStrategy
@@ -59,7 +61,7 @@ _TASK_PLANNER_INSTRUCTIONS = (
)
def task_planner_agent(chat_client: ChatClientProtocol) -> AgentFrameworkAgent:
def task_planner_agent(chat_client: ChatClientProtocol[Any]) -> AgentFrameworkAgent:
"""Create a task planner agent with user approval for actions.
Args:
python/packages/ag-ui/agent_framework_ag_ui_examples/agents/task_steps_agent.py
+3 -3
View File
@@ -52,7 +52,7 @@ def generate_task_steps(steps: list[TaskStep]) -> str:
return "Steps generated."
def _create_task_steps_agent(chat_client: ChatClientProtocol) -> AgentFrameworkAgent:
def _create_task_steps_agent(chat_client: ChatClientProtocol[Any]) -> AgentFrameworkAgent:
"""Create the task steps agent using tool-based approach for streaming.
Args:
@@ -61,7 +61,7 @@ def _create_task_steps_agent(chat_client: ChatClientProtocol) -> AgentFrameworkA
Returns:
A configured AgentFrameworkAgent instance
"""
agent = ChatAgent(
agent = ChatAgent[Any](
name="task_steps_agent",
instructions="""You are a helpful assistant that breaks down tasks into actionable steps.
@@ -331,7 +331,7 @@ class TaskStepsAgentWithExecution:
yield run_finished_event
def task_steps_agent_wrapped(chat_client: ChatClientProtocol) -> TaskStepsAgentWithExecution:
def task_steps_agent_wrapped(chat_client: ChatClientProtocol[Any]) -> TaskStepsAgentWithExecution:
"""Create a task steps agent with execution simulation.
Args:
python/packages/ag-ui/agent_framework_ag_ui_examples/agents/ui_generator_agent.py
+14 -6
View File
@@ -2,11 +2,17 @@
"""Example agent demonstrating Tool-based Generative UI (Feature 5)."""
from typing import Any
import sys
from typing import Any, TypedDict
from agent_framework import AIFunction, ChatAgent, ChatClientProtocol
from agent_framework import AIFunction, ChatAgent, ChatClientProtocol, ChatOptions
from agent_framework.ag_ui import AgentFrameworkAgent
if sys.version_info >= (3, 13):
from typing import TypeVar # type: ignore # pragma: no cover
else:
from typing_extensions import TypeVar # type: ignore # pragma: no cover
# Declaration-only tools (func=None) - actual rendering happens on the client side
generate_haiku = AIFunction[Any, str](
name="generate_haiku",
@@ -150,15 +156,17 @@ _UI_GENERATOR_INSTRUCTIONS = """You MUST use the provided tools to generate cont
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]
def ui_generator_agent(chat_client: ChatClientProtocol) -> AgentFrameworkAgent:
"""Create a UI generator agent with frontend rendering tools.
def ui_generator_agent(chat_client: ChatClientProtocol[TOptions]) -> AgentFrameworkAgent:
"""Create a UI generator agent with custom React component rendering.
Args:
chat_client: The chat client to use for the agent
Returns:
A configured AgentFrameworkAgent instance with UI generation tools
A configured AgentFrameworkAgent instance with UI generation capabilities
"""
agent = ChatAgent(
name="ui_generator",
@@ -166,7 +174,7 @@ def ui_generator_agent(chat_client: ChatClientProtocol) -> AgentFrameworkAgent:
chat_client=chat_client,
tools=[generate_haiku, create_chart, display_timeline, show_comparison_table],
# Force tool usage - the LLM MUST call a tool, cannot respond with plain text
chat_options={"tool_choice": "required"},
default_options={"tool_choice": "required"}, # type: ignore
)
return AgentFrameworkAgent(
python/packages/ag-ui/agent_framework_ag_ui_examples/agents/weather_agent.py
+2 -2
View File
@@ -57,7 +57,7 @@ def get_forecast(location: str, days: int = 3) -> str:
return f"{days}-day forecast for {location}:\n" + "\n".join(forecast)
def weather_agent(chat_client: ChatClientProtocol) -> ChatAgent:
def weather_agent(chat_client: ChatClientProtocol[Any]) -> ChatAgent[Any]:
"""Create a weather agent with get_weather and get_forecast tools.
Args:
@@ -66,7 +66,7 @@ def weather_agent(chat_client: ChatClientProtocol) -> ChatAgent:
Returns:
A configured ChatAgent instance with weather tools
"""
return ChatAgent(
return ChatAgent[Any](
name="weather_agent",
instructions=(
"You are a helpful weather assistant. "
python/packages/ag-ui/agent_framework_ag_ui_examples/server/main.py
+6 -1
View File
@@ -4,6 +4,7 @@
import logging
import os
from typing import TYPE_CHECKING
import uvicorn
from agent_framework.ag_ui import add_agent_framework_fastapi_endpoint
@@ -19,6 +20,10 @@ from ..agents.task_steps_agent import task_steps_agent_wrapped
from ..agents.ui_generator_agent import ui_generator_agent
from ..agents.weather_agent import weather_agent
if TYPE_CHECKING:
from agent_framework import ChatOptions
from agent_framework._clients import BaseChatClient
# Configure logging to file and console (disabled by default - set ENABLE_DEBUG_LOGGING=1 to enable)
if os.getenv("ENABLE_DEBUG_LOGGING"):
log_file = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "..", "ag_ui_server.log")
@@ -60,7 +65,7 @@ app.add_middleware(
# Create a shared chat client for all agents
# You can use different chat clients for different agents if needed
chat_client = AzureOpenAIChatClient()
chat_client: BaseChatClient[ChatOptions] = AzureOpenAIChatClient()
# Agentic Chat - basic chat agent
add_agent_framework_fastapi_endpoint(
python/packages/ag-ui/tests/test_ag_ui_client.py
+13 -15
View File
@@ -40,22 +40,22 @@ class TestableAGUIChatClient(AGUIChatClient):
"""Expose message conversion helper."""
return self._convert_messages_to_agui_format(messages)
def get_thread_id(self, chat_options: ChatOptions) -> str:
def get_thread_id(self, options: dict[str, Any]) -> str:
"""Expose thread id helper."""
return self._get_thread_id(chat_options)
return self._get_thread_id(options)
async def inner_get_streaming_response(
self, *, messages: MutableSequence[ChatMessage], chat_options: ChatOptions
self, *, messages: MutableSequence[ChatMessage], options: dict[str, Any]
) -> AsyncIterable[ChatResponseUpdate]:
"""Proxy to protected streaming call."""
async for update in self._inner_get_streaming_response(messages=messages, chat_options=chat_options):
async for update in self._inner_get_streaming_response(messages=messages, options=options):
yield update
async def inner_get_response(
self, *, messages: MutableSequence[ChatMessage], chat_options: ChatOptions
self, *, messages: MutableSequence[ChatMessage], options: dict[str, Any]
) -> ChatResponse:
"""Proxy to protected response call."""
return await self._inner_get_response(messages=messages, chat_options=chat_options)
return await self._inner_get_response(messages=messages, options=options)
class TestAGUIChatClient:
@@ -191,7 +191,7 @@ class TestAGUIChatClient:
chat_options = ChatOptions()
updates: list[ChatResponseUpdate] = []
async for update in client.inner_get_streaming_response(messages=messages, chat_options=chat_options):
async for update in client.inner_get_streaming_response(messages=messages, options=chat_options):
updates.append(update)
assert len(updates) == 4
@@ -221,9 +221,9 @@ class TestAGUIChatClient:
monkeypatch.setattr(client.http_service, "post_run", mock_post_run)
messages = [ChatMessage(role="user", text="Test message")]
chat_options = ChatOptions()
chat_options = {}
response = await client.inner_get_response(messages=messages, chat_options=chat_options)
response = await client.inner_get_response(messages=messages, options=chat_options)
assert response is not None
assert len(response.messages) > 0
@@ -266,7 +266,7 @@ class TestAGUIChatClient:
messages = [ChatMessage(role="user", text="Test with tools")]
chat_options = ChatOptions(tools=[test_tool])
response = await client.inner_get_response(messages=messages, chat_options=chat_options)
response = await client.inner_get_response(messages=messages, options=chat_options)
assert response is not None
@@ -288,10 +288,9 @@ class TestAGUIChatClient:
monkeypatch.setattr(client.http_service, "post_run", mock_post_run)
messages = [ChatMessage(role="user", text="Test server tool execution")]
chat_options = ChatOptions()
updates: list[ChatResponseUpdate] = []
async for update in client.get_streaming_response(messages, chat_options=chat_options):
async for update in client.get_streaming_response(messages):
updates.append(update)
function_calls = [
@@ -332,9 +331,8 @@ class TestAGUIChatClient:
monkeypatch.setattr(client.http_service, "post_run", mock_post_run)
messages = [ChatMessage(role="user", text="Test server tool execution")]
chat_options = ChatOptions(tool_choice="auto", tools=[client_tool])
async for _ in client.get_streaming_response(messages, chat_options=chat_options):
async for _ in client.get_streaming_response(messages, options={"tool_choice": "auto", "tools": [client_tool]}):
pass
async def test_state_transmission(self, monkeypatch: MonkeyPatch) -> None:
@@ -370,6 +368,6 @@ class TestAGUIChatClient:
chat_options = ChatOptions()
response = await client.inner_get_response(messages=messages, chat_options=chat_options)
response = await client.inner_get_response(messages=messages, options=chat_options)
assert response is not None
python/packages/ag-ui/tests/test_agent_wrapper_comprehensive.py
+36 -30
View File
@@ -21,11 +21,15 @@ async def test_agent_initialization_basic():
from agent_framework.ag_ui import AgentFrameworkAgent
async def stream_fn(
messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=[TextContent(text="Hello")])
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
agent = ChatAgent[ChatOptions](
chat_client=StreamingChatClientStub(stream_fn),
name="test_agent",
instructions="Test",
)
wrapper = AgentFrameworkAgent(agent=agent)
assert wrapper.name == "test_agent"
@@ -39,7 +43,7 @@ async def test_agent_initialization_with_state_schema():
from agent_framework.ag_ui import AgentFrameworkAgent
async def stream_fn(
messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=[TextContent(text="Hello")])
@@ -55,7 +59,7 @@ async def test_agent_initialization_with_predict_state_config():
from agent_framework.ag_ui import AgentFrameworkAgent
async def stream_fn(
messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=[TextContent(text="Hello")])
@@ -71,7 +75,7 @@ async def test_agent_initialization_with_pydantic_state_schema():
from agent_framework.ag_ui import AgentFrameworkAgent
async def stream_fn(
messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=[TextContent(text="Hello")])
@@ -94,7 +98,7 @@ async def test_run_started_event_emission():
from agent_framework.ag_ui import AgentFrameworkAgent
async def stream_fn(
messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=[TextContent(text="Hello")])
@@ -118,7 +122,7 @@ async def test_predict_state_custom_event_emission():
from agent_framework.ag_ui import AgentFrameworkAgent
async def stream_fn(
messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=[TextContent(text="Hello")])
@@ -150,7 +154,7 @@ async def test_initial_state_snapshot_with_schema():
from agent_framework.ag_ui import AgentFrameworkAgent
async def stream_fn(
messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=[TextContent(text="Hello")])
@@ -180,7 +184,7 @@ async def test_state_initialization_object_type():
from agent_framework.ag_ui import AgentFrameworkAgent
async def stream_fn(
messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=[TextContent(text="Hello")])
@@ -207,7 +211,7 @@ async def test_state_initialization_array_type():
from agent_framework.ag_ui import AgentFrameworkAgent
async def stream_fn(
messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=[TextContent(text="Hello")])
@@ -234,7 +238,7 @@ async def test_run_finished_event_emission():
from agent_framework.ag_ui import AgentFrameworkAgent
async def stream_fn(
messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=[TextContent(text="Hello")])
@@ -256,7 +260,7 @@ async def test_tool_result_confirm_changes_accepted():
from agent_framework.ag_ui import AgentFrameworkAgent
async def stream_fn(
messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=[TextContent(text="Document updated")])
@@ -303,7 +307,7 @@ async def test_tool_result_confirm_changes_rejected():
from agent_framework.ag_ui import AgentFrameworkAgent
async def stream_fn(
messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=[TextContent(text="OK")])
@@ -337,7 +341,7 @@ async def test_tool_result_function_approval_accepted():
from agent_framework.ag_ui import AgentFrameworkAgent
async def stream_fn(
messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=[TextContent(text="OK")])
@@ -383,7 +387,7 @@ async def test_tool_result_function_approval_rejected():
from agent_framework.ag_ui import AgentFrameworkAgent
async def stream_fn(
messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=[TextContent(text="OK")])
@@ -422,10 +426,11 @@ async def test_thread_metadata_tracking():
thread_metadata: dict[str, Any] = {}
async def stream_fn(
messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
if chat_options.metadata:
thread_metadata.update(chat_options.metadata)
metadata = options.get("metadata")
if metadata:
thread_metadata.update(metadata)
yield ChatResponseUpdate(contents=[TextContent(text="Hello")])
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
@@ -447,15 +452,16 @@ async def test_thread_metadata_tracking():
async def test_state_context_injection():
"""Test that current state is injected into thread metadata."""
from agent_framework.ag_ui import AgentFrameworkAgent
from agent_framework_ag_ui import AgentFrameworkAgent
thread_metadata: dict[str, Any] = {}
async def stream_fn(
messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
if chat_options.metadata:
thread_metadata.update(chat_options.metadata)
metadata = options.get("metadata")
if metadata:
thread_metadata.update(metadata)
yield ChatResponseUpdate(contents=[TextContent(text="Hello")])
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
@@ -484,7 +490,7 @@ async def test_no_messages_provided():
from agent_framework.ag_ui import AgentFrameworkAgent
async def stream_fn(
messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=[TextContent(text="Hello")])
@@ -508,7 +514,7 @@ async def test_message_end_event_emission():
from agent_framework.ag_ui import AgentFrameworkAgent
async def stream_fn(
messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=[TextContent(text="Hello world")])
@@ -536,7 +542,7 @@ async def test_error_handling_with_exception():
from agent_framework.ag_ui import AgentFrameworkAgent
async def stream_fn(
messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
if False:
yield ChatResponseUpdate(contents=[])
@@ -557,7 +563,7 @@ async def test_json_decode_error_in_tool_result():
from agent_framework.ag_ui import AgentFrameworkAgent
async def stream_fn(
messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
if False:
yield ChatResponseUpdate(contents=[])
@@ -594,7 +600,7 @@ async def test_suppressed_summary_with_document_state():
from agent_framework.ag_ui import AgentFrameworkAgent, DocumentWriterConfirmationStrategy
async def stream_fn(
messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=[TextContent(text="Response")])
@@ -647,7 +653,7 @@ async def test_function_approval_mode_executes_tool():
return "2025/12/01 12:00:00"
async def stream_fn(
messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
messages: MutableSequence[ChatMessage], options: ChatOptions, **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
# Capture the messages received by the chat client
messages_received.clear()
@@ -655,9 +661,9 @@ async def test_function_approval_mode_executes_tool():
yield ChatResponseUpdate(contents=[TextContent(text="Processing completed")])
agent = ChatAgent(
chat_client=StreamingChatClientStub(stream_fn),
name="test_agent",
instructions="Test",
chat_client=StreamingChatClientStub(stream_fn),
tools=[get_datetime],
)
wrapper = AgentFrameworkAgent(agent=agent)
@@ -738,7 +744,7 @@ async def test_function_approval_mode_rejection():
return "All data deleted"
async def stream_fn(
messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
messages: MutableSequence[ChatMessage], options: ChatOptions, **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
# Capture the messages received by the chat client
messages_received.clear()
python/packages/ag-ui/tests/test_orchestrators.py
+1 -1
View File
@@ -22,7 +22,7 @@ class DummyAgent:
"""Minimal agent stub to capture run_stream parameters."""
def __init__(self) -> None:
self.chat_options = SimpleNamespace(tools=[server_tool], response_format=None)
self.default_options: dict[str, Any] = {"tools": [server_tool], "response_format": None}
self.tools = [server_tool]
self.chat_client = SimpleNamespace(
function_invocation_configuration=FunctionInvocationConfiguration(),
python/packages/ag-ui/tests/test_orchestrators_coverage.py
+25 -25
View File
@@ -29,7 +29,7 @@ def approval_tool(param: str) -> str:
return f"executed: {param}"
DEFAULT_CHAT_OPTIONS = SimpleNamespace(tools=[approval_tool], response_format=None)
DEFAULT_OPTIONS: dict[str, Any] = {"tools": [approval_tool], "response_format": None}
async def test_human_in_the_loop_json_decode_error() -> None:
@@ -54,7 +54,7 @@ async def test_human_in_the_loop_json_decode_error() -> None:
]
agent = StubAgent(
chat_options=SimpleNamespace(tools=[approval_tool], response_format=None),
default_options={"tools": [approval_tool], "response_format": None},
updates=[AgentRunResponseUpdate(contents=[TextContent(text="response")], role="assistant")],
)
context = TestExecutionContext(
@@ -106,7 +106,7 @@ async def test_sanitize_tool_history_confirm_changes() -> None:
input_data: dict[str, Any] = {"messages": []}
agent = StubAgent(
chat_options=DEFAULT_CHAT_OPTIONS,
default_options=DEFAULT_OPTIONS,
)
context = TestExecutionContext(
input_data=input_data,
@@ -151,7 +151,7 @@ async def test_sanitize_tool_history_orphaned_tool_result() -> None:
orchestrator = DefaultOrchestrator()
input_data: dict[str, Any] = {"messages": []}
agent = StubAgent(
chat_options=DEFAULT_CHAT_OPTIONS,
default_options=DEFAULT_OPTIONS,
)
context = TestExecutionContext(
input_data=input_data,
@@ -191,7 +191,7 @@ async def test_orphaned_tool_result_sanitization() -> None:
}
agent = StubAgent(
chat_options=DEFAULT_CHAT_OPTIONS,
default_options=DEFAULT_OPTIONS,
)
context = TestExecutionContext(
input_data=input_data,
@@ -234,7 +234,7 @@ async def test_deduplicate_messages_empty_tool_results() -> None:
orchestrator = DefaultOrchestrator()
input_data: dict[str, Any] = {"messages": []}
agent = StubAgent(
chat_options=DEFAULT_CHAT_OPTIONS,
default_options=DEFAULT_OPTIONS,
)
context = TestExecutionContext(
input_data=input_data,
@@ -279,7 +279,7 @@ async def test_deduplicate_messages_duplicate_assistant_tool_calls() -> None:
orchestrator = DefaultOrchestrator()
input_data: dict[str, Any] = {"messages": []}
agent = StubAgent(
chat_options=DEFAULT_CHAT_OPTIONS,
default_options=DEFAULT_OPTIONS,
)
context = TestExecutionContext(
input_data=input_data,
@@ -323,7 +323,7 @@ async def test_deduplicate_messages_duplicate_system_messages() -> None:
orchestrator = DefaultOrchestrator()
input_data: dict[str, Any] = {"messages": []}
agent = StubAgent(
chat_options=DEFAULT_CHAT_OPTIONS,
default_options=DEFAULT_OPTIONS,
)
context = TestExecutionContext(
input_data=input_data,
@@ -362,7 +362,7 @@ async def test_state_context_injection() -> None:
}
agent = StubAgent(
chat_options=DEFAULT_CHAT_OPTIONS,
default_options=DEFAULT_OPTIONS,
)
context = TestExecutionContext(
input_data=input_data,
@@ -407,7 +407,7 @@ async def test_state_context_injection_with_tool_calls_and_input_state() -> None
orchestrator = DefaultOrchestrator()
input_data: dict[str, Any] = {"messages": [], "state": {"weather": "sunny"}}
agent = StubAgent(
chat_options=DEFAULT_CHAT_OPTIONS,
default_options=DEFAULT_OPTIONS,
)
context = TestExecutionContext(
input_data=input_data,
@@ -449,7 +449,7 @@ async def test_structured_output_processing() -> None:
# Agent with structured output
agent = StubAgent(
chat_options=DEFAULT_CHAT_OPTIONS,
default_options=DEFAULT_OPTIONS,
updates=[
AgentRunResponseUpdate(
contents=[TextContent(text='{"ingredients": ["tomato"], "message": "Added tomato"}')],
@@ -457,7 +457,7 @@ async def test_structured_output_processing() -> None:
)
],
)
agent.chat_options.response_format = RecipeState
agent.default_options["response_format"] = RecipeState
context = TestExecutionContext(
input_data=input_data,
@@ -510,9 +510,9 @@ async def test_duplicate_client_tools_filtered() -> None:
}
agent = StubAgent(
chat_options=DEFAULT_CHAT_OPTIONS,
default_options=DEFAULT_OPTIONS,
)
agent.chat_options.tools = [get_weather]
agent.default_options["tools"] = [get_weather]
context = TestExecutionContext(
input_data=input_data,
@@ -559,9 +559,9 @@ async def test_unique_client_tools_merged() -> None:
}
agent = StubAgent(
chat_options=DEFAULT_CHAT_OPTIONS,
default_options=DEFAULT_OPTIONS,
)
agent.chat_options.tools = [server_tool]
agent.default_options["tools"] = [server_tool]
context = TestExecutionContext(
input_data=input_data,
@@ -587,7 +587,7 @@ async def test_empty_messages_handling() -> None:
input_data: dict[str, Any] = {"messages": []}
agent = StubAgent(
chat_options=DEFAULT_CHAT_OPTIONS,
default_options=DEFAULT_OPTIONS,
)
context = TestExecutionContext(
input_data=input_data,
@@ -621,7 +621,7 @@ async def test_all_messages_filtered_handling() -> None:
}
agent = StubAgent(
chat_options=DEFAULT_CHAT_OPTIONS,
default_options=DEFAULT_OPTIONS,
)
context = TestExecutionContext(
input_data=input_data,
@@ -663,7 +663,7 @@ async def test_confirm_changes_with_invalid_json_fallback() -> None:
orchestrator = DefaultOrchestrator()
input_data: dict[str, Any] = {"messages": []}
agent = StubAgent(
chat_options=DEFAULT_CHAT_OPTIONS,
default_options=DEFAULT_OPTIONS,
)
context = TestExecutionContext(
input_data=input_data,
@@ -706,7 +706,7 @@ async def test_confirm_changes_closes_active_message_before_finish() -> None:
orchestrator = DefaultOrchestrator()
input_data: dict[str, Any] = {"messages": [{"role": "user", "content": "Start"}]}
agent = StubAgent(
chat_options=DEFAULT_CHAT_OPTIONS,
default_options=DEFAULT_OPTIONS,
updates=updates,
)
context = TestExecutionContext(
@@ -751,7 +751,7 @@ async def test_tool_result_kept_when_call_id_matches() -> None:
orchestrator = DefaultOrchestrator()
input_data: dict[str, Any] = {"messages": []}
agent = StubAgent(
chat_options=DEFAULT_CHAT_OPTIONS,
default_options=DEFAULT_OPTIONS,
)
context = TestExecutionContext(
input_data=input_data,
@@ -781,7 +781,7 @@ async def test_agent_protocol_fallback_paths() -> None:
"""Custom agent without ChatAgent type."""
def __init__(self) -> None:
self.chat_options = SimpleNamespace(tools=[], response_format=None)
self.default_options: dict[str, Any] = {"tools": [], "response_format": None}
self.chat_client = SimpleNamespace(function_invocation_configuration=SimpleNamespace())
self.messages_received: list[Any] = []
@@ -827,7 +827,7 @@ async def test_initial_state_snapshot_with_array_schema() -> None:
orchestrator = DefaultOrchestrator()
input_data: dict[str, Any] = {"messages": [], "state": {}}
agent = StubAgent(
chat_options=DEFAULT_CHAT_OPTIONS,
default_options=DEFAULT_OPTIONS,
)
context = TestExecutionContext(
input_data=input_data,
@@ -859,9 +859,9 @@ async def test_response_format_skip_text_content() -> None:
input_data: dict[str, Any] = {"messages": []}
agent = StubAgent(
chat_options=DEFAULT_CHAT_OPTIONS,
default_options=DEFAULT_OPTIONS,
)
agent.chat_options.response_format = OutputModel
agent.default_options["response_format"] = OutputModel
context = TestExecutionContext(
input_data=input_data,
python/packages/ag-ui/tests/test_structured_output.py
+11 -11
View File
@@ -40,14 +40,14 @@ async def test_structured_output_with_recipe():
from agent_framework.ag_ui import AgentFrameworkAgent
async def stream_fn(
messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
messages: MutableSequence[ChatMessage], options: ChatOptions, **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
yield ChatResponseUpdate(
contents=[TextContent(text='{"recipe": {"name": "Pasta"}, "message": "Here is your recipe"}')]
)
agent = ChatAgent(name="test", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
agent.chat_options = ChatOptions(response_format=RecipeOutput)
agent.default_options = ChatOptions(response_format=RecipeOutput)
wrapper = AgentFrameworkAgent(
agent=agent,
@@ -78,7 +78,7 @@ async def test_structured_output_with_steps():
from agent_framework.ag_ui import AgentFrameworkAgent
async def stream_fn(
messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
messages: MutableSequence[ChatMessage], options: ChatOptions, **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
steps_data = {
"steps": [
@@ -89,7 +89,7 @@ async def test_structured_output_with_steps():
yield ChatResponseUpdate(contents=[TextContent(text=json.dumps(steps_data))])
agent = ChatAgent(name="test", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
agent.chat_options = ChatOptions(response_format=StepsOutput)
agent.default_options = ChatOptions(response_format=StepsOutput)
wrapper = AgentFrameworkAgent(
agent=agent,
@@ -124,7 +124,7 @@ async def test_structured_output_with_no_schema_match():
agent = ChatAgent(
name="test", instructions="Test", chat_client=StreamingChatClientStub(stream_from_updates(updates))
)
agent.chat_options = ChatOptions(response_format=GenericOutput)
agent.default_options = ChatOptions(response_format=GenericOutput)
wrapper = AgentFrameworkAgent(
agent=agent,
@@ -154,12 +154,12 @@ async def test_structured_output_without_schema():
info: str
async def stream_fn(
messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
messages: MutableSequence[ChatMessage], options: ChatOptions, **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=[TextContent(text='{"data": {"key": "value"}, "info": "processed"}')])
agent = ChatAgent(name="test", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
agent.chat_options = ChatOptions(response_format=DataOutput)
agent.default_options = ChatOptions(response_format=DataOutput)
wrapper = AgentFrameworkAgent(
agent=agent,
@@ -213,13 +213,13 @@ async def test_structured_output_with_message_field():
from agent_framework.ag_ui import AgentFrameworkAgent
async def stream_fn(
messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
messages: MutableSequence[ChatMessage], options: ChatOptions, **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
output_data = {"recipe": {"name": "Salad"}, "message": "Fresh salad recipe ready"}
yield ChatResponseUpdate(contents=[TextContent(text=json.dumps(output_data))])
agent = ChatAgent(name="test", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
agent.chat_options = ChatOptions(response_format=RecipeOutput)
agent.default_options = ChatOptions(response_format=RecipeOutput)
wrapper = AgentFrameworkAgent(
agent=agent,
@@ -248,13 +248,13 @@ async def test_empty_updates_no_structured_processing():
from agent_framework.ag_ui import AgentFrameworkAgent
async def stream_fn(
messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
messages: MutableSequence[ChatMessage], options: ChatOptions, **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
if False:
yield ChatResponseUpdate(contents=[])
agent = ChatAgent(name="test", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
agent.chat_options = ChatOptions(response_format=RecipeOutput)
agent.default_options = ChatOptions(response_format=RecipeOutput)
wrapper = AgentFrameworkAgent(agent=agent)
python/packages/ag-ui/tests/utils_test_ag_ui.py
+21 -11
View File
@@ -2,9 +2,10 @@
"""Shared test stubs for AG-UI tests."""
import sys
from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Callable, MutableSequence
from types import SimpleNamespace
from typing import Any
from typing import Any, Generic
from agent_framework import (
AgentProtocol,
@@ -13,20 +14,25 @@ from agent_framework import (
AgentThread,
BaseChatClient,
ChatMessage,
ChatOptions,
ChatResponse,
ChatResponseUpdate,
TextContent,
)
from agent_framework._clients import TOptions_co
from agent_framework_ag_ui._message_adapters import _deduplicate_messages, _sanitize_tool_history
from agent_framework_ag_ui._orchestrators import ExecutionContext
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
StreamFn = Callable[..., AsyncIterator[ChatResponseUpdate]]
ResponseFn = Callable[..., Awaitable[ChatResponse]]
class StreamingChatClientStub(BaseChatClient):
class StreamingChatClientStub(BaseChatClient[TOptions_co], Generic[TOptions_co]):
"""Typed streaming stub that satisfies ChatClientProtocol."""
def __init__(self, stream_fn: StreamFn, response_fn: ResponseFn | None = None) -> None:
@@ -34,20 +40,22 @@ class StreamingChatClientStub(BaseChatClient):
self._stream_fn = stream_fn
self._response_fn = response_fn
@override
async def _inner_get_streaming_response(
self, *, messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
self, *, messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
async for update in self._stream_fn(messages, chat_options, **kwargs):
async for update in self._stream_fn(messages, options, **kwargs):
yield update
@override
async def _inner_get_response(
self, *, messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
self, *, messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
) -> ChatResponse:
if self._response_fn is not None:
return await self._response_fn(messages, chat_options, **kwargs)
return await self._response_fn(messages, options, **kwargs)
contents: list[Any] = []
async for update in self._stream_fn(messages, chat_options, **kwargs):
async for update in self._stream_fn(messages, options, **kwargs):
contents.extend(update.contents)
return ChatResponse(
@@ -60,7 +68,7 @@ def stream_from_updates(updates: list[ChatResponseUpdate]) -> StreamFn:
"""Create a stream function that yields from a static list of updates."""
async def _stream(
messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
for update in updates:
yield update
@@ -77,14 +85,16 @@ class StubAgent(AgentProtocol):
*,
agent_id: str = "stub-agent",
agent_name: str | None = "stub-agent",
chat_options: Any | None = None,
default_options: Any | None = None,
chat_client: Any | None = None,
) -> None:
self.id = agent_id
self.name = agent_name
self.description = "stub agent"
self.updates = updates or [AgentRunResponseUpdate(contents=[TextContent(text="response")], role="assistant")]
self.chat_options = chat_options or SimpleNamespace(tools=None, response_format=None)
self.default_options: dict[str, Any] = (
default_options if isinstance(default_options, dict) else {"tools": None, "response_format": None}
)
self.chat_client = chat_client or SimpleNamespace(function_invocation_configuration=None)
self.messages_received: list[Any] = []
self.tools_received: list[Any] | None = None