mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
Python: Fix AG-UI message handling and MCP tool double-call bug (#3635)
* AG-UI bug fixes * Fixes * Fixes * Revert human_in_the_loop_agent.py changes * Address copilot feedback * PR feedback addressed
This commit is contained in:
Evan Mattson
@@ -98,7 +98,14 @@ def test_agui_tool_result_to_agent_framework():
|
||||
|
||||
|
||||
def test_agui_tool_approval_updates_tool_call_arguments():
|
||||
"""Tool approval updates matching tool call arguments for snapshots and agent context."""
|
||||
"""Tool approval updates matching tool call arguments for snapshots and agent context.
|
||||
|
||||
The LLM context (ChatMessage) should contain only enabled steps, so the LLM
|
||||
generates responses based on what was actually approved/executed.
|
||||
|
||||
The raw messages (for MESSAGES_SNAPSHOT) should contain all steps with status,
|
||||
so the UI can show which steps were enabled/disabled.
|
||||
"""
|
||||
messages_input = [
|
||||
{
|
||||
"role": "assistant",
|
||||
@@ -142,13 +149,14 @@ def test_agui_tool_approval_updates_tool_call_arguments():
|
||||
assert len(messages) == 2
|
||||
assistant_msg = messages[0]
|
||||
func_call = next(content for content in assistant_msg.contents if content.type == "function_call")
|
||||
# LLM context should only have enabled steps (what was actually approved)
|
||||
assert func_call.arguments == {
|
||||
"steps": [
|
||||
{"description": "Boil water", "status": "enabled"},
|
||||
{"description": "Brew coffee", "status": "disabled"},
|
||||
{"description": "Serve coffee", "status": "enabled"},
|
||||
]
|
||||
}
|
||||
# Raw messages (for MESSAGES_SNAPSHOT) should have all steps with status
|
||||
assert messages_input[0]["tool_calls"][0]["function"]["arguments"] == {
|
||||
"steps": [
|
||||
{"description": "Boil water", "status": "enabled"},
|
||||
|
||||
@@ -5,7 +5,13 @@ from agent_framework import ChatMessage, Content
|
||||
from agent_framework_ag_ui._message_adapters import _deduplicate_messages, _sanitize_tool_history
|
||||
|
||||
|
||||
def test_sanitize_tool_history_injects_confirm_changes_result() -> None:
|
||||
def test_sanitize_tool_history_filters_out_confirm_changes_only_message() -> None:
|
||||
"""Test that assistant messages with ONLY confirm_changes are filtered out entirely.
|
||||
|
||||
When an assistant message contains only a confirm_changes tool call (no other tools),
|
||||
the entire message should be filtered out because confirm_changes is a synthetic
|
||||
tool for the approval UI flow that shouldn't be sent to the LLM.
|
||||
"""
|
||||
messages = [
|
||||
ChatMessage(
|
||||
role="assistant",
|
||||
@@ -25,10 +31,17 @@ 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 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"
|
||||
# Assistant message with only confirm_changes should be filtered out
|
||||
assistant_messages = [
|
||||
msg for msg in sanitized if (msg.role.value if hasattr(msg.role, "value") else str(msg.role)) == "assistant"
|
||||
]
|
||||
assert len(assistant_messages) == 0
|
||||
|
||||
# No synthetic tool result should be injected since confirm_changes was filtered out
|
||||
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) == 0
|
||||
|
||||
|
||||
def test_deduplicate_messages_prefers_non_empty_tool_results() -> None:
|
||||
@@ -46,3 +59,212 @@ def test_deduplicate_messages_prefers_non_empty_tool_results() -> None:
|
||||
deduped = _deduplicate_messages(messages)
|
||||
assert len(deduped) == 1
|
||||
assert deduped[0].contents[0].result == "result data"
|
||||
|
||||
|
||||
def test_convert_approval_results_to_tool_messages() -> None:
|
||||
"""Test that function_result content in user messages gets converted to tool messages.
|
||||
|
||||
This is a regression test for the MCP tool double-call bug where approved tool
|
||||
results ended up in user messages instead of tool messages, causing OpenAI to
|
||||
reject the request with 'tool_call_ids did not have response messages'.
|
||||
"""
|
||||
from agent_framework_ag_ui._run import _convert_approval_results_to_tool_messages
|
||||
|
||||
# Simulate what happens after _resolve_approval_responses:
|
||||
# A user message contains function_result content (the executed tool result)
|
||||
messages = [
|
||||
ChatMessage(
|
||||
role="assistant",
|
||||
contents=[
|
||||
Content.from_function_call(call_id="call_123", name="my_mcp_tool", arguments="{}"),
|
||||
],
|
||||
),
|
||||
ChatMessage(
|
||||
role="user",
|
||||
contents=[
|
||||
Content.from_function_result(call_id="call_123", result="tool execution result"),
|
||||
],
|
||||
),
|
||||
]
|
||||
|
||||
_convert_approval_results_to_tool_messages(messages)
|
||||
|
||||
# After conversion, the function result should be in a tool message, not user message
|
||||
assert len(messages) == 2
|
||||
|
||||
# First message unchanged
|
||||
assert messages[0].role == "assistant"
|
||||
|
||||
# Second message should now be role="tool"
|
||||
assert messages[1].role == "tool"
|
||||
assert messages[1].contents[0].type == "function_result"
|
||||
assert messages[1].contents[0].call_id == "call_123"
|
||||
|
||||
|
||||
def test_convert_approval_results_preserves_other_user_content() -> None:
|
||||
"""Test that user messages with mixed content are handled correctly.
|
||||
|
||||
If a user message has both function_result content and other content (like text),
|
||||
the function_result content should be extracted to a tool message while the
|
||||
remaining content stays in the user message.
|
||||
"""
|
||||
from agent_framework_ag_ui._run import _convert_approval_results_to_tool_messages
|
||||
|
||||
messages = [
|
||||
ChatMessage(
|
||||
role="assistant",
|
||||
contents=[
|
||||
Content.from_function_call(call_id="call_123", name="my_tool", arguments="{}"),
|
||||
],
|
||||
),
|
||||
ChatMessage(
|
||||
role="user",
|
||||
contents=[
|
||||
Content.from_text(text="User also said something"),
|
||||
Content.from_function_result(call_id="call_123", result="tool result"),
|
||||
],
|
||||
),
|
||||
]
|
||||
|
||||
_convert_approval_results_to_tool_messages(messages)
|
||||
|
||||
# Should have 3 messages now: assistant, tool (with result), user (with text)
|
||||
# OpenAI requires tool messages immediately after the assistant message with the tool call
|
||||
assert len(messages) == 3
|
||||
|
||||
# First message unchanged
|
||||
assert messages[0].role == "assistant"
|
||||
|
||||
# Second message should be tool with result (must come right after assistant per OpenAI requirements)
|
||||
assert messages[1].role == "tool"
|
||||
assert messages[1].contents[0].type == "function_result"
|
||||
|
||||
# Third message should be user with just text
|
||||
assert messages[2].role == "user"
|
||||
assert len(messages[2].contents) == 1
|
||||
assert messages[2].contents[0].type == "text"
|
||||
|
||||
|
||||
def test_sanitize_tool_history_filters_confirm_changes_keeps_other_tools() -> None:
|
||||
"""Test that confirm_changes is filtered but other tools are preserved.
|
||||
|
||||
When an assistant message contains both a real tool call and confirm_changes,
|
||||
confirm_changes should be filtered out while the real tool call is kept.
|
||||
No synthetic result is injected for confirm_changes since it's filtered.
|
||||
"""
|
||||
messages = [
|
||||
# User asks something
|
||||
ChatMessage(
|
||||
role="user",
|
||||
contents=[Content.from_text(text="What time is it?")],
|
||||
),
|
||||
# Assistant calls MCP tool + confirm_changes
|
||||
ChatMessage(
|
||||
role="assistant",
|
||||
contents=[
|
||||
Content.from_function_call(call_id="call_1", name="get_datetime", arguments="{}"),
|
||||
Content.from_function_call(call_id="call_c1", name="confirm_changes", arguments="{}"),
|
||||
],
|
||||
),
|
||||
# Tool result for the actual MCP tool
|
||||
ChatMessage(
|
||||
role="tool",
|
||||
contents=[Content.from_function_result(call_id="call_1", result="2024-01-01 12:00:00")],
|
||||
),
|
||||
# User asks something else
|
||||
ChatMessage(
|
||||
role="user",
|
||||
contents=[Content.from_text(text="What's the date?")],
|
||||
),
|
||||
]
|
||||
|
||||
sanitized = _sanitize_tool_history(messages)
|
||||
|
||||
# Find the assistant message
|
||||
assistant_messages = [
|
||||
msg for msg in sanitized if (msg.role.value if hasattr(msg.role, "value") else str(msg.role)) == "assistant"
|
||||
]
|
||||
assert len(assistant_messages) == 1
|
||||
|
||||
# Assistant message should only have get_datetime, not confirm_changes
|
||||
function_call_names = [c.name for c in assistant_messages[0].contents if c.type == "function_call"]
|
||||
assert "get_datetime" in function_call_names
|
||||
assert "confirm_changes" not in function_call_names
|
||||
|
||||
# Only one tool message (for call_1), no synthetic for confirm_changes
|
||||
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_1"
|
||||
|
||||
|
||||
def test_sanitize_tool_history_filters_confirm_changes_from_assistant_messages() -> None:
|
||||
"""Test that confirm_changes is removed from assistant messages sent to LLM.
|
||||
|
||||
This is a regression test for the human-in-the-loop bug where the LLM would see
|
||||
confirm_changes with function_arguments containing the original steps (e.g., 5 steps)
|
||||
even when the user only approved a subset (e.g., 2 steps), causing the LLM to
|
||||
respond with "Here's your 5-step plan" instead of "Here's your 2-step plan".
|
||||
"""
|
||||
messages = [
|
||||
ChatMessage(
|
||||
role="user",
|
||||
contents=[Content.from_text(text="Build a robot")],
|
||||
),
|
||||
# Assistant message with both generate_task_steps and confirm_changes
|
||||
ChatMessage(
|
||||
role="assistant",
|
||||
contents=[
|
||||
Content.from_function_call(
|
||||
call_id="call_1",
|
||||
name="generate_task_steps",
|
||||
arguments='{"steps": [{"description": "Step 1"}, {"description": "Step 2"}]}',
|
||||
),
|
||||
Content.from_function_call(
|
||||
call_id="call_c1",
|
||||
name="confirm_changes",
|
||||
arguments='{"function_arguments": {"steps": [{"description": "Step 1"}, {"description": "Step 2"}]}}',
|
||||
),
|
||||
],
|
||||
),
|
||||
# Approval response
|
||||
ChatMessage(
|
||||
role="user",
|
||||
contents=[
|
||||
Content.from_function_approval_response(
|
||||
approved=True,
|
||||
id="call_1",
|
||||
function_call=Content.from_function_call(
|
||||
call_id="call_1",
|
||||
name="generate_task_steps",
|
||||
arguments='{"steps": [{"description": "Step 1"}]}', # Only 1 step approved
|
||||
),
|
||||
),
|
||||
],
|
||||
),
|
||||
]
|
||||
|
||||
sanitized = _sanitize_tool_history(messages)
|
||||
|
||||
# Find the assistant message in sanitized output
|
||||
assistant_messages = [
|
||||
msg for msg in sanitized if (msg.role.value if hasattr(msg.role, "value") else str(msg.role)) == "assistant"
|
||||
]
|
||||
|
||||
assert len(assistant_messages) == 1
|
||||
|
||||
# The assistant message should NOT contain confirm_changes
|
||||
assistant_contents = assistant_messages[0].contents or []
|
||||
function_call_names = [c.name for c in assistant_contents if c.type == "function_call"]
|
||||
assert "generate_task_steps" in function_call_names
|
||||
assert "confirm_changes" not in function_call_names
|
||||
|
||||
# No synthetic tool result for confirm_changes (it was filtered from the message)
|
||||
tool_messages = [
|
||||
msg for msg in sanitized if (msg.role.value if hasattr(msg.role, "value") else str(msg.role)) == "tool"
|
||||
]
|
||||
# No tool results expected since there are no completed tool calls
|
||||
# (the approval response is handled separately by the framework)
|
||||
tool_call_ids = {str(msg.contents[0].call_id) for msg in tool_messages}
|
||||
assert "call_c1" not in tool_call_ids # No synthetic result for confirm_changes
|
||||
|
||||
@@ -2,12 +2,18 @@
|
||||
|
||||
"""Tests for _run.py helper functions and FlowState."""
|
||||
|
||||
from ag_ui.core import (
|
||||
TextMessageEndEvent,
|
||||
TextMessageStartEvent,
|
||||
)
|
||||
from agent_framework import ChatMessage, Content
|
||||
|
||||
from agent_framework_ag_ui._run import (
|
||||
FlowState,
|
||||
_build_safe_metadata,
|
||||
_create_state_context_message,
|
||||
_emit_content,
|
||||
_emit_tool_result,
|
||||
_has_only_tool_calls,
|
||||
_inject_state_context,
|
||||
_should_suppress_intermediate_snapshot,
|
||||
@@ -351,6 +357,50 @@ def test_emit_tool_call_generates_id():
|
||||
assert flow.tool_call_id is not None # ID should be generated
|
||||
|
||||
|
||||
def test_emit_tool_result_closes_open_message():
|
||||
"""Test _emit_tool_result emits TextMessageEndEvent for open text message.
|
||||
|
||||
This is a regression test for where TEXT_MESSAGE_END was not
|
||||
emitted when using MCP tools because the message_id was reset without
|
||||
closing the message first.
|
||||
"""
|
||||
flow = FlowState()
|
||||
# Simulate an open text message (e.g., from Feature #4 tool-only detection)
|
||||
flow.message_id = "open-msg-123"
|
||||
flow.tool_call_id = "call_456"
|
||||
|
||||
content = Content.from_function_result(call_id="call_456", result="tool result")
|
||||
|
||||
events = _emit_tool_result(content, flow, predictive_handler=None)
|
||||
|
||||
# Should have: ToolCallEndEvent, ToolCallResultEvent, TextMessageEndEvent
|
||||
assert len(events) == 3
|
||||
|
||||
# Verify TextMessageEndEvent is emitted for the open message
|
||||
text_end_events = [e for e in events if isinstance(e, TextMessageEndEvent)]
|
||||
assert len(text_end_events) == 1
|
||||
assert text_end_events[0].message_id == "open-msg-123"
|
||||
|
||||
# Verify message_id is reset after
|
||||
assert flow.message_id is None
|
||||
|
||||
|
||||
def test_emit_tool_result_no_open_message():
|
||||
"""Test _emit_tool_result works when there's no open text message."""
|
||||
flow = FlowState()
|
||||
# No open message
|
||||
flow.message_id = None
|
||||
flow.tool_call_id = "call_456"
|
||||
|
||||
content = Content.from_function_result(call_id="call_456", result="tool result")
|
||||
|
||||
events = _emit_tool_result(content, flow, predictive_handler=None)
|
||||
|
||||
# Should have: ToolCallEndEvent, ToolCallResultEvent (no TextMessageEndEvent)
|
||||
text_end_events = [e for e in events if isinstance(e, TextMessageEndEvent)]
|
||||
assert len(text_end_events) == 0
|
||||
|
||||
|
||||
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
|
||||
@@ -369,3 +419,268 @@ def test_extract_approved_state_updates_no_approval():
|
||||
messages = [ChatMessage("user", [Content.from_text("Hello")])]
|
||||
result = _extract_approved_state_updates(messages, handler)
|
||||
assert result == {}
|
||||
|
||||
|
||||
class TestBuildMessagesSnapshot:
|
||||
"""Tests for _build_messages_snapshot function."""
|
||||
|
||||
def test_tool_calls_and_text_are_separate_messages(self):
|
||||
"""Test that tool calls and text content are emitted as separate messages.
|
||||
|
||||
This is a regression test for issue #3619 where tool calls and content
|
||||
were incorrectly merged into a single assistant message.
|
||||
"""
|
||||
from agent_framework_ag_ui._run import FlowState, _build_messages_snapshot
|
||||
|
||||
flow = FlowState()
|
||||
flow.message_id = "msg-123"
|
||||
flow.pending_tool_calls = [
|
||||
{"id": "call_1", "function": {"name": "get_weather", "arguments": '{"city": "NYC"}'}},
|
||||
]
|
||||
flow.accumulated_text = "Here is the weather information."
|
||||
flow.tool_results = [{"id": "result-1", "role": "tool", "content": '{"temp": 72}', "toolCallId": "call_1"}]
|
||||
|
||||
result = _build_messages_snapshot(flow, [])
|
||||
|
||||
# Should have 3 messages: tool call msg, tool result, text content msg
|
||||
assert len(result.messages) == 3
|
||||
|
||||
# First message: assistant with tool calls only (no content)
|
||||
assistant_tool_msg = result.messages[0]
|
||||
assert assistant_tool_msg.role == "assistant"
|
||||
assert assistant_tool_msg.tool_calls is not None
|
||||
assert len(assistant_tool_msg.tool_calls) == 1
|
||||
assert assistant_tool_msg.content is None
|
||||
|
||||
# Second message: tool result
|
||||
tool_result_msg = result.messages[1]
|
||||
assert tool_result_msg.role == "tool"
|
||||
|
||||
# Third message: assistant with content only (no tool calls)
|
||||
assistant_text_msg = result.messages[2]
|
||||
assert assistant_text_msg.role == "assistant"
|
||||
assert assistant_text_msg.content == "Here is the weather information."
|
||||
assert assistant_text_msg.tool_calls is None
|
||||
|
||||
# The text message should have a different ID than the tool call message
|
||||
assert assistant_text_msg.id != assistant_tool_msg.id
|
||||
|
||||
def test_only_tool_calls_no_text(self):
|
||||
"""Test snapshot with only tool calls and no accumulated text."""
|
||||
from agent_framework_ag_ui._run import FlowState, _build_messages_snapshot
|
||||
|
||||
flow = FlowState()
|
||||
flow.message_id = "msg-123"
|
||||
flow.pending_tool_calls = [
|
||||
{"id": "call_1", "function": {"name": "get_weather", "arguments": "{}"}},
|
||||
]
|
||||
flow.accumulated_text = ""
|
||||
flow.tool_results = []
|
||||
|
||||
result = _build_messages_snapshot(flow, [])
|
||||
|
||||
# Should have 1 message: tool call msg only
|
||||
assert len(result.messages) == 1
|
||||
assert result.messages[0].role == "assistant"
|
||||
assert result.messages[0].tool_calls is not None
|
||||
assert result.messages[0].content is None
|
||||
|
||||
def test_only_text_no_tool_calls(self):
|
||||
"""Test snapshot with only text and no tool calls."""
|
||||
from agent_framework_ag_ui._run import FlowState, _build_messages_snapshot
|
||||
|
||||
flow = FlowState()
|
||||
flow.message_id = "msg-123"
|
||||
flow.pending_tool_calls = []
|
||||
flow.accumulated_text = "Hello world"
|
||||
flow.tool_results = []
|
||||
|
||||
result = _build_messages_snapshot(flow, [])
|
||||
|
||||
# Should have 1 message: text content msg only
|
||||
assert len(result.messages) == 1
|
||||
assert result.messages[0].role == "assistant"
|
||||
assert result.messages[0].content == "Hello world"
|
||||
assert result.messages[0].tool_calls is None
|
||||
# Should use the existing message_id
|
||||
assert result.messages[0].id == "msg-123"
|
||||
|
||||
def test_preserves_snapshot_messages(self):
|
||||
"""Test that existing snapshot messages are preserved."""
|
||||
from agent_framework_ag_ui._run import FlowState, _build_messages_snapshot
|
||||
|
||||
flow = FlowState()
|
||||
flow.pending_tool_calls = []
|
||||
flow.accumulated_text = ""
|
||||
|
||||
existing_messages = [
|
||||
{"id": "user-1", "role": "user", "content": "Hello"},
|
||||
{"id": "assist-1", "role": "assistant", "content": "Hi there"},
|
||||
]
|
||||
|
||||
result = _build_messages_snapshot(flow, existing_messages)
|
||||
|
||||
assert len(result.messages) == 2
|
||||
assert result.messages[0].id == "user-1"
|
||||
assert result.messages[1].id == "assist-1"
|
||||
|
||||
|
||||
def test_malformed_json_in_confirm_args_skips_confirmation():
|
||||
"""Test that malformed JSON in tool arguments skips confirm_changes flow.
|
||||
|
||||
This is a regression test to ensure that when tool arguments contain malformed
|
||||
JSON, the code skips the confirmation flow entirely rather than crashing or
|
||||
showing incomplete data to the user.
|
||||
"""
|
||||
import json
|
||||
|
||||
# Simulate the parsing logic - malformed JSON should trigger skip
|
||||
malformed_arguments = "{ invalid json }"
|
||||
tool_call = {"function": {"name": "write_doc", "arguments": malformed_arguments}}
|
||||
|
||||
# This is what the code should do - detect parsing failure and skip
|
||||
should_skip_confirmation = False
|
||||
try:
|
||||
json.loads(tool_call.get("function", {}).get("arguments", "{}"))
|
||||
except json.JSONDecodeError:
|
||||
should_skip_confirmation = True
|
||||
|
||||
# Should skip confirmation when JSON is malformed
|
||||
assert should_skip_confirmation is True
|
||||
|
||||
# Valid JSON should proceed with confirmation
|
||||
valid_arguments = '{"content": "hello"}'
|
||||
tool_call_valid = {"function": {"name": "write_doc", "arguments": valid_arguments}}
|
||||
should_skip_confirmation = False
|
||||
try:
|
||||
function_arguments = json.loads(tool_call_valid.get("function", {}).get("arguments", "{}"))
|
||||
except json.JSONDecodeError:
|
||||
should_skip_confirmation = True
|
||||
|
||||
assert should_skip_confirmation is False
|
||||
assert function_arguments == {"content": "hello"}
|
||||
|
||||
|
||||
class TestTextMessageEventBalancing:
|
||||
"""Tests for proper TEXT_MESSAGE_START/END event balancing.
|
||||
|
||||
These tests verify that the streaming flow produces balanced pairs of
|
||||
TextMessageStartEvent and TextMessageEndEvent, especially when tool
|
||||
execution is involved.
|
||||
"""
|
||||
|
||||
def test_tool_only_flow_produces_balanced_events(self):
|
||||
"""Test that a tool-only response produces balanced TEXT_MESSAGE events.
|
||||
|
||||
This simulates the scenario where the LLM immediately calls a tool
|
||||
without any initial text, then returns text after the tool result.
|
||||
"""
|
||||
flow = FlowState()
|
||||
all_events: list = []
|
||||
|
||||
# Step 1: LLM outputs function_call only (no text)
|
||||
func_call_content = Content.from_function_call(
|
||||
call_id="call_weather",
|
||||
name="get_weather",
|
||||
arguments='{"city": "Seattle"}',
|
||||
)
|
||||
|
||||
# Feature #4 check: this should trigger TextMessageStartEvent
|
||||
contents = [func_call_content]
|
||||
if not flow.message_id and _has_only_tool_calls(contents):
|
||||
flow.message_id = "tool-msg-1"
|
||||
all_events.append(TextMessageStartEvent(message_id=flow.message_id, role="assistant"))
|
||||
|
||||
# Emit tool call events
|
||||
all_events.extend(_emit_content(func_call_content, flow))
|
||||
|
||||
# Step 2: Tool executes and returns result
|
||||
func_result_content = Content.from_function_result(
|
||||
call_id="call_weather",
|
||||
result='{"temp": 55, "conditions": "rainy"}',
|
||||
)
|
||||
|
||||
# This should close the text message
|
||||
all_events.extend(_emit_tool_result(func_result_content, flow))
|
||||
|
||||
# Verify message_id was reset
|
||||
assert flow.message_id is None, "message_id should be reset after tool result"
|
||||
|
||||
# Step 3: LLM outputs text response
|
||||
text_content = Content.from_text("The weather in Seattle is 55°F and rainy.")
|
||||
|
||||
# Since message_id is None, _emit_text should create a new one
|
||||
for event in _emit_content(text_content, flow):
|
||||
all_events.append(event)
|
||||
|
||||
# Step 4: End of stream - emit final TextMessageEndEvent
|
||||
if flow.message_id:
|
||||
all_events.append(TextMessageEndEvent(message_id=flow.message_id))
|
||||
|
||||
# Verify event counts
|
||||
start_events = [e for e in all_events if isinstance(e, TextMessageStartEvent)]
|
||||
end_events = [e for e in all_events if isinstance(e, TextMessageEndEvent)]
|
||||
|
||||
# Should have 2 TextMessageStartEvent and 2 TextMessageEndEvent
|
||||
assert len(start_events) == 2, f"Expected 2 start events, got {len(start_events)}"
|
||||
assert len(end_events) == 2, f"Expected 2 end events, got {len(end_events)}"
|
||||
|
||||
# Verify order: first message should start and end before second starts
|
||||
# Find indices
|
||||
start_indices = [i for i, e in enumerate(all_events) if isinstance(e, TextMessageStartEvent)]
|
||||
end_indices = [i for i, e in enumerate(all_events) if isinstance(e, TextMessageEndEvent)]
|
||||
|
||||
# First end should come before second start
|
||||
assert end_indices[0] < start_indices[1], (
|
||||
f"First TextMessageEndEvent (index {end_indices[0]}) should come "
|
||||
f"before second TextMessageStartEvent (index {start_indices[1]})"
|
||||
)
|
||||
|
||||
def test_text_then_tool_flow(self):
|
||||
"""Test flow where LLM outputs text first, then calls a tool.
|
||||
|
||||
This simulates: "Let me check the weather..." -> tool call -> tool result -> "The weather is..."
|
||||
"""
|
||||
flow = FlowState()
|
||||
all_events: list = []
|
||||
|
||||
# Step 1: LLM outputs text first
|
||||
text1 = Content.from_text("Let me check the weather for you.")
|
||||
all_events.extend(_emit_content(text1, flow))
|
||||
|
||||
# Verify message_id is set
|
||||
assert flow.message_id is not None, "message_id should be set after text"
|
||||
first_msg_id = flow.message_id
|
||||
|
||||
# Step 2: LLM outputs function_call
|
||||
func_call = Content.from_function_call(
|
||||
call_id="call_1",
|
||||
name="get_weather",
|
||||
arguments="{}",
|
||||
)
|
||||
all_events.extend(_emit_content(func_call, flow))
|
||||
|
||||
# Step 3: Tool result comes back
|
||||
func_result = Content.from_function_result(call_id="call_1", result="sunny")
|
||||
all_events.extend(_emit_tool_result(func_result, flow))
|
||||
|
||||
# Verify message_id was reset and first message was closed
|
||||
assert flow.message_id is None
|
||||
end_events_so_far = [e for e in all_events if isinstance(e, TextMessageEndEvent)]
|
||||
assert len(end_events_so_far) == 1
|
||||
assert end_events_so_far[0].message_id == first_msg_id
|
||||
|
||||
# Step 4: LLM outputs follow-up text
|
||||
text2 = Content.from_text("The weather is sunny!")
|
||||
all_events.extend(_emit_content(text2, flow))
|
||||
|
||||
# Step 5: End of stream
|
||||
if flow.message_id:
|
||||
all_events.append(TextMessageEndEvent(message_id=flow.message_id))
|
||||
|
||||
# Verify balance
|
||||
start_events = [e for e in all_events if isinstance(e, TextMessageStartEvent)]
|
||||
end_events = [e for e in all_events if isinstance(e, TextMessageEndEvent)]
|
||||
|
||||
assert len(start_events) == 2
|
||||
assert len(end_events) == 2
|
||||
|
||||
Reference in New Issue
Block a user