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/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