mirrors/claude-code-sdk-python
1
0
Fork 0
mirror of https://github.com/anthropics/claude-code-sdk-python.git synced 2025-12-23 09:19:52 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Improve UserMessage types to include ToolResultBlock (#101)

Browse source 
Fixes https://github.com/anthropics/claude-code-sdk-python/issues/90
This commit is contained in:
Dickson Tsai 2025-07-31 07:51:39 -07:00 • committed by GitHub
parent df94948edc
commit fa3962de3f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 176 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

47
examples/streaming_mode.py
View file

@ -27,6 +27,8 @@ from claude_code_sdk import (
ResultMessage,
SystemMessage,
TextBlock,
ToolResultBlock,
ToolUseBlock,
UserMessage,
)
@ -303,6 +305,50 @@ async def example_async_iterable_prompt():
print("\n")
async def example_bash_command():
"""Example showing tool use blocks when running bash commands."""
print("=== Bash Command Example ===")
async with ClaudeSDKClient() as client:
print("User: Run a bash echo command")
await client.query("Run a bash echo command that says 'Hello from bash!'")
# Track all message types received
message_types = []
async for msg in client.receive_messages():
message_types.append(type(msg).__name__)
if isinstance(msg, UserMessage):
# User messages can contain tool results
for block in msg.content:
if isinstance(block, TextBlock):
print(f"User: {block.text}")
elif isinstance(block, ToolResultBlock):
print(f"Tool Result (id: {block.tool_use_id}): {block.content[:100] if block.content else 'None'}...")
elif isinstance(msg, AssistantMessage):
# Assistant messages can contain tool use blocks
for block in msg.content:
if isinstance(block, TextBlock):
print(f"Claude: {block.text}")
elif isinstance(block, ToolUseBlock):
print(f"Tool Use: {block.name} (id: {block.id})")
if block.name == "Bash":
command = block.input.get("command", "")
print(f" Command: {command}")
elif isinstance(msg, ResultMessage):
print("Result ended")
if msg.total_cost_usd:
print(f"Cost: ${msg.total_cost_usd:.4f}")
break
print(f"\nMessage types received: {', '.join(set(message_types))}")
print("\n")
async def example_error_handling():
"""Demonstrate proper error handling."""
print("=== Error Handling Example ===")
@ -359,6 +405,7 @@ async def main():
"manual_message_handling": example_manual_message_handling,
"with_options": example_with_options,
"async_iterable_prompt": example_async_iterable_prompt,
"bash_command": example_bash_command,
"error_handling": example_error_handling,
}

25
src/claude_code_sdk/_internal/message_parser.py
View file

@ -45,6 +45,31 @@ def parse_message(data: dict[str, Any]) -> Message:
match message_type:
case "user":
try:
if isinstance(data["message"]["content"], list):
user_content_blocks: list[ContentBlock] = []
for block in data["message"]["content"]:
match block["type"]:
case "text":
user_content_blocks.append(
TextBlock(text=block["text"])
)
case "tool_use":
user_content_blocks.append(
ToolUseBlock(
id=block["id"],
name=block["name"],
input=block["input"],
)
)
case "tool_result":
user_content_blocks.append(
ToolResultBlock(
tool_use_id=block["tool_use_id"],
content=block.get("content"),
is_error=block.get("is_error"),
)
)
return UserMessage(content=user_content_blocks)
return UserMessage(content=data["message"]["content"])
except KeyError as e:
raise MessageParseError(

2
src/claude_code_sdk/types.py
View file

@ -73,7 +73,7 @@ ContentBlock = TextBlock | ToolUseBlock | ToolResultBlock
class UserMessage:
"""User message."""
content: str
content: str | list[ContentBlock]
@dataclass

103
tests/test_message_parser.py
View file

@ -9,6 +9,7 @@ from claude_code_sdk.types import (
ResultMessage,
SystemMessage,
TextBlock,
ToolResultBlock,
ToolUseBlock,
UserMessage,
)
@ -25,6 +26,108 @@ class TestMessageParser:
}
message = parse_message(data)
assert isinstance(message, UserMessage)
assert len(message.content) == 1
assert isinstance(message.content[0], TextBlock)
assert message.content[0].text == "Hello"
def test_parse_user_message_with_tool_use(self):
"""Test parsing a user message with tool_use block."""
data = {
"type": "user",
"message": {
"content": [
{"type": "text", "text": "Let me read this file"},
{
"type": "tool_use",
"id": "tool_456",
"name": "Read",
"input": {"file_path": "/example.txt"},
},
]
},
}
message = parse_message(data)
assert isinstance(message, UserMessage)
assert len(message.content) == 2
assert isinstance(message.content[0], TextBlock)
assert isinstance(message.content[1], ToolUseBlock)
assert message.content[1].id == "tool_456"
assert message.content[1].name == "Read"
assert message.content[1].input == {"file_path": "/example.txt"}
def test_parse_user_message_with_tool_result(self):
"""Test parsing a user message with tool_result block."""
data = {
"type": "user",
"message": {
"content": [
{
"type": "tool_result",
"tool_use_id": "tool_789",
"content": "File contents here",
}
]
},
}
message = parse_message(data)
assert isinstance(message, UserMessage)
assert len(message.content) == 1
assert isinstance(message.content[0], ToolResultBlock)
assert message.content[0].tool_use_id == "tool_789"
assert message.content[0].content == "File contents here"
def test_parse_user_message_with_tool_result_error(self):
"""Test parsing a user message with error tool_result block."""
data = {
"type": "user",
"message": {
"content": [
{
"type": "tool_result",
"tool_use_id": "tool_error",
"content": "File not found",
"is_error": True,
}
]
},
}
message = parse_message(data)
assert isinstance(message, UserMessage)
assert len(message.content) == 1
assert isinstance(message.content[0], ToolResultBlock)
assert message.content[0].tool_use_id == "tool_error"
assert message.content[0].content == "File not found"
assert message.content[0].is_error is True
def test_parse_user_message_with_mixed_content(self):
"""Test parsing a user message with mixed content blocks."""
data = {
"type": "user",
"message": {
"content": [
{"type": "text", "text": "Here's what I found:"},
{
"type": "tool_use",
"id": "use_1",
"name": "Search",
"input": {"query": "test"},
},
{
"type": "tool_result",
"tool_use_id": "use_1",
"content": "Search results",
},
{"type": "text", "text": "What do you think?"},
]
},
}
message = parse_message(data)
assert isinstance(message, UserMessage)
assert len(message.content) == 4
assert isinstance(message.content[0], TextBlock)
assert isinstance(message.content[1], ToolUseBlock)
assert isinstance(message.content[2], ToolResultBlock)
assert isinstance(message.content[3], TextBlock)
def test_parse_valid_assistant_message(self):
"""Test parsing a valid assistant message."""