Files
agent-framework/python/packages/ag-ui/tests/test_message_hygiene.py
T
Evan Mattson 8cf8b0f995 Python: Refactor ag-ui to clean up some patterns (#2363)
* Refactor ag-ui to clean up some patterns

* Mypy fixes

* Fix imports, typing, tests, logging.

* Fix test import error

* Fix imports again

* Fix thread handling
2025-11-27 02:13:03 +00:00

54 lines
1.6 KiB
Python

# Copyright (c) Microsoft. All rights reserved.
from agent_framework import ChatMessage, FunctionCallContent, FunctionResultContent, TextContent
from agent_framework_ag_ui._orchestration._message_hygiene import (
deduplicate_messages,
sanitize_tool_history,
)
def test_sanitize_tool_history_injects_confirm_changes_result() -> None:
messages = [
ChatMessage(
role="assistant",
contents=[
FunctionCallContent(
name="confirm_changes",
call_id="call_confirm_123",
arguments='{"changes": "test"}',
)
],
),
ChatMessage(
role="user",
contents=[TextContent(text='{"accepted": true}')],
),
]
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"
]
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"
def test_deduplicate_messages_prefers_non_empty_tool_results() -> None:
messages = [
ChatMessage(
role="tool",
contents=[FunctionResultContent(call_id="call1", result="")],
),
ChatMessage(
role="tool",
contents=[FunctionResultContent(call_id="call1", result="result data")],
),
]
deduped = deduplicate_messages(messages)
assert len(deduped) == 1
assert deduped[0].contents[0].result == "result data"