import asyncio from typing import Any import pytest from llama_index.core.base.llms.types import ChatMessage, MessageRole from llama_index.core.base.llms.types import ( TextBlock as LITextBlock, ) from llama_index.core.schema import NodeWithScore, TextNode from pydantic import ValidationError from private_gpt.chat.input_models import MessageInput, ToolSpecBody from private_gpt.components.chunk.models import Chunk from private_gpt.components.engines.citations.utils import process_history_citations from private_gpt.events.models import ( ImageBlock, MidConvSystemBlock, SourceBlock, TextBlock, ThinkingBlock, TLDRBlock, ToolResultBlock, ToolUseBlock, ) def test_anthropic_message_default_role() -> None: with pytest.raises(ValidationError): MessageInput(content="test") @pytest.mark.parametrize( ("role", "content", "expected_role", "expected_content"), [ ("user", "Hello", MessageRole.USER, "Hello"), ("assistant", "Response", MessageRole.ASSISTANT, "Response"), ], ) def test_simple_message_conversion( role: str, content: str, expected_role: MessageRole, expected_content: str ) -> None: message = MessageInput(role=role, content=content) # type: ignore[arg-type] result, tool_uses = message._convert_into_llama_index_messages() assert len(result) == 1 assert result[0].role == expected_role assert result[0].content == expected_content assert tool_uses == {} def test_user_message_with_mid_conv_system_block() -> None: message = MessageInput( role="user", content=[ TextBlock(type="text", text="Hello"), MidConvSystemBlock( content=[ TextBlock(type="text", text="Be concise."), TextBlock(type="text", text="Use bullet points."), ] ), ], ) result, _ = message._convert_into_llama_index_messages() assert len(result) == 1 assert result[0].role == MessageRole.USER blocks = result[0].blocks assert len(blocks) == 2 assert isinstance(blocks[0], LITextBlock) assert blocks[0].text == "Hello" assert isinstance(blocks[1], LITextBlock) assert blocks[1].text == "Be concise.\nUse bullet points." def test_user_message_mid_conv_system_single_text_block() -> None: message = MessageInput( role="user", content=[ MidConvSystemBlock( content=[TextBlock(type="text", text="Focus on safety.")] ), ], ) result, _ = message._convert_into_llama_index_messages() assert len(result) == 1 blocks = result[0].blocks assert len(blocks) == 1 assert isinstance(blocks[0], LITextBlock) assert blocks[0].text == "Focus on safety." def test_assistant_message_with_text_block() -> None: text_block = TextBlock(type="text", text="Hello from assistant") message = MessageInput(role="assistant", content=[text_block]) result, tool_uses = message._convert_into_llama_index_messages() assert len(result) == 1 assert result[0].role == MessageRole.ASSISTANT assert result[0].content == "Hello from assistant" assert tool_uses == {} def test_assistant_message_with_tool_blocks() -> None: text_block = TextBlock(type="text", text="Using tools") tool_block1 = ToolUseBlock( type="tool_use", id="tool1", name="calculator", input={"operation": "add", "x": 1, "y": 2}, ) tool_block2 = ToolUseBlock( type="tool_use", id="tool2", name="search", input={"query": "weather"} ) message = MessageInput( role="assistant", content=[text_block, tool_block1, tool_block2] ) result, tool_uses = message._convert_into_llama_index_messages() assert len(result) == 2 # First message has text content and first tool assert result[0].role == MessageRole.ASSISTANT assert result[0].content == "Using tools" assert len(result[0].additional_kwargs["tool_calls"]) == 1 tool_call1 = result[0].additional_kwargs["tool_calls"][0] assert tool_call1.tool_id == "tool1" assert tool_call1.tool_name == "calculator" assert tool_call1.tool_kwargs == {"operation": "add", "x": 1, "y": 2} # Second message has no text content and second tool assert result[1].role == MessageRole.ASSISTANT assert result[1].content is None assert len(result[1].additional_kwargs["tool_calls"]) == 1 tool_call2 = result[1].additional_kwargs["tool_calls"][0] assert tool_call2.tool_id == "tool2" assert tool_call2.tool_name == "search" assert tool_call2.tool_kwargs == {"query": "weather"} # Check tool_uses mapping assert len(tool_uses) == 2 assert tool_uses["tool1"] == tool_block1 assert tool_uses["tool2"] == tool_block2 def test_tool_message_with_result_blocks() -> None: result_text = TextBlock(type="text", text="Result: 3") tool_result = ToolResultBlock( type="tool_result", tool_use_id="tool1", content=[result_text] ) message = MessageInput(role="assistant", content=[tool_result]) # Mock the tool_uses to simulate previous tool usage tool_uses = { "tool1": ToolUseBlock(type="tool_use", id="tool1", name="calculator", input={}) } result, _ = message._convert_into_llama_index_messages(tool_uses) assert len(result) == 1 assert result[0].role == MessageRole.TOOL assert result[0].content == "Result: 3" assert result[0].additional_kwargs["tool_call_id"] == "tool1" assert result[0].additional_kwargs["tool_call_name"] == "calculator" def test_tool_result_with_empty_content() -> None: tool_result = ToolResultBlock(type="tool_result", tool_use_id="tool1", content=[]) message = MessageInput(role="assistant", content=[tool_result]) tool_uses = { "tool1": ToolUseBlock(type="tool_use", id="tool1", name="calculator", input={}) } result, _ = message._convert_into_llama_index_messages(tool_uses) assert len(result) == 1 assert result[0].role == MessageRole.TOOL assert result[0].content == "No content" assert result[0].additional_kwargs["tool_call_id"] == "tool1" assert result[0].additional_kwargs["tool_call_name"] == "calculator" def test_extract_text_content_with_various_inputs() -> None: message = MessageInput(role="user", content="test") # Test with string blocks, custom_blocks = message._extract_content("Hello") assert blocks is not None text_block = next( (b for b in blocks if isinstance(b, LITextBlock)), None, ) assert text_block.text == "Hello" # Test with TextBlock in list text_block = TextBlock(type="text", text="Block text") blocks, custom_blocks = message._extract_content([text_block]) assert blocks is not None text_block = next( (b for b in blocks if isinstance(b, LITextBlock)), None, ) assert text_block.text == "Block text" # Test with non-text blocks tool_block = ToolUseBlock(type="tool_use", id="1", name="tool", input={}) blocks, custom_blocks = message._extract_content([tool_block]) assert blocks is None assert custom_blocks is not None assert len(custom_blocks.items()) == 1 # Test with None blocks, custom_blocks = message._extract_content(None) assert blocks is None assert custom_blocks is None def test_convert_from_llama_index_messages_with_reordering() -> None: # Create messages that need reordering messages = [ MessageInput(role="user", content="Calculate 2+2"), MessageInput( role="assistant", content=[ TextBlock(type="text", text="Let me calculate"), ToolUseBlock( type="tool_use", id="calc1", name="calculator", input={"operation": "add"}, ), ToolUseBlock( type="tool_use", id="calc2", name="calculator", input={"operation": "add"}, ), ToolResultBlock( type="tool_result", tool_use_id="calc1", content=[TextBlock(type="text", text="4")], ), ToolResultBlock( type="tool_result", tool_use_id="calc2", content=[TextBlock(type="text", text="4")], ), ], ), ] result = MessageInput.convert_from_llama_index_messages(messages) # Check that tool response is moved after its assistant message assert result[0].role == MessageRole.USER assert result[1].role == MessageRole.ASSISTANT assert result[2].role == MessageRole.TOOL # Moved from position 3 assert result[3].role == MessageRole.ASSISTANT assert result[4].role == MessageRole.TOOL # Verify tool connection assert result[1].additional_kwargs["tool_calls"][0].tool_id == "calc1" assert result[2].additional_kwargs["tool_call_id"] == "calc1" @pytest.mark.parametrize( ("name", "description", "input_schema"), [ ("calculator", "Performs calculations", {"type": "object", "properties": {}}), ("search", None, {"type": "object", "required": ["query"]}), ( "data_fetcher", "Fetches data", {"type": "object", "properties": {"_id": {"type": "string"}}}, ), ], ) def test_anthropic_tool_creation( name: str, description: str | None, input_schema: dict[str, Any] ) -> None: tool = ToolSpecBody(name=name, description=description, input_schema=input_schema) assert tool.name == name assert tool.description == description assert tool.input_schema == input_schema def test_empty_content_handling() -> None: with pytest.raises(ValidationError): MessageInput(role="assistant", content=None) def test_complex_tool_workflow() -> None: # Simulate a complete workflow with tools messages = [ MessageInput(role="user", content="What's the weather?"), MessageInput( role="assistant", content=[ TextBlock(type="text", text="I'll check the weather"), ToolUseBlock( type="tool_use", id="weather1", name="weather_api", input={"location": "New York"}, ), ], ), MessageInput( role="assistant", content=[ ToolResultBlock( type="tool_result", tool_use_id="weather1", content=[TextBlock(type="text", text="Sunny, 72°F")], ) ], ), MessageInput( role="assistant", content="The weather in New York is sunny and 72°F." ), ] result = MessageInput.convert_from_llama_index_messages(messages) # Verify the flow assert len(result) == 4 assert result[0].role == MessageRole.USER assert result[1].role == MessageRole.ASSISTANT assert result[1].additional_kwargs["tool_calls"][0].tool_name == "weather_api" assert result[2].role == MessageRole.TOOL assert result[2].content == "Sunny, 72°F" assert result[3].role == MessageRole.ASSISTANT assert result[3].content == "The weather in New York is sunny and 72°F." def test_multiple_tools_same_message() -> None: messages = [ MessageInput( role="user", content="What's the weather in NYC and search for restaurants there?", ), MessageInput( role="assistant", content=[ TextBlock( type="text", text="I'll check the weather and find restaurants" ), ToolUseBlock( type="tool_use", id="weather1", name="weather_api", input={"location": "New York"}, ), ToolUseBlock( type="tool_use", id="search1", name="search_api", input={"query": "restaurants in New York", "type": "local"}, ), ToolUseBlock( type="tool_use", id="maps1", name="maps_api", input={"location": "New York", "category": "restaurants"}, ), ], ), MessageInput( role="assistant", content=[ ToolResultBlock( type="tool_result", tool_use_id="weather1", content=[TextBlock(type="text", text="Sunny, 72°F")], ) ], ), MessageInput( role="assistant", content=[ ToolResultBlock( type="tool_result", tool_use_id="search1", content=[ TextBlock( type="text", text="Found 50 restaurants: Italian, Chinese, Mexican...", ) ], ) ], ), MessageInput( role="assistant", content=[ ToolResultBlock( type="tool_result", tool_use_id="maps1", content=[ TextBlock( type="text", text="Top rated: Joe's Pizza (4.5★), Dragon Palace (4.7★)", ) ], ) ], ), MessageInput( role="assistant", content="The weather in NYC is sunny and 72°F. I found many restaurants including Joe's Pizza and Dragon Palace which are highly rated.", ), ] result = MessageInput.convert_from_llama_index_messages(messages) # Should create 3 separate assistant messages for each tool assert len(result) == 8 # user + 3*(assistant+tool) + final assistant # Check user message assert result[0].role == MessageRole.USER # Check first tool (weather) assert result[1].role == MessageRole.ASSISTANT assert result[1].content == "I'll check the weather and find restaurants" assert result[1].additional_kwargs["tool_calls"][0].tool_name == "weather_api" assert result[2].role == MessageRole.TOOL assert result[2].additional_kwargs["tool_call_id"] == "weather1" # Check second tool (search) assert result[3].role == MessageRole.ASSISTANT assert result[3].content is None # Only first tool has text content assert result[3].additional_kwargs["tool_calls"][0].tool_name == "search_api" assert result[4].role == MessageRole.TOOL assert result[4].additional_kwargs["tool_call_id"] == "search1" # Check third tool (maps) assert result[5].role == MessageRole.ASSISTANT assert result[5].content is None assert result[5].additional_kwargs["tool_calls"][0].tool_name == "maps_api" assert result[6].role == MessageRole.TOOL assert result[6].additional_kwargs["tool_call_id"] == "maps1" # Check final assistant response assert result[7].role == MessageRole.ASSISTANT assert "sunny and 72°F" in result[7].content assert "Joe's Pizza" in result[7].content def test_multiple_tools_with_followups() -> None: messages = [ MessageInput(role="user", content="Plan my trip to Tokyo"), MessageInput( role="assistant", content=[ TextBlock( type="text", text="I'll help plan your trip. Let me check flights and weather", ), ToolUseBlock( type="tool_use", id="flight1", name="flight_search", input={"destination": "Tokyo", "dates": "flexible"}, ), ToolUseBlock( type="tool_use", id="weather1", name="weather_api", input={"location": "Tokyo", "forecast": "7days"}, ), ], ), MessageInput( role="assistant", content=[ ToolResultBlock( type="tool_result", tool_use_id="flight1", content=[ TextBlock( type="text", text="Best flights: $800-1200, 12-14 hours" ) ], ) ], ), MessageInput( role="assistant", content=[ ToolResultBlock( type="tool_result", tool_use_id="weather1", content=[ TextBlock(type="text", text="Next week: Partly cloudy, 18-22°C") ], ) ], ), # Second round - split into two messages again MessageInput( role="assistant", content=[ TextBlock( type="text", text="Great! Flights are $800-1200. Weather looks good. Let me find hotels and attractions", ), ToolUseBlock( type="tool_use", id="hotel1", name="hotel_search", input={"city": "Tokyo", "budget": "moderate"}, ), ], ), MessageInput( role="assistant", content=[ ToolUseBlock( type="tool_use", id="attraction1", name="attraction_search", input={"city": "Tokyo", "type": "popular"}, ) ], ), # Tool results can be combined in a single message MessageInput( role="user", content=[ ToolResultBlock( type="tool_result", tool_use_id="hotel1", content=[ TextBlock( type="text", text="Hotels: Shinjuku Grand ($150/night), Tokyo Bay ($120/night)", ), SourceBlock( type="source", sources=[ Chunk.from_node( NodeWithScore(node=TextNode(text="Shinjuku Grand")) ), Chunk.from_node( NodeWithScore(node=TextNode(text="Tokyo Bay")) ), ], ), ], ), ToolResultBlock( type="tool_result", tool_use_id="attraction1", content=[ TextBlock( type="text", text="Must see: Senso-ji Temple, Tokyo Skytree, Shibuya Crossing", ), SourceBlock( type="source", sources=[ Chunk.from_node( NodeWithScore(node=TextNode(text="Shinjuku Grand")) ), Chunk.from_node( NodeWithScore(node=TextNode(text="Tokyo Bay")) ), ], ), ], ), ], ), MessageInput( role="assistant", content="Your Tokyo trip is planned! Flights: $800-1200, Weather: 18-22°C partly cloudy." "Hotels: Shinjuku Grand or Tokyo Bay. Must-see: Senso-ji Temple, Tokyo Skytree, and Shibuya Crossing.", ), ] result = MessageInput.convert_from_llama_index_messages(messages) # Count messages: user + 2 pairs (assistant+tool) # + 2 pairs (assistant+tool) + final assistant assert len(result) == 10 # After reordering: user, assistant1, tool1, assistant2, tool2, # assistant3, tool3, assistant4, tool4, final assistant assert result[0].role == MessageRole.USER # First assistant+tool pair assert result[1].role == MessageRole.ASSISTANT assert ( result[1].content == "I'll help plan your trip. Let me check flights and weather" ) assert result[1].additional_kwargs["tool_calls"][0].tool_name == "flight_search" assert result[2].role == MessageRole.TOOL assert result[2].additional_kwargs["tool_call_id"] == "flight1" # Second assistant+tool pair assert result[3].role == MessageRole.ASSISTANT assert result[3].content is None assert result[3].additional_kwargs["tool_calls"][0].tool_name == "weather_api" assert result[4].role == MessageRole.TOOL assert result[4].additional_kwargs["tool_call_id"] == "weather1" # Third assistant+tool pair assert result[5].role == MessageRole.ASSISTANT assert "Flights are $800-1200" in result[5].content assert result[5].additional_kwargs["tool_calls"][0].tool_name == "hotel_search" assert result[6].role == MessageRole.TOOL assert result[6].additional_kwargs["tool_call_id"] == "hotel1" # Fourth assistant+tool pair assert result[7].role == MessageRole.ASSISTANT assert result[7].content is None assert result[7].additional_kwargs["tool_calls"][0].tool_name == "attraction_search" assert result[8].role == MessageRole.TOOL assert result[8].additional_kwargs["tool_call_id"] == "attraction1" # Check that we can extract sources from the history _, sources, _ = asyncio.run(process_history_citations(result)) assert len(sources) == 4 # Final response assert result[9].role == MessageRole.ASSISTANT assert "Tokyo trip is planned" in result[9].content # Verify all tool calls have unique IDs tool_ids = set() for msg in result: if msg.role == MessageRole.ASSISTANT and "tool_calls" in msg.additional_kwargs: tool_id = msg.additional_kwargs["tool_calls"][0].tool_id assert tool_id not in tool_ids, f"Duplicate tool ID: {tool_id}" tool_ids.add(tool_id) assert tool_ids == {"flight1", "weather1", "hotel1", "attraction1"} def test_tool_results_multiple_in_single_message() -> None: # Test when tool results are combined in a single tool message messages = [ MessageInput(role="user", content="Find me a restaurant and check the weather"), MessageInput( role="assistant", content=[ TextBlock( type="text", text="I'll search for restaurants and check weather" ), ToolUseBlock( type="tool_use", id="rest1", name="restaurant_search", input={"cuisine": "Italian", "price": "moderate"}, ), ], ), MessageInput( role="assistant", content=[ ToolUseBlock( type="tool_use", id="weather2", name="weather_api", input={"location": "current"}, ) ], ), # Both tool results in a single message MessageInput( role="assistant", content=[ ToolResultBlock( type="tool_result", tool_use_id="rest1", content=[ TextBlock( type="text", text="Found: Luigi's Italian (4.5★), Pasta Palace (4.2★)", ) ], ), ToolResultBlock( type="tool_result", tool_use_id="weather2", content=[TextBlock(type="text", text="Sunny, 68°F")], ), ], ), MessageInput( role="assistant", content="Great options! Luigi's Italian has 4.5 stars. The weather is sunny and 68°F - perfect for dining out!", ), ] result = MessageInput.convert_from_llama_index_messages(messages) # After reordering: user, assistant1, tool1, assistant2, tool2, final assistant assert len(result) == 6 # User message assert result[0].role == MessageRole.USER # First assistant+tool pair (restaurant search) assert result[1].role == MessageRole.ASSISTANT assert result[1].additional_kwargs["tool_calls"][0].tool_name == "restaurant_search" assert result[2].role == MessageRole.TOOL assert result[2].additional_kwargs["tool_call_id"] == "rest1" assert "Luigi's Italian" in result[2].content # Second assistant+tool pair (weather) assert result[3].role == MessageRole.ASSISTANT assert result[3].additional_kwargs["tool_calls"][0].tool_name == "weather_api" assert result[4].role == MessageRole.TOOL assert result[4].additional_kwargs["tool_call_id"] == "weather2" assert "Sunny, 68°F" in result[4].content # Final assistant response assert result[5].role == MessageRole.ASSISTANT assert "Luigi's Italian" in result[5].content def test_assistant_message_only_tools_no_text() -> None: """Test assistant message with only tool blocks, no text content.""" tool_block1 = ToolUseBlock( type="tool_use", id="tool1", name="calculator", input={"operation": "add", "x": 1, "y": 2}, ) tool_block2 = ToolUseBlock( type="tool_use", id="tool2", name="search", input={"query": "weather"} ) message = MessageInput(role="assistant", content=[tool_block1, tool_block2]) result, tool_uses = message._convert_into_llama_index_messages() assert len(result) == 2 # First message has no text content, just first tool assert result[0].role == MessageRole.ASSISTANT assert result[0].content is None assert len(result[0].additional_kwargs["tool_calls"]) == 1 assert result[0].additional_kwargs["tool_calls"][0].tool_id == "tool1" # Second message has no text content, just second tool assert result[1].role == MessageRole.ASSISTANT assert result[1].content is None assert len(result[1].additional_kwargs["tool_calls"]) == 1 assert result[1].additional_kwargs["tool_calls"][0].tool_id == "tool2" assert len(tool_uses) == 2 def test_mixed_content_blocks_with_tools() -> None: """Test assistant message with mixed content: text, tool, result.""" text_block1 = TextBlock(type="text", text="First text") tool_block1 = ToolUseBlock(type="tool_use", id="tool1", name="calc", input={"x": 1}) tool_result1 = ToolResultBlock( type="tool_result", tool_use_id="tool1", content=[TextBlock(type="text", text="2")], ) message = MessageInput( role="assistant", content=[text_block1, tool_block1, tool_result1] ) result, _ = message._convert_into_llama_index_messages() assert len(result) == 2 # First message: first text + first tool assert result[0].role == MessageRole.ASSISTANT assert result[0].content == "First text" assert result[0].additional_kwargs["tool_calls"][0].tool_id == "tool1" # Second message: second message is a tool result assert result[1].role == MessageRole.TOOL assert result[1].additional_kwargs["tool_call_id"] == "tool1" def test_tool_result_without_prior_tool_use() -> None: """Test tool result block when tool_use is not in the current context.""" tool_result = ToolResultBlock( type="tool_result", tool_use_id="unknown_tool", content=[TextBlock(type="text", text="Some result")], ) message = MessageInput(role="assistant", content=[tool_result]) result, _ = message._convert_into_llama_index_messages({}) assert len(result) == 1 assert result[0].role == MessageRole.TOOL assert result[0].additional_kwargs["tool_call_id"] == "unknown_tool" assert result[0].additional_kwargs["tool_call_name"] is None # No prior tool use def test_empty_tool_use_block() -> None: """Test handling of empty/None tool use blocks.""" text_block = TextBlock(type="text", text="Test") # Invalid tool block is rejected by strict pydantic constraints. with pytest.raises(ValidationError): ToolUseBlock(type="tool_use", id="", name="", input={}) message = MessageInput(role="assistant", content=[text_block]) result, _ = message._convert_into_llama_index_messages() assert len(result) == 1 assert result[0].role == MessageRole.ASSISTANT def test_tool_result_with_multiple_content_blocks() -> None: """Test tool result with various content block types.""" from llama_index.core.schema import NodeWithScore, TextNode from private_gpt.events.models import SourceBlock text_content = TextBlock(type="text", text="Result text") source_content = SourceBlock( type="source", sources=[ Chunk.from_node(NodeWithScore(node=TextNode(text="Source 1"))), Chunk.from_node(NodeWithScore(node=TextNode(text="Source 2"))), ], ) tool_result = ToolResultBlock( type="tool_result", tool_use_id="tool1", content=[text_content, source_content] ) message = MessageInput(role="assistant", content=[tool_result]) # Mock tool_uses tool_uses = { "tool1": ToolUseBlock(type="tool_use", id="tool1", name="search", input={}) } result, _ = message._convert_into_llama_index_messages(tool_uses) assert len(result) == 1 assert result[0].role == MessageRole.TOOL assert result[0].content == "Result text" # Should have source blocks in additional_kwargs assert "source" in result[0].additional_kwargs def test_large_number_of_tools() -> None: """Test assistant message with many tool blocks.""" text_block = TextBlock(type="text", text="Using many tools") tool_blocks = [ ToolUseBlock( type="tool_use", id=f"tool{i}", name=f"tool_{i}", input={"param": i} ) for i in range(5) ] message = MessageInput(role="assistant", content=[text_block, *tool_blocks]) result, tool_uses = message._convert_into_llama_index_messages() # Should create 5 separate assistant messages assert len(result) == 5 # First message has text content assert result[0].content == "Using many tools" assert result[0].additional_kwargs["tool_calls"][0].tool_id == "tool0" # Subsequent messages have no text content for i in range(1, 5): assert result[i].content is None assert result[i].additional_kwargs["tool_calls"][0].tool_id == f"tool{i}" # All tools should be in tool_uses assert len(tool_uses) == 5 for i in range(5): assert f"tool{i}" in tool_uses def test_convert_from_llama_index_messages_preserves_order() -> None: """Test that the overall message order is preserved correctly.""" messages = [ MessageInput(role="user", content="Request 1"), MessageInput(role="assistant", content="Response 1"), MessageInput(role="user", content="Request 2"), MessageInput( role="assistant", content=[ TextBlock(type="text", text="Using tool"), ToolUseBlock(type="tool_use", id="t1", name="tool", input={}), ], ), MessageInput( role="assistant", content=[ ToolResultBlock( type="tool_result", tool_use_id="t1", content=[TextBlock(type="text", text="Tool result")], ), ], ), MessageInput(role="assistant", content="Final response"), MessageInput(role="user", content="Request 4"), MessageInput( role="assistant", content=[ ToolResultBlock( type="tool_result", tool_use_id="t2", content=[TextBlock(type="text", text="Tool result")], ), ToolUseBlock(type="tool_use", id="t2", name="tool", input={}), TextBlock(type="text", text="Using tool"), ], ), MessageInput(role="user", content="Request 5"), MessageInput( role="assistant", content=[ TextBlock(type="text", text="Using tool"), ], ), MessageInput(role="user", content="Request 6"), ] result = MessageInput.convert_from_llama_index_messages(messages) # Verify the sequence: user, assistant, user, assistant, tool, user, assistant expected_roles = [ MessageRole.USER, # Request 1 MessageRole.ASSISTANT, # Response 1 MessageRole.USER, # Request 2 MessageRole.ASSISTANT, # Using tool MessageRole.TOOL, # Tool result (reordered) MessageRole.ASSISTANT, # Final response MessageRole.USER, # Request 3 MessageRole.ASSISTANT, # Final response MessageRole.TOOL, MessageRole.ASSISTANT, # Using tool (reordered) MessageRole.USER, # Request 4 MessageRole.ASSISTANT, # Using tool (reordered) MessageRole.USER, # Request 5 ] assert len(result) == len(expected_roles) for i, expected_role in enumerate(expected_roles): assert result[i].role == expected_role, ( f"Position {i}: expected {expected_role}, got {result[i].role}" ) @pytest.mark.parametrize("invalid_role", ["tool", "invalid", "ASSISTANT", "USER"]) def test_invalid_message_roles(invalid_role: str) -> None: """Test handling of invalid message roles.""" with pytest.raises(ValidationError): MessageInput(role=invalid_role, content="test") # type: ignore def test_image_block_handling() -> None: """Test handling of image blocks in content.""" data = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg==" mime_type = "image/png" image_block = ImageBlock.from_base64(data=data, mime_type=mime_type) text_block = TextBlock(type="text", text="Check this image") message = MessageInput(role="user", content=[text_block, image_block]) result, _ = message._convert_into_llama_index_messages() assert len(result) == 1 assert result[0].role == MessageRole.USER content_blocks = result[0].content if isinstance(content_blocks, list): text_blocks = [b for b in content_blocks if isinstance(b, LITextBlock)] image_blocks = [ b for b in content_blocks if isinstance(b, type(content_blocks[0])) and hasattr(b, "image") ] assert len(text_blocks) >= 1 assert len(image_blocks) >= 1 def test_convert_multiple_tool_uses_with_interleaved_results() -> None: """Test assistant message with multiple tool uses and their results interleaved.""" messages = [ MessageInput(role="user", content="Get weather for Madrid and Paris"), MessageInput( role="assistant", content=[ TextBlock(type="text", text="Using tool1"), ToolUseBlock( type="tool_use", id="toolu_madrid", name="get_weather", input={"location": "Madrid"}, ), ToolResultBlock( type="tool_result", tool_use_id="toolu_madrid", content=[TextBlock(type="text", text="Madrid: -20°C")], is_error=False, ), TextBlock(type="text", text="Using tool2"), ToolUseBlock( type="tool_use", id="toolu_paris", name="get_weather", input={"location": "Paris"}, ), ToolResultBlock( type="tool_result", tool_use_id="toolu_paris", content=[TextBlock(type="text", text="Paris: -20°C")], is_error=False, ), TextBlock(type="text", text="Weather retrieved for both cities."), ], ), MessageInput(role="user", content="Send email with this info"), ] result = MessageInput.convert_from_llama_index_messages(messages) assert len(result) == 7 # user, assistant, tool1, tool2, user assert result[0].role == MessageRole.USER assert result[1].role == MessageRole.ASSISTANT assert result[2].role == MessageRole.TOOL assert result[3].role == MessageRole.ASSISTANT assert result[4].role == MessageRole.TOOL assert result[5].role == MessageRole.ASSISTANT assert result[6].role == MessageRole.USER # Verify tool calls are properly connected tool_calls = result[1].additional_kwargs["tool_calls"] assert len(tool_calls) == 1 assert tool_calls[0].tool_id == "toolu_madrid" tool_calls = result[3].additional_kwargs["tool_calls"] assert tool_calls[0].tool_id == "toolu_paris" # Verify tool results maintain proper IDs assert result[2].additional_kwargs["tool_call_id"] == "toolu_madrid" assert result[4].additional_kwargs["tool_call_id"] == "toolu_paris" # Verify text blocks are in the right places assert result[1].content == "Using tool1" assert result[3].content == "Using tool2" assert result[5].content == "Weather retrieved for both cities." def test_convert_assistant_with_text_and_tool_results_only() -> None: """Test assistant message containing only tool results and text (no tool uses).""" messages = [ MessageInput(role="user", content="Previous request"), MessageInput(role="assistant", content="I'll help with that"), MessageInput( role="assistant", content=[ ToolUseBlock( type="tool_use", id="previous_tool_id", name="get_weather", input={"location": "InvalidPlace"}, ), ToolResultBlock( type="tool_result", tool_use_id="previous_tool_id", content=[TextBlock(type="text", text="Result data")], is_error=False, ), TextBlock(type="text", text="Based on the results..."), ], ), ] result = MessageInput.convert_from_llama_index_messages(messages) assert len(result) == 4 assert result[0].role == MessageRole.USER assert result[1].role == MessageRole.ASSISTANT assert result[2].role == MessageRole.TOOL assert result[3].role == MessageRole.ASSISTANT def test_convert_orphaned_tool_results() -> None: """Test handling of tool results without corresponding tool uses.""" messages = [ MessageInput(role="user", content="Test orphaned results"), MessageInput( role="assistant", content=[ ToolResultBlock( type="tool_result", tool_use_id="missing_tool_id", content=[TextBlock(type="text", text="Orphaned result")], is_error=False, ), TextBlock(type="text", text="Some response text"), ], ), ] result = MessageInput.convert_from_llama_index_messages(messages) assert len(result) == 2 assert len(result[0].additional_kwargs) == 0 assert len(result[1].additional_kwargs) == 0 def test_convert_multiple_assistant_messages_with_tools() -> None: """Test multiple consecutive assistant messages with tool interactions.""" messages = [ MessageInput(role="user", content="Multi-step request"), MessageInput( role="assistant", content=[ ToolUseBlock( type="tool_use", id="step1", name="tool1", input={"param": "value1"} ) ], ), MessageInput( role="assistant", content=[ ToolResultBlock( type="tool_result", tool_use_id="step1", content=[TextBlock(type="text", text="Step 1 complete")], is_error=False, ), ToolUseBlock( type="tool_use", id="step2", name="tool2", input={"param": "value2"} ), ], ), MessageInput( role="assistant", content=[ ToolResultBlock( type="tool_result", tool_use_id="step2", content=[TextBlock(type="text", text="Step 2 complete")], is_error=False, ), TextBlock(type="text", text="All steps completed"), ], ), ] result = MessageInput.convert_from_llama_index_messages(messages) assert len(result) == 6 # user, asst, tool, asst, tool, asst assert result[0].role == MessageRole.USER assert result[1].role == MessageRole.ASSISTANT assert result[2].role == MessageRole.TOOL assert result[3].role == MessageRole.ASSISTANT assert result[4].role == MessageRole.TOOL assert result[5].role == MessageRole.ASSISTANT def test_convert_with_meta_fields() -> None: """Test handling of _meta fields in messages.""" messages = [ MessageInput( role="user", content="Test with meta", ), MessageInput( role="assistant", content="Response", ), ] result = MessageInput.convert_from_llama_index_messages(messages) assert len(result) == 2 # Verify meta fields are handled appropriately assert result[0].role == MessageRole.USER assert result[1].role == MessageRole.ASSISTANT @pytest.mark.parametrize( ("message_inputs", "expected_roles_and_tldr"), [ # Case 1: Simple TLDR at end of conversation ( [ MessageInput(role="user", content="Calculate 2+2"), MessageInput( role="assistant", content=[ TextBlock(text="Let me calculate"), ToolUseBlock( id="calc1", name="calculator", input={"expr": "2+2"} ), ToolResultBlock(tool_use_id="calc1", content="4"), TLDRBlock( content=[ TextBlock( text="Request calculation for 2+2", metadata={"type": "tldr", "role": "assistant"}, ), TextBlock( text="4", metadata={"type": "tldr", "role": "tool"}, ), ] ), ], ), ], [ ("user", False, "Calculate 2+2"), ("assistant", True, "Request calculation for 2+2"), ("tool", True, "4"), ], ), # Case 2: Multiple TLDR blocks in same assistant message ( [ MessageInput(role="user", content="Multi-step task"), MessageInput( role="assistant", content=[ TextBlock(text="Step 1: Addition"), ToolUseBlock( id="add1", name="calculator", input={"expr": "10+5"} ), ToolResultBlock( tool_use_id="add1", content="Calculated 10+5=15" ), TLDRBlock( content=[ TextBlock( text="Request calculation for 10+5", metadata={"type": "tldr", "role": "assistant"}, ), TextBlock( text="15", metadata={"type": "tldr", "role": "tool"}, ), ] ), TextBlock(text="Step 2: Multiplication"), ToolUseBlock( id="mult2", name="calculator", input={"expr": "15*2"} ), ToolResultBlock(tool_use_id="mult2", content="Calculated 30"), TLDRBlock( content=[ TextBlock( text="Request calculation for 10+5", metadata={"type": "tldr", "role": "assistant"}, ), TextBlock( text="15", metadata={"type": "tldr", "role": "tool"}, ), TextBlock( text="Request calculation for 15*2", metadata={"type": "tldr", "role": "assistant"}, ), TextBlock( text="30", metadata={"type": "tldr", "role": "tool"}, ), ] ), ], ), ], [ ("user", False, "Multi-step task"), ("assistant", True, "Request calculation for 10+5"), ("tool", True, "15"), ("assistant", True, "Request calculation for 15*2"), ("tool", True, "30"), ], ), # Case 3: TLDR followed by another user query ( [ MessageInput(role="user", content="First calculation"), MessageInput( role="assistant", content=[ TextBlock(text="Calculating..."), ToolUseBlock( id="calc3", name="calculator", input={"expr": "7*8"} ), ToolResultBlock(tool_use_id="calc3", content="56"), TLDRBlock( content=[ TextBlock( text="Request calculation for 7*8", metadata={"type": "tldr", "role": "assistant"}, ), TextBlock( text="56", metadata={"type": "tldr", "role": "tool"}, ), ] ), TextBlock(text="The result is 56"), ], ), MessageInput(role="user", content="What about division?"), ], [ ("user", False, "First calculation"), ("assistant", True, "Request calculation for 7*8"), ("tool", True, "56"), ("assistant", False, "The result is 56"), ("user", False, "What about division?"), ], ), # Case 4: TLDR followed by an user query ( [ MessageInput(role="user", content="This message has to be dropped"), MessageInput( role="assistant", content="This message has to be dropped" ), MessageInput(role="user", content="This message has to be dropped"), MessageInput( role="assistant", content="This message has to be dropped" ), MessageInput(role="user", content="First calculation"), MessageInput( role="assistant", content=[ TLDRBlock( content=[ TextBlock( text="Drop 1", metadata={"type": "tldr", "role": "user"}, ), TextBlock( text="Drop 2", metadata={"type": "tldr", "role": "assistant"}, ), TextBlock( text="Drop 3", metadata={"type": "tldr", "role": "user"}, ), TextBlock( text="Drop 4", metadata={"type": "tldr", "role": "assistant"}, ), ] ), TextBlock(text="Calculating..."), ToolUseBlock( id="calc3", name="calculator", input={"expr": "7*8"} ), ToolResultBlock(tool_use_id="calc3", content="56"), TextBlock(text="The result is 56"), ], ), MessageInput(role="user", content="What about division?"), ], [ ("user", True, "Drop 1"), ("assistant", True, "Drop 2"), ("user", True, "Drop 3"), ("assistant", True, "Drop 4"), ("user", False, "First calculation"), ("assistant", False, "Calculating..."), ("tool", False, "56"), ("assistant", False, "The result is 56"), ("user", False, "What about division?"), ], ), # Case 5: TLDR in left and right ( [ MessageInput(role="user", content="This message has to be dropped"), MessageInput( role="assistant", content="This message has to be dropped" ), MessageInput(role="user", content="This message has to be dropped"), MessageInput( role="assistant", content="This message has to be dropped" ), MessageInput(role="user", content="First calculation"), MessageInput( role="assistant", content=[ TLDRBlock( content=[ TextBlock( text="Drop 1", metadata={"type": "tldr", "role": "user"}, ), TextBlock( text="Drop 2", metadata={"type": "tldr", "role": "assistant"}, ), TextBlock( text="Drop 3", metadata={"type": "tldr", "role": "user"}, ), TextBlock( text="Drop 4", metadata={"type": "tldr", "role": "assistant"}, ), ] ), TextBlock(text="Calculating..."), ToolUseBlock( id="calc3", name="calculator", input={"expr": "7*8"} ), ToolResultBlock(tool_use_id="calc3", content="56"), TLDRBlock( content=[ TextBlock( text="Request calculation for 7*8", metadata={"type": "tldr", "role": "assistant"}, ), TextBlock( text="56", metadata={"type": "tldr", "role": "tool"}, ), ] ), TextBlock(text="The result is 56"), ], ), MessageInput(role="user", content="What about division?"), ], [ ("user", True, "Drop 1"), ("assistant", True, "Drop 2"), ("user", True, "Drop 3"), ("assistant", True, "Drop 4"), ("user", False, "First calculation"), ("assistant", True, "Request calculation for 7*8"), ("tool", True, "56"), ("assistant", False, "The result is 56"), ("user", False, "What about division?"), ], ), # Case 6: Several TLDR in left and right (even by level) ( [ MessageInput(role="user", content="This message has to be dropped"), MessageInput( role="assistant", content="This message has to be dropped" ), MessageInput(role="user", content="This message has to be dropped"), MessageInput( role="assistant", content="This message has to be dropped" ), MessageInput(role="user", content="First calculation"), MessageInput( role="assistant", content=[ TLDRBlock( content=[ TextBlock( text="Drop 1", metadata={"type": "tldr", "role": "user"}, ), TextBlock( text="Drop 2", metadata={"type": "tldr", "role": "assistant"}, ), TextBlock( text="Drop 3", metadata={"type": "tldr", "role": "user"}, ), TextBlock( text="Drop 4", metadata={"type": "tldr", "role": "assistant"}, ), ] ), TextBlock(text="Calculating..."), ToolUseBlock( id="calc3", name="calculator", input={"expr": "7*8"} ), ToolResultBlock(tool_use_id="calc3", content="56"), TLDRBlock( content=[ TextBlock( text="Request calculation for 7*8", metadata={"type": "tldr", "role": "assistant"}, ), TextBlock( text="56", metadata={"type": "tldr", "role": "tool"}, ), ] ), TextBlock(text="The result is 56"), ], ), MessageInput(role="user", content="What about division?"), MessageInput( role="assistant", content=[ TLDRBlock( content=[ TextBlock( text="Drop 1", metadata={"type": "tldr", "role": "user"}, ), TextBlock( text="Drop 2", metadata={"type": "tldr", "role": "assistant"}, ), TextBlock( text="Drop 3", metadata={"type": "tldr", "role": "user"}, ), TextBlock( text="Drop 4", metadata={"type": "tldr", "role": "assistant"}, ), TextBlock( text="What about division? Dropped", metadata={"type": "tldr", "role": "user"}, ), TextBlock( text="The result is 56 Dropped", metadata={"type": "tldr", "role": "assistant"}, ), ] ), TextBlock(text="Calculating..."), ToolUseBlock( id="calc3", name="calculator", input={"expr": "7*8"} ), ToolResultBlock(tool_use_id="calc3", content="56"), TLDRBlock( content=[ TextBlock( text="Request calculation for 7*8", metadata={"type": "tldr", "role": "assistant"}, ), TextBlock( text="56", metadata={"type": "tldr", "role": "tool"}, ), ] ), TextBlock(text="The result is 56"), ], ), ], [ ("user", True, "Drop 1"), ("assistant", True, "Drop 2"), ("user", True, "Drop 3"), ("assistant", True, "Drop 4"), ("user", True, "What about division? Dropped"), ("assistant", True, "The result is 56 Dropped"), ("user", False, "What about division?"), ("assistant", True, "Request calculation for 7*8"), ("tool", True, "56"), ("assistant", False, "The result is 56"), ], ), # Case 4: No TLDR blocks ( [ MessageInput(role="user", content="Simple query"), MessageInput( role="assistant", content=[ TextBlock(text="Simple response"), ToolUseBlock( id="tool1", name="helper", input={"param": "value"} ), ToolResultBlock(tool_use_id="tool1", content="Tool output"), TextBlock(text="Second simple response"), ], ), MessageInput(role="user", content="Follow up"), ], [ ("user", False, "Simple query"), ("assistant", False, "Simple response"), ("tool", False, "Tool output"), ("assistant", False, "Second simple response"), ("user", False, "Follow up"), ], ), # Case 5: Empty TLDR block ( [ MessageInput(role="user", content="Test empty TLDR"), MessageInput( role="assistant", content=[ TextBlock(text="Some response"), TLDRBlock(content=[]), ], ), ], [ ("user", False, "Test empty TLDR"), ("assistant", False, "Some response"), ], ), # Case 7: Explicit right-side TLDR only ( [ MessageInput(role="user", content="Calculate 5+5"), MessageInput( role="assistant", content=[ TextBlock(text="Computing..."), ToolUseBlock( id="calc_right", name="calculator", input={"expr": "5+5"} ), ToolResultBlock(tool_use_id="calc_right", content="10"), TLDRBlock( content=[ TextBlock( text="Requested 5+5", metadata={ "type": "tldr", "role": "assistant", "tldr_side": "right", }, ), TextBlock( text="Result: 10", metadata={ "type": "tldr", "role": "tool", "tldr_side": "right", }, ), ], tldr_side="right", ), TextBlock(text="Done!"), ], ), ], [ ("user", False, "Calculate 5+5"), ("assistant", True, "Requested 5+5"), ("tool", True, "Result: 10"), ("assistant", False, "Done!"), ], ), # Case 8: Explicit left-side TLDR only ( [ MessageInput(role="user", content="Old message 1"), MessageInput(role="assistant", content="Old response 1"), MessageInput(role="user", content="Old message 2"), MessageInput(role="assistant", content="Old response 2"), MessageInput(role="user", content="Current query"), MessageInput( role="assistant", content=[ TLDRBlock( content=[ TextBlock( text="Summary: Old message 1", metadata={ "type": "tldr", "role": "user", "tldr_side": "left", }, ), TextBlock( text="Summary: Old response 1", metadata={ "type": "tldr", "role": "assistant", "tldr_side": "left", }, ), TextBlock( text="Summary: Old message 2", metadata={ "type": "tldr", "role": "user", "tldr_side": "left", }, ), TextBlock( text="Summary: Old response 2", metadata={ "type": "tldr", "role": "assistant", "tldr_side": "left", }, ), ], tldr_side="left", ), TextBlock(text="Current response"), ], ), ], [ ("user", True, "Summary: Old message 1"), ("assistant", True, "Summary: Old response 1"), ("user", True, "Summary: Old message 2"), ("assistant", True, "Summary: Old response 2"), ("user", False, "Current query"), ("assistant", False, "Current response"), ], ), # Case 9: Explicit combination of left and right TLDR ( [ MessageInput(role="user", content="Historic query 1"), MessageInput(role="assistant", content="Historic response 1"), MessageInput(role="user", content="Historic query 2"), MessageInput(role="assistant", content="Historic response 2"), MessageInput(role="user", content="Recent query"), MessageInput( role="assistant", content=[ TLDRBlock( content=[ TextBlock( text="Left summary 1", metadata={ "type": "tldr", "role": "user", "tldr_side": "left", }, ), TextBlock( text="Left summary 2", metadata={ "type": "tldr", "role": "assistant", "tldr_side": "left", }, ), ], tldr_side="left", ), TextBlock(text="Processing..."), ToolUseBlock( id="tool_combo", name="processor", input={"data": "test"} ), ToolResultBlock(tool_use_id="tool_combo", content="Processed"), TLDRBlock( content=[ TextBlock( text="Right summary 1", metadata={ "type": "tldr", "role": "assistant", "tldr_side": "right", }, ), TextBlock( text="Right summary 2", metadata={ "type": "tldr", "role": "tool", "tldr_side": "right", }, ), ], tldr_side="right", ), TextBlock(text="Complete"), ], ), ], [ ("user", True, "Left summary 1"), ("assistant", True, "Left summary 2"), ("user", False, "Recent query"), ("assistant", True, "Right summary 1"), ("tool", True, "Right summary 2"), ("assistant", False, "Complete"), ], ), # Case 10: Multiple right-side TLDRs in sequence ( [ MessageInput(role="user", content="Do multiple tasks"), MessageInput( role="assistant", content=[ TextBlock(text="Task 1"), ToolUseBlock(id="t1", name="task1", input={}), ToolResultBlock(tool_use_id="t1", content="T1 done"), TLDRBlock( content=[ TextBlock( text="Task 1 summarize", metadata={ "type": "tldr", "role": "assistant", "tldr_side": "right", }, ), TextBlock( text="Task 1 summarized", metadata={ "type": "tldr", "role": "tool", "tldr_side": "right", }, ), ], tldr_side="right", ), TextBlock(text="Task 2"), ToolUseBlock(id="t2", name="task2", input={}), ToolResultBlock(tool_use_id="t2", content="T2 done"), TLDRBlock( content=[ TextBlock( text="Task 2 summarize", metadata={ "type": "tldr", "role": "assistant", "tldr_side": "right", }, ), TextBlock( text="Task 2 summarized", metadata={ "type": "tldr", "role": "tool", "tldr_side": "right", }, ), ], tldr_side="right", ), TextBlock(text="Complete"), ], ), ], [ ("user", False, "Do multiple tasks"), ("assistant", True, "Task 1 summarize"), ("tool", True, "Task 1 summarized"), ("assistant", True, "Task 2 summarize"), ("tool", True, "Task 2 summarized"), ("assistant", False, "Complete"), ], ), # Case 11: Double right TLDR with prior conversational history ( [ MessageInput(role="user", content="Old query"), MessageInput(role="assistant", content=[TextBlock(text="Old answer")]), MessageInput(role="user", content="Current query"), MessageInput( role="assistant", content=[ TextBlock(text="Phase 1"), ToolUseBlock( id="dr1", name="calculator", input={"expr": "3+4"} ), ToolResultBlock(tool_use_id="dr1", content="7"), TLDRBlock( tldr_side="right", content=[ TextBlock( text="Do 3+4", metadata={ "type": "tldr", "role": "assistant", "tldr_side": "right", }, ), TextBlock( text="7", metadata={ "type": "tldr", "role": "tool", "tldr_side": "right", }, ), ], ), TextBlock(text="Phase 2"), ToolUseBlock( id="dr2", name="calculator", input={"expr": "7*3"} ), ToolResultBlock(tool_use_id="dr2", content="21"), TLDRBlock( tldr_side="right", content=[ TextBlock( text="Do 3+4", metadata={ "type": "tldr", "role": "assistant", "tldr_side": "right", }, ), TextBlock( text="7", metadata={ "type": "tldr", "role": "tool", "tldr_side": "right", }, ), TextBlock( text="Do 7*3", metadata={ "type": "tldr", "role": "assistant", "tldr_side": "right", }, ), TextBlock( text="21", metadata={ "type": "tldr", "role": "tool", "tldr_side": "right", }, ), ], ), TextBlock(text="Final response"), ], ), ], [ ("user", False, "Old query"), ("assistant", False, "Old answer"), ("user", False, "Current query"), ("assistant", True, "Do 3+4"), ("tool", True, "7"), ("assistant", True, "Do 7*3"), ("tool", True, "21"), ("assistant", False, "Final response"), ], ), # Case 12: Double right TLDR after left TLDR in same assistant message ( [ MessageInput(role="user", content="Legacy 1"), MessageInput(role="assistant", content="Legacy A"), MessageInput(role="user", content="Legacy 2"), MessageInput(role="assistant", content="Legacy B"), MessageInput(role="user", content="Active task"), MessageInput( role="assistant", content=[ TLDRBlock( tldr_side="left", content=[ TextBlock( text="L-user-1", metadata={ "type": "tldr", "role": "user", "tldr_side": "left", }, ), TextBlock( text="L-asst-1", metadata={ "type": "tldr", "role": "assistant", "tldr_side": "left", }, ), ], ), TextBlock(text="Step one"), ToolUseBlock(id="mix1", name="helper", input={}), ToolResultBlock(tool_use_id="mix1", content="ok1"), TLDRBlock( tldr_side="right", content=[ TextBlock( text="R-step-1", metadata={ "type": "tldr", "role": "assistant", "tldr_side": "right", }, ), TextBlock( text="R-ok-1", metadata={ "type": "tldr", "role": "tool", "tldr_side": "right", }, ), ], ), TextBlock(text="Step two"), ToolUseBlock(id="mix2", name="helper", input={}), ToolResultBlock(tool_use_id="mix2", content="ok2"), TLDRBlock( tldr_side="right", content=[ TextBlock( text="R-step-1", metadata={ "type": "tldr", "role": "assistant", "tldr_side": "right", }, ), TextBlock( text="R-ok-1", metadata={ "type": "tldr", "role": "tool", "tldr_side": "right", }, ), TextBlock( text="R-step-2", metadata={ "type": "tldr", "role": "assistant", "tldr_side": "right", }, ), TextBlock( text="R-ok-2", metadata={ "type": "tldr", "role": "tool", "tldr_side": "right", }, ), ], ), TextBlock(text="Done active task"), ], ), ], [ ("user", True, "L-user-1"), ("assistant", True, "L-asst-1"), ("user", False, "Active task"), ("assistant", True, "R-step-1"), ("tool", True, "R-ok-1"), ("assistant", True, "R-step-2"), ("tool", True, "R-ok-2"), ("assistant", False, "Done active task"), ], ), # Case 13: Double right TLDR with previous user block still preserved ( [ MessageInput(role="user", content="First user"), MessageInput( role="assistant", content=[ TextBlock(text="First answer"), ], ), MessageInput(role="user", content="Second user"), MessageInput( role="assistant", content=[ TextBlock(text="Run alpha"), ToolUseBlock(id="p1", name="pipe", input={}), ToolResultBlock(tool_use_id="p1", content="alpha"), TLDRBlock( tldr_side="right", content=[ TextBlock( text="alpha step", metadata={ "type": "tldr", "role": "assistant", "tldr_side": "right", }, ), TextBlock( text="alpha out", metadata={ "type": "tldr", "role": "tool", "tldr_side": "right", }, ), ], ), TextBlock(text="Run beta"), ToolUseBlock(id="p2", name="pipe", input={}), ToolResultBlock(tool_use_id="p2", content="beta"), TLDRBlock( tldr_side="right", content=[ TextBlock( text="alpha step", metadata={ "type": "tldr", "role": "assistant", "tldr_side": "right", }, ), TextBlock( text="alpha out", metadata={ "type": "tldr", "role": "tool", "tldr_side": "right", }, ), TextBlock( text="beta step", metadata={ "type": "tldr", "role": "assistant", "tldr_side": "right", }, ), TextBlock( text="beta out", metadata={ "type": "tldr", "role": "tool", "tldr_side": "right", }, ), ], ), ], ), ], [ ("user", False, "First user"), ("assistant", False, "First answer"), ("user", False, "Second user"), ("assistant", True, "alpha step"), ("tool", True, "alpha out"), ("assistant", True, "beta step"), ("tool", True, "beta out"), ], ), ], ) def test_reorder_tldr_messages_with_message_input( message_inputs: list[MessageInput], expected_roles_and_tldr: list[tuple[str, bool, str]], ) -> None: result = MessageInput.convert_from_llama_index_messages(message_inputs) assert len(result) == len(expected_roles_and_tldr), ( f"Expected {len(expected_roles_and_tldr)} messages, got {len(result)}" ) for i, (actual_msg, (expected_role, has_tldr, expected_content)) in enumerate( zip(result, expected_roles_and_tldr, strict=False) ): assert actual_msg.role.value == expected_role, ( f"Message {i}: expected role {expected_role}, got {actual_msg.role.value}" ) assert bool("tldr" in actual_msg.additional_kwargs) == has_tldr, ( f"Message {i}: TLDR flag mismatch" ) assert actual_msg.content == expected_content, ( f"Message {i}: expected content '{expected_content}', got '{actual_msg.content}'" ) def test_bug_three_consecutive_same_role_with_tldr_pattern() -> None: messages = [ ChatMessage( role=MessageRole.USER, content="User message 1", ), ChatMessage( role=MessageRole.ASSISTANT, content="Regular message 1", additional_kwargs={}, # No TLDR ), ChatMessage( role=MessageRole.ASSISTANT, content="TLDR message 2", additional_kwargs={"tldr": True}, ), ChatMessage( role=MessageRole.ASSISTANT, content="TLDR message 3", additional_kwargs={"tldr": True}, ), ChatMessage( role=MessageRole.ASSISTANT, content="TLDR message 4", additional_kwargs={"tldr": True}, ), ] result = MessageInput._reorder_tldr_messages(messages) assert len(result) > 0 def test_accumulated_right_tldr_with_thinking_blocks_causes_consecutive_assistant_messages() -> ( None ): """Regression: exact structure from production crash. thinking → t1 → result1 thinking → t2 → result2 TLDR1: [ASST(t1), TOOL(-), ASST(t2), TOOL(-)] thinking → t3 → result3 TLDR2: [ASST(t1), TOOL(-), ASST(t2), TOOL(-), ASST(t3), TOOL(-)] thinking → t4 → result4 (no TLDR) TOOL('-') entries are identical ChatMessages → add_result deduplicates them. TLDR2 skips ASST(t1), TOOL(-), ASST(t2) as already added, skips TOOL(-), adds ASST(t3), then skips TOOL(-) → ASST(t2) followed by ASST(t3) directly → consecutive ASSISTANT → ValueError in _validate_message_order. """ messages = [ MessageInput(role="user", content="query"), MessageInput( role="assistant", content=[ ThinkingBlock( type="thinking", thinking="thinking 1", signature="sig_1" ), ToolUseBlock(type="tool_use", id="t1", name="tool1", input={}), ToolResultBlock( type="tool_result", tool_use_id="t1", content=[TextBlock(type="text", text="result1")], ), ThinkingBlock( type="thinking", thinking="thinking 2", signature="sig_2" ), ToolUseBlock(type="tool_use", id="t2", name="tool2", input={}), ToolResultBlock( type="tool_result", tool_use_id="t2", content=[TextBlock(type="text", text="result2")], ), TLDRBlock( tldr_side="right", content=[ TextBlock( text="[t1]", metadata={"type": "tldr", "role": "assistant"} ), TextBlock(text="-", metadata={"type": "tldr", "role": "tool"}), TextBlock( text="[t2]", metadata={"type": "tldr", "role": "assistant"} ), TextBlock(text="-", metadata={"type": "tldr", "role": "tool"}), ], ), ThinkingBlock( type="thinking", thinking="thinking 3", signature="sig_3" ), ToolUseBlock(type="tool_use", id="t3", name="tool3", input={}), ToolResultBlock( type="tool_result", tool_use_id="t3", content=[TextBlock(type="text", text="result3")], ), TLDRBlock( tldr_side="right", content=[ TextBlock( text="[t1]", metadata={"type": "tldr", "role": "assistant"} ), TextBlock(text="-", metadata={"type": "tldr", "role": "tool"}), TextBlock( text="[t2]", metadata={"type": "tldr", "role": "assistant"} ), TextBlock(text="-", metadata={"type": "tldr", "role": "tool"}), TextBlock( text="[t3]", metadata={"type": "tldr", "role": "assistant"} ), TextBlock(text="-", metadata={"type": "tldr", "role": "tool"}), ], ), ThinkingBlock( type="thinking", thinking="thinking 4", signature="sig_4" ), ToolUseBlock(type="tool_use", id="t4", name="tool4", input={}), ToolResultBlock( type="tool_result", tool_use_id="t4", content=[TextBlock(type="text", text="result4")], ), ], ), ] result = MessageInput.convert_from_llama_index_messages(messages) roles = [m.role for m in result] for i in range(len(roles) - 1): assert roles[i] != roles[i + 1], ( f"Consecutive same-role at positions {i}/{i + 1}: {roles}" )