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] Types API Review improvements (#3647)

Browse Source 
* Replace Role and FinishReason classes with NewType + Literal

- Remove EnumLike metaclass from _types.py
- Replace Role class with NewType('Role', str) + RoleLiteral
- Replace FinishReason class with NewType('FinishReason', str) + FinishReasonLiteral
- Update all usages across codebase to use string literals
- Remove .value access patterns (direct string comparison now works)
- Add backward compatibility for legacy dict serialization format
- Update tests to reflect new string-based types

Addresses #3591, #3615

* Simplify ChatResponse and AgentResponse type hints (#3592)

- Remove overloads from ChatResponse.__init__
- Remove text parameter from ChatResponse.__init__
- Remove | dict[str, Any] from finish_reason and usage_details params
- Remove **kwargs from AgentResponse.__init__
- Both now accept ChatMessage | Sequence[ChatMessage] | None for messages
- Update docstrings and examples to reflect changes
- Fix tests that were using removed kwargs
- Fix Role type hint usage in ag-ui utils

* Remove text parameter from ChatResponseUpdate and AgentResponseUpdate (#3597)

- Remove text parameter from ChatResponseUpdate.__init__
- Remove text parameter from AgentResponseUpdate.__init__
- Remove **kwargs from both update classes
- Simplify contents parameter type to Sequence[Content] | None
- Update all usages to use contents=[Content.from_text(...)] pattern
- Fix imports in test files
- Update docstrings and examples

* Rename from_chat_response_updates to from_updates (#3593)

- ChatResponse.from_chat_response_updates → ChatResponse.from_updates
- ChatResponse.from_chat_response_generator → ChatResponse.from_update_generator
- AgentResponse.from_agent_run_response_updates → AgentResponse.from_updates

* Remove try_parse_value method from ChatResponse and AgentResponse (#3595)

- Remove try_parse_value method from ChatResponse
- Remove try_parse_value method from AgentResponse
- Remove try_parse_value calls from from_updates and from_update_generator methods
- Update samples to use try/except with response.value instead
- Update tests to use response.value pattern
- Users should now use response.value with try/except for safe parsing

* Add agent_id to AgentResponse and clarify author_name documentation (#3596)

- Add agent_id parameter to AgentResponse class
- Document that author_name is on ChatMessage objects, not responses
- Update ChatResponse docstring with author_name note
- Update AgentResponse docstring with author_name note

* Simplify ChatMessage.__init__ signature (#3618)

- Make contents a positional argument accepting Sequence[Content | str]
- Auto-convert strings in contents to TextContent
- Remove overloads, keep text kwarg for backward compatibility with serialization
- Update _parse_content_list to handle string items
- Update all usages across codebase to use new format: ChatMessage("role", ["text"])

* Allow Content as input on run and get_response

- Update prepare_messages and normalize_messages to accept Content
- Update type signatures in _agents.py and _clients.py
- Add tests for Content input handling

* Fix ChatMessage usage across packages and samples

Update all remaining ChatMessage(role=..., text=...) to use new
ChatMessage('role', ['text']) signature.

* Fix Role string usage and response format parsing

- Fix redis provider: remove .value access on string literals
- Fix durabletask ensure_response_format: set _response_format before accessing .value

* Fix ollama .value and ai_model_id issues, handle None in content list

- Fix ollama _chat_client: remove .value on string literals
- Fix ollama _chat_client: rename ai_model_id to model_id
- Fix _parse_content_list: skip None values gracefully

* Fix A2AAgent type signature to include Content

* Fix Role/FinishReason NewType dict annotations and improve test coverage to 95%

* Fix mypy errors for Role/FinishReason NewType usage

* Fix Role.TOOL and Role.ASSISTANT usage in _orchestrator_helpers.py

* Fix Role NewType usage in durabletask _models.py
This commit is contained in:
Eduard van Valkenburg
2026-02-04 10:13:23 +00:00
committed by GitHub
parent ef798629e5
commit 838a7fd61d
341 changed files with 3766 additions and 3228 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/packages/ag-ui/agent_framework_ag_ui/_client.py
+1 -1
View File
@@ -334,7 +334,7 @@ class AGUIChatClient(BaseChatClient[TAGUIChatOptions], Generic[TAGUIChatOptions]
Returns:
ChatResponse object
"""
return await ChatResponse.from_chat_response_generator(
return await ChatResponse.from_update_generator(
self._inner_get_streaming_response(
messages=messages,
options=options,
python/packages/ag-ui/agent_framework_ag_ui/_event_converters.py
+10 -12
View File
@@ -7,8 +7,6 @@ from typing import Any
from agent_framework import (
ChatResponseUpdate,
Content,
FinishReason,
Role,
)
@@ -86,7 +84,7 @@ class AGUIEventConverter:
self.run_id = event.get("runId")
return ChatResponseUpdate(
role=Role.ASSISTANT,
role="assistant",
contents=[],
additional_properties={
"thread_id": self.thread_id,
@@ -98,7 +96,7 @@ class AGUIEventConverter:
"""Handle TEXT_MESSAGE_START event."""
self.current_message_id = event.get("messageId")
return ChatResponseUpdate(
role=Role.ASSISTANT,
role="assistant",
message_id=self.current_message_id,
contents=[],
)
@@ -112,7 +110,7 @@ class AGUIEventConverter:
self.current_message_id = message_id
return ChatResponseUpdate(
role=Role.ASSISTANT,
role="assistant",
message_id=self.current_message_id,
contents=[Content.from_text(text=delta)],
)
@@ -128,7 +126,7 @@ class AGUIEventConverter:
self.accumulated_tool_args = ""
return ChatResponseUpdate(
role=Role.ASSISTANT,
role="assistant",
contents=[
Content.from_function_call(
call_id=self.current_tool_call_id or "",
@@ -144,7 +142,7 @@ class AGUIEventConverter:
self.accumulated_tool_args += delta
return ChatResponseUpdate(
role=Role.ASSISTANT,
role="assistant",
contents=[
Content.from_function_call(
call_id=self.current_tool_call_id or "",
@@ -165,7 +163,7 @@ class AGUIEventConverter:
result = event.get("result") if event.get("result") is not None else event.get("content")
return ChatResponseUpdate(
role=Role.TOOL,
role="tool",
contents=[
Content.from_function_result(
call_id=tool_call_id,
@@ -177,8 +175,8 @@ class AGUIEventConverter:
def _handle_run_finished(self, event: dict[str, Any]) -> ChatResponseUpdate:
"""Handle RUN_FINISHED event."""
return ChatResponseUpdate(
role=Role.ASSISTANT,
finish_reason=FinishReason.STOP,
role="assistant",
finish_reason="stop",
contents=[],
additional_properties={
"thread_id": self.thread_id,
@@ -191,8 +189,8 @@ class AGUIEventConverter:
error_message = event.get("message", "Unknown error")
return ChatResponseUpdate(
role=Role.ASSISTANT,
finish_reason=FinishReason.CONTENT_FILTER,
role="assistant",
finish_reason="content_filter",
contents=[
Content.from_error(
message=error_message,
python/packages/ag-ui/agent_framework_ag_ui/_message_adapters.py
+12 -13
View File
@@ -9,7 +9,6 @@ from typing import Any, cast
from agent_framework import (
ChatMessage,
Content,
Role,
prepare_function_call_results,
)
@@ -269,7 +268,7 @@ def agui_messages_to_agent_framework(messages: list[dict[str, Any]]) -> list[Cha
def _find_matching_func_call(call_id: str) -> Content | None:
for prev_msg in result:
role_val = prev_msg.role.value if hasattr(prev_msg.role, "value") else str(prev_msg.role)
role_val = prev_msg.role if hasattr(prev_msg.role, "value") else str(prev_msg.role)
if role_val != "assistant":
continue
for content in prev_msg.contents or []:
@@ -287,7 +286,7 @@ def agui_messages_to_agent_framework(messages: list[dict[str, Any]]) -> list[Cha
return str(explicit_call_id)
for prev_msg in result:
role_val = prev_msg.role.value if hasattr(prev_msg.role, "value") else str(prev_msg.role)
role_val = prev_msg.role if hasattr(prev_msg.role, "value") else str(prev_msg.role)
if role_val != "assistant":
continue
direct_call = None
@@ -396,7 +395,7 @@ def agui_messages_to_agent_framework(messages: list[dict[str, Any]]) -> list[Cha
m
for m in result
if not (
(m.role.value if hasattr(m.role, "value") else str(m.role)) == "tool"
(m.role if hasattr(m.role, "value") else str(m.role)) == "tool"
and any(
c.type == "function_result" and c.call_id == approval_call_id
for c in (m.contents or [])
@@ -473,14 +472,14 @@ def agui_messages_to_agent_framework(messages: list[dict[str, Any]]) -> list[Cha
additional_properties={"ag_ui_state_args": state_args} if state_args else None,
)
chat_msg = ChatMessage(
role=Role.USER,
role="user",
contents=[approval_response],
)
else:
# No matching function call found - this is likely a confirm_changes approval
# Keep the old behavior for backwards compatibility
chat_msg = ChatMessage(
role=Role.USER,
role="user",
contents=[Content.from_text(text=approval_payload_text)],
additional_properties={"is_tool_result": True, "tool_call_id": str(tool_call_id or "")},
)
@@ -500,7 +499,7 @@ def agui_messages_to_agent_framework(messages: list[dict[str, Any]]) -> list[Cha
else:
func_result = str(result_content)
chat_msg = ChatMessage(
role=Role.TOOL,
role="tool",
contents=[Content.from_function_result(call_id=str(tool_call_id), result=func_result)],
)
if "id" in msg:
@@ -516,7 +515,7 @@ def agui_messages_to_agent_framework(messages: list[dict[str, Any]]) -> list[Cha
result_content = msg.get("result", msg.get("content", ""))
chat_msg = ChatMessage(
role=Role.TOOL,
role="tool",
contents=[Content.from_function_result(call_id=str(tool_call_id), result=result_content)],
)
if "id" in msg:
@@ -554,7 +553,7 @@ def agui_messages_to_agent_framework(messages: list[dict[str, Any]]) -> list[Cha
arguments=arguments,
)
)
chat_msg = ChatMessage(role=Role.ASSISTANT, contents=contents)
chat_msg = ChatMessage("assistant", contents)
if "id" in msg:
chat_msg.message_id = msg["id"]
result.append(chat_msg)
@@ -562,7 +561,7 @@ def agui_messages_to_agent_framework(messages: list[dict[str, Any]]) -> list[Cha
# No special handling required for assistant/plain messages here
role = AGUI_TO_FRAMEWORK_ROLE.get(role_str, Role.USER)
role = AGUI_TO_FRAMEWORK_ROLE.get(role_str, "user")
# Check if this message contains function approvals
if "function_approvals" in msg and msg["function_approvals"]:
@@ -584,14 +583,14 @@ def agui_messages_to_agent_framework(messages: list[dict[str, Any]]) -> list[Cha
)
approval_contents.append(approval_response)
chat_msg = ChatMessage(role=role, contents=approval_contents) # type: ignore[arg-type]
chat_msg = ChatMessage(role, approval_contents) # type: ignore[arg-type]
else:
# Regular text message
content = msg.get("content", "")
if isinstance(content, str):
chat_msg = ChatMessage(role=role, contents=[Content.from_text(text=content)])
chat_msg = ChatMessage(role, [Content.from_text(text=content)])
else:
chat_msg = ChatMessage(role=role, contents=[Content.from_text(text=str(content))])
chat_msg = ChatMessage(role, [Content.from_text(text=str(content))])
if "id" in msg:
chat_msg.message_id = msg["id"]
python/packages/ag-ui/agent_framework_ag_ui/_run.py
+1 -1
View File
@@ -862,7 +862,7 @@ async def run_agent_stream(
from pydantic import BaseModel
logger.info(f"Processing structured output, update count: {len(all_updates)}")
final_response = AgentResponse.from_agent_run_response_updates(all_updates, output_format_type=response_format)
final_response = AgentResponse.from_updates(all_updates, output_format_type=response_format)
if final_response.value and isinstance(final_response.value, BaseModel):
response_dict = final_response.value.model_dump(mode="json", exclude_none=True)
python/packages/ag-ui/agent_framework_ag_ui/_utils.py
+9 -9
View File
@@ -10,19 +10,19 @@ from dataclasses import asdict, is_dataclass
from datetime import date, datetime
from typing import Any
from agent_framework import AgentResponseUpdate, ChatResponseUpdate, FunctionTool, Role, ToolProtocol
from agent_framework import AgentResponseUpdate, ChatResponseUpdate, FunctionTool, ToolProtocol
# Role mapping constants
AGUI_TO_FRAMEWORK_ROLE: dict[str, Role] = {
"user": Role.USER,
"assistant": Role.ASSISTANT,
"system": Role.SYSTEM,
AGUI_TO_FRAMEWORK_ROLE: dict[str, str] = {
"user": "user",
"assistant": "assistant",
"system": "system",
}
FRAMEWORK_TO_AGUI_ROLE: dict[Role, str] = {
Role.USER: "user",
Role.ASSISTANT: "assistant",
Role.SYSTEM: "system",
FRAMEWORK_TO_AGUI_ROLE: dict[str, str] = {
"user": "user",
"assistant": "assistant",
"system": "system",
}
ALLOWED_AGUI_ROLES: set[str] = {"user", "assistant", "system", "tool"}
python/packages/ag-ui/getting_started/client_with_agent.py
+1 -1
View File
@@ -171,7 +171,7 @@ async def main():
messages = await thread.message_store.list_messages()
print(f"\n[THREAD STATE] {len(messages)} messages in thread's message_store")
for i, msg in enumerate(messages[-6:], 1): # Show last 6
role = msg.role.value if hasattr(msg.role, "value") else str(msg.role)
role = msg.role if hasattr(msg.role, "value") else str(msg.role)
text_preview = _preview_for_message(msg)
print(f" {i}. [{role}]: {text_preview}")
python/packages/ag-ui/tests/test_ag_ui_client.py
+11 -12
View File
@@ -12,7 +12,6 @@ from agent_framework import (
ChatResponse,
ChatResponseUpdate,
Content,
Role,
tool,
)
from pytest import MonkeyPatch
@@ -76,8 +75,8 @@ class TestAGUIChatClient:
"""Test state extraction when no state is present."""
client = TestableAGUIChatClient(endpoint="http://localhost:8888/")
messages = [
ChatMessage(role="user", text="Hello"),
ChatMessage(role="assistant", text="Hi there"),
ChatMessage("user", ["Hello"]),
ChatMessage("assistant", ["Hi there"]),
]
result_messages, state = client.extract_state_from_messages(messages)
@@ -96,7 +95,7 @@ class TestAGUIChatClient:
state_b64 = base64.b64encode(state_json.encode("utf-8")).decode("utf-8")
messages = [
ChatMessage(role="user", text="Hello"),
ChatMessage("user", ["Hello"]),
ChatMessage(
role="user",
contents=[Content.from_uri(uri=f"data:application/json;base64,{state_b64}")],
@@ -134,8 +133,8 @@ class TestAGUIChatClient:
"""Test message conversion to AG-UI format."""
client = TestableAGUIChatClient(endpoint="http://localhost:8888/")
messages = [
ChatMessage(role=Role.USER, text="What is the weather?"),
ChatMessage(role=Role.ASSISTANT, text="Let me check.", message_id="msg_123"),
ChatMessage("user", ["What is the weather?"]),
ChatMessage("assistant", ["Let me check."], message_id="msg_123"),
]
agui_messages = client.convert_messages_to_agui_format(messages)
@@ -182,7 +181,7 @@ class TestAGUIChatClient:
client = TestableAGUIChatClient(endpoint="http://localhost:8888/")
monkeypatch.setattr(client.http_service, "post_run", mock_post_run)
messages = [ChatMessage(role="user", text="Test message")]
messages = [ChatMessage("user", ["Test message"])]
chat_options = ChatOptions()
updates: list[ChatResponseUpdate] = []
@@ -215,7 +214,7 @@ class TestAGUIChatClient:
client = TestableAGUIChatClient(endpoint="http://localhost:8888/")
monkeypatch.setattr(client.http_service, "post_run", mock_post_run)
messages = [ChatMessage(role="user", text="Test message")]
messages = [ChatMessage("user", ["Test message"])]
chat_options = {}
response = await client.inner_get_response(messages=messages, options=chat_options)
@@ -258,7 +257,7 @@ class TestAGUIChatClient:
client = TestableAGUIChatClient(endpoint="http://localhost:8888/")
monkeypatch.setattr(client.http_service, "post_run", mock_post_run)
messages = [ChatMessage(role="user", text="Test with tools")]
messages = [ChatMessage("user", ["Test with tools"])]
chat_options = ChatOptions(tools=[test_tool])
response = await client.inner_get_response(messages=messages, options=chat_options)
@@ -282,7 +281,7 @@ class TestAGUIChatClient:
client = TestableAGUIChatClient(endpoint="http://localhost:8888/")
monkeypatch.setattr(client.http_service, "post_run", mock_post_run)
messages = [ChatMessage(role="user", text="Test server tool execution")]
messages = [ChatMessage("user", ["Test server tool execution"])]
updates: list[ChatResponseUpdate] = []
async for update in client.get_streaming_response(messages):
@@ -324,7 +323,7 @@ class TestAGUIChatClient:
client = TestableAGUIChatClient(endpoint="http://localhost:8888/")
monkeypatch.setattr(client.http_service, "post_run", mock_post_run)
messages = [ChatMessage(role="user", text="Test server tool execution")]
messages = [ChatMessage("user", ["Test server tool execution"])]
async for _ in client.get_streaming_response(messages, options={"tool_choice": "auto", "tools": [client_tool]}):
pass
@@ -338,7 +337,7 @@ class TestAGUIChatClient:
state_b64 = base64.b64encode(state_json.encode("utf-8")).decode("utf-8")
messages = [
ChatMessage(role="user", text="Hello"),
ChatMessage("user", ["Hello"]),
ChatMessage(
role="user",
contents=[Content.from_uri(uri=f"data:application/json;base64,{state_b64}")],
python/packages/ag-ui/tests/test_event_converters.py
+9 -11
View File
@@ -2,8 +2,6 @@
"""Tests for AG-UI event converter."""
from agent_framework import FinishReason, Role
from agent_framework_ag_ui._event_converters import AGUIEventConverter
@@ -22,7 +20,7 @@ class TestAGUIEventConverter:
update = converter.convert_event(event)
assert update is not None
assert update.role == Role.ASSISTANT
assert update.role == "assistant"
assert update.additional_properties["thread_id"] == "thread_123"
assert update.additional_properties["run_id"] == "run_456"
assert converter.thread_id == "thread_123"
@@ -39,7 +37,7 @@ class TestAGUIEventConverter:
update = converter.convert_event(event)
assert update is not None
assert update.role == Role.ASSISTANT
assert update.role == "assistant"
assert update.message_id == "msg_789"
assert converter.current_message_id == "msg_789"
@@ -55,7 +53,7 @@ class TestAGUIEventConverter:
update = converter.convert_event(event)
assert update is not None
assert update.role == Role.ASSISTANT
assert update.role == "assistant"
assert update.message_id == "msg_1"
assert len(update.contents) == 1
assert update.contents[0].text == "Hello"
@@ -101,7 +99,7 @@ class TestAGUIEventConverter:
update = converter.convert_event(event)
assert update is not None
assert update.role == Role.ASSISTANT
assert update.role == "assistant"
assert len(update.contents) == 1
assert update.contents[0].call_id == "call_123"
assert update.contents[0].name == "get_weather"
@@ -184,7 +182,7 @@ class TestAGUIEventConverter:
update = converter.convert_event(event)
assert update is not None
assert update.role == Role.TOOL
assert update.role == "tool"
assert len(update.contents) == 1
assert update.contents[0].call_id == "call_123"
assert update.contents[0].result == {"temperature": 22, "condition": "sunny"}
@@ -204,8 +202,8 @@ class TestAGUIEventConverter:
update = converter.convert_event(event)
assert update is not None
assert update.role == Role.ASSISTANT
assert update.finish_reason == FinishReason.STOP
assert update.role == "assistant"
assert update.finish_reason == "stop"
assert update.additional_properties["thread_id"] == "thread_123"
assert update.additional_properties["run_id"] == "run_456"
@@ -223,8 +221,8 @@ class TestAGUIEventConverter:
update = converter.convert_event(event)
assert update is not None
assert update.role == Role.ASSISTANT
assert update.finish_reason == FinishReason.CONTENT_FILTER
assert update.role == "assistant"
assert update.finish_reason == "content_filter"
assert len(update.contents) == 1
assert update.contents[0].message == "Connection timeout"
assert update.contents[0].error_code == "RUN_ERROR"
python/packages/ag-ui/tests/test_helpers.py
+5 -5
View File
@@ -29,8 +29,8 @@ class TestPendingToolCallIds:
def test_no_tool_calls(self):
"""Returns empty set when no tool calls in messages."""
messages = [
ChatMessage(role="user", contents=[Content.from_text("Hello")]),
ChatMessage(role="assistant", contents=[Content.from_text("Hi there")]),
ChatMessage("user", [Content.from_text("Hello")]),
ChatMessage("assistant", [Content.from_text("Hi there")]),
]
result = pending_tool_call_ids(messages)
assert result == set()
@@ -114,7 +114,7 @@ class TestIsStateContextMessage:
def test_empty_contents(self):
"""Returns False for message with empty contents."""
message = ChatMessage(role="system", contents=[])
message = ChatMessage("system", [])
assert is_state_context_message(message) is False
@@ -342,7 +342,7 @@ class TestLatestApprovalResponse:
def test_no_approval_response(self):
"""Returns None when no approval response in last message."""
messages = [
ChatMessage(role="assistant", contents=[Content.from_text("Hello")]),
ChatMessage("assistant", [Content.from_text("Hello")]),
]
result = latest_approval_response(messages)
assert result is None
@@ -357,7 +357,7 @@ class TestLatestApprovalResponse:
function_call=fc,
)
messages = [
ChatMessage(role="user", contents=[approval_content]),
ChatMessage("user", [approval_content]),
]
result = latest_approval_response(messages)
assert result is approval_content
python/packages/ag-ui/tests/test_message_adapters.py
+21 -21
View File
@@ -5,7 +5,7 @@
import json
import pytest
from agent_framework import ChatMessage, Content, Role
from agent_framework import ChatMessage, Content
from agent_framework_ag_ui._message_adapters import (
agent_framework_messages_to_agui,
@@ -24,7 +24,7 @@ def sample_agui_message():
@pytest.fixture
def sample_agent_framework_message():
"""Create a sample Agent Framework message."""
return ChatMessage(role=Role.USER, contents=[Content.from_text(text="Hello")], message_id="msg-123")
return ChatMessage("user", [Content.from_text(text="Hello")], message_id="msg-123")
def test_agui_to_agent_framework_basic(sample_agui_message):
@@ -32,7 +32,7 @@ def test_agui_to_agent_framework_basic(sample_agui_message):
messages = agui_messages_to_agent_framework([sample_agui_message])
assert len(messages) == 1
assert messages[0].role == Role.USER
assert messages[0].role == "user"
assert messages[0].message_id == "msg-123"
@@ -86,7 +86,7 @@ def test_agui_tool_result_to_agent_framework():
assert len(messages) == 1
message = messages[0]
assert message.role == Role.USER
assert message.role == "user"
assert len(message.contents) == 1
assert message.contents[0].type == "text"
@@ -328,9 +328,9 @@ def test_agui_multiple_messages_to_agent_framework():
messages = agui_messages_to_agent_framework(messages_input)
assert len(messages) == 3
assert messages[0].role == Role.USER
assert messages[1].role == Role.ASSISTANT
assert messages[2].role == Role.USER
assert messages[0].role == "user"
assert messages[1].role == "assistant"
assert messages[2].role == "user"
def test_agui_empty_messages():
@@ -366,7 +366,7 @@ def test_agui_function_approvals():
assert len(messages) == 1
msg = messages[0]
assert msg.role == Role.USER
assert msg.role == "user"
assert len(msg.contents) == 2
assert msg.contents[0].type == "function_approval_response"
@@ -385,7 +385,7 @@ def test_agui_system_role():
messages = agui_messages_to_agent_framework([{"role": "system", "content": "System prompt"}])
assert len(messages) == 1
assert messages[0].role == Role.SYSTEM
assert messages[0].role == "system"
def test_agui_non_string_content():
@@ -425,7 +425,7 @@ def test_agui_with_tool_calls_to_agent_framework():
assert len(messages) == 1
msg = messages[0]
assert msg.role == Role.ASSISTANT
assert msg.role == "assistant"
assert msg.message_id == "msg-789"
# First content is text, second is the function call
assert msg.contents[0].type == "text"
@@ -439,7 +439,7 @@ def test_agui_with_tool_calls_to_agent_framework():
def test_agent_framework_to_agui_with_tool_calls():
"""Test converting Agent Framework message with tool calls to AG-UI."""
msg = ChatMessage(
role=Role.ASSISTANT,
role="assistant",
contents=[
Content.from_text(text="Calling tool"),
Content.from_function_call(call_id="call-123", name="search", arguments={"query": "test"}),
@@ -464,7 +464,7 @@ def test_agent_framework_to_agui_with_tool_calls():
def test_agent_framework_to_agui_multiple_text_contents():
"""Test concatenating multiple text contents."""
msg = ChatMessage(
role=Role.ASSISTANT,
role="assistant",
contents=[Content.from_text(text="Part 1 "), Content.from_text(text="Part 2")],
)
@@ -476,7 +476,7 @@ def test_agent_framework_to_agui_multiple_text_contents():
def test_agent_framework_to_agui_no_message_id():
"""Test message without message_id - should auto-generate ID."""
msg = ChatMessage(role=Role.USER, contents=[Content.from_text(text="Hello")])
msg = ChatMessage("user", [Content.from_text(text="Hello")])
messages = agent_framework_messages_to_agui([msg])
@@ -488,7 +488,7 @@ def test_agent_framework_to_agui_no_message_id():
def test_agent_framework_to_agui_system_role():
"""Test system role conversion."""
msg = ChatMessage(role=Role.SYSTEM, contents=[Content.from_text(text="System")])
msg = ChatMessage("system", [Content.from_text(text="System")])
messages = agent_framework_messages_to_agui([msg])
@@ -534,7 +534,7 @@ def test_extract_text_from_custom_contents():
def test_agent_framework_to_agui_function_result_dict():
"""Test converting FunctionResultContent with dict result to AG-UI."""
msg = ChatMessage(
role=Role.TOOL,
role="tool",
contents=[Content.from_function_result(call_id="call-123", result={"key": "value", "count": 42})],
message_id="msg-789",
)
@@ -551,7 +551,7 @@ def test_agent_framework_to_agui_function_result_dict():
def test_agent_framework_to_agui_function_result_none():
"""Test converting FunctionResultContent with None result to AG-UI."""
msg = ChatMessage(
role=Role.TOOL,
role="tool",
contents=[Content.from_function_result(call_id="call-123", result=None)],
message_id="msg-789",
)
@@ -567,7 +567,7 @@ def test_agent_framework_to_agui_function_result_none():
def test_agent_framework_to_agui_function_result_string():
"""Test converting FunctionResultContent with string result to AG-UI."""
msg = ChatMessage(
role=Role.TOOL,
role="tool",
contents=[Content.from_function_result(call_id="call-123", result="plain text result")],
message_id="msg-789",
)
@@ -582,7 +582,7 @@ def test_agent_framework_to_agui_function_result_string():
def test_agent_framework_to_agui_function_result_empty_list():
"""Test converting FunctionResultContent with empty list result to AG-UI."""
msg = ChatMessage(
role=Role.TOOL,
role="tool",
contents=[Content.from_function_result(call_id="call-123", result=[])],
message_id="msg-789",
)
@@ -604,7 +604,7 @@ def test_agent_framework_to_agui_function_result_single_text_content():
text: str
msg = ChatMessage(
role=Role.TOOL,
role="tool",
contents=[Content.from_function_result(call_id="call-123", result=[MockTextContent("Hello from MCP!")])],
message_id="msg-789",
)
@@ -626,7 +626,7 @@ def test_agent_framework_to_agui_function_result_multiple_text_contents():
text: str
msg = ChatMessage(
role=Role.TOOL,
role="tool",
contents=[
Content.from_function_result(
call_id="call-123",
@@ -723,7 +723,7 @@ def test_agui_to_agent_framework_tool_result():
assert len(result) == 2
# Second message should be tool result
tool_msg = result[1]
assert tool_msg.role == Role.TOOL
assert tool_msg.role == "tool"
assert tool_msg.contents[0].type == "function_result"
assert tool_msg.contents[0].result == "Sunny"
python/packages/ag-ui/tests/test_message_hygiene.py
+1 -3
View File
@@ -25,9 +25,7 @@ def test_sanitize_tool_history_injects_confirm_changes_result() -> None:
sanitized = _sanitize_tool_history(messages)
tool_messages = [
msg for msg in sanitized if (msg.role.value if hasattr(msg.role, "value") else str(msg.role)) == "tool"
]
tool_messages = [msg for msg in sanitized if (msg.role if hasattr(msg.role, "value") else str(msg.role)) == "tool"]
assert len(tool_messages) == 1
assert str(tool_messages[0].contents[0].call_id) == "call_confirm_123"
assert tool_messages[0].contents[0].result == "Confirmed"
python/packages/ag-ui/tests/test_run.py
+11 -13
View File
@@ -188,7 +188,6 @@ class TestCreateStateContextMessage:
def test_creates_message(self):
"""Creates state context message."""
from agent_framework import Role
state = {"document": "Hello world"}
schema = {"properties": {"document": {"type": "string"}}}
@@ -196,7 +195,7 @@ class TestCreateStateContextMessage:
result = _create_state_context_message(state, schema)
assert result is not None
assert result.role == Role.SYSTEM
assert result.role == "system"
assert len(result.contents) == 1
assert "Hello world" in result.contents[0].text
assert "Current state" in result.contents[0].text
@@ -207,7 +206,7 @@ class TestInjectStateContext:
def test_no_state_message(self):
"""Returns original messages when no state context needed."""
messages = [ChatMessage(role="user", contents=[Content.from_text("Hello")])]
messages = [ChatMessage("user", [Content.from_text("Hello")])]
result = _inject_state_context(messages, {}, {})
assert result == messages
@@ -219,8 +218,8 @@ class TestInjectStateContext:
def test_last_message_not_user(self):
"""Returns original messages when last message is not from user."""
messages = [
ChatMessage(role="user", contents=[Content.from_text("Hello")]),
ChatMessage(role="assistant", contents=[Content.from_text("Hi")]),
ChatMessage("user", [Content.from_text("Hello")]),
ChatMessage("assistant", [Content.from_text("Hi")]),
]
state = {"key": "value"}
schema = {"properties": {"key": {"type": "string"}}}
@@ -230,11 +229,10 @@ class TestInjectStateContext:
def test_injects_before_last_user_message(self):
"""Injects state context before last user message."""
from agent_framework import Role
messages = [
ChatMessage(role="system", contents=[Content.from_text("You are helpful")]),
ChatMessage(role="user", contents=[Content.from_text("Hello")]),
ChatMessage("system", [Content.from_text("You are helpful")]),
ChatMessage("user", [Content.from_text("Hello")]),
]
state = {"document": "content"}
schema = {"properties": {"document": {"type": "string"}}}
@@ -243,13 +241,13 @@ class TestInjectStateContext:
assert len(result) == 3
# System message first
assert result[0].role == Role.SYSTEM
assert result[0].role == "system"
assert "helpful" in result[0].contents[0].text
# State context second
assert result[1].role == Role.SYSTEM
assert result[1].role == "system"
assert "Current state" in result[1].contents[0].text
# User message last
assert result[2].role == Role.USER
assert result[2].role == "user"
assert "Hello" in result[2].contents[0].text
@@ -357,7 +355,7 @@ def test_extract_approved_state_updates_no_handler():
"""Test _extract_approved_state_updates returns empty with no handler."""
from agent_framework_ag_ui._run import _extract_approved_state_updates
messages = [ChatMessage(role="user", contents=[Content.from_text("Hello")])]
messages = [ChatMessage("user", [Content.from_text("Hello")])]
result = _extract_approved_state_updates(messages, None)
assert result == {}
@@ -368,6 +366,6 @@ def test_extract_approved_state_updates_no_approval():
from agent_framework_ag_ui._run import _extract_approved_state_updates
handler = PredictiveStateHandler(predict_state_config={"doc": {"tool": "write", "tool_argument": "content"}})
messages = [ChatMessage(role="user", contents=[Content.from_text("Hello")])]
messages = [ChatMessage("user", [Content.from_text("Hello")])]
result = _extract_approved_state_updates(messages, handler)
assert result == {}
python/packages/ag-ui/tests/test_utils.py
+2 -2
View File
@@ -404,11 +404,11 @@ def test_safe_json_parse_with_none():
def test_get_role_value_with_enum():
"""Test get_role_value with enum role."""
from agent_framework import ChatMessage, Content, Role
from agent_framework import ChatMessage, Content
from agent_framework_ag_ui._utils import get_role_value
message = ChatMessage(role=Role.USER, contents=[Content.from_text("test")])
message = ChatMessage("user", [Content.from_text("test")])
result = get_role_value(message)
assert result == "user"
python/packages/ag-ui/tests/utils_test_ag_ui.py
+1 -1
View File
@@ -56,7 +56,7 @@ class StreamingChatClientStub(BaseChatClient[TOptions_co], Generic[TOptions_co])
contents.extend(update.contents)
return ChatResponse(
messages=[ChatMessage(role="assistant", contents=contents)],
messages=[ChatMessage("assistant", contents)],
response_id="stub-response",
)