github/privateGPT
1
0
Fork 1
mirror of https://github.com/imartinez/privateGPT.git synced 2026-07-17 01:48:03 +00:00
Code Issues Packages Projects Releases Wiki Activity
Files
feat/improve-web-search-tool
privateGPT/tests/server/chat/test_input_models.py
Javier Martinez feb11e45cb fix: update claude api spec (#2246) 
* feat: allow to update claude api spec automatically

* feat: script to update sdk / openapi from claude

* feat: add xhigh reasoning

* feat: add mid system content block
2026-06-03 09:23:05 +02:00

2316 lines
90 KiB
Python
Raw Permalink Blame History

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}"
Reference in New Issue View Git Blame Copy Permalink