Skip to content

Commit fa3962d

Browse files
dicksontsaidicksontsai
authored
Improve UserMessage types to include ToolResultBlock (#101)
Fixes #90
1 parent df94948 commit fa3962d

File tree

4 files changed

+176
-1
lines changed

4 files changed

+176
-1
lines changed

‎examples/streaming_mode.py‎

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,8 @@
2727
ResultMessage,
2828
SystemMessage,
2929
TextBlock,
30+
ToolResultBlock,
31+
ToolUseBlock,
3032
UserMessage,
3133
)
3234

@@ -303,6 +305,50 @@ async def create_message_stream():
303305
print("\n")
304306

305307

308+
async def example_bash_command():
309+
"""Example showing tool use blocks when running bash commands."""
310+
print("=== Bash Command Example ===")
311+
312+
async with ClaudeSDKClient() as client:
313+
print("User: Run a bash echo command")
314+
await client.query("Run a bash echo command that says 'Hello from bash!'")
315+
316+
# Track all message types received
317+
message_types = []
318+
319+
async for msg in client.receive_messages():
320+
message_types.append(type(msg).__name__)
321+
322+
if isinstance(msg, UserMessage):
323+
# User messages can contain tool results
324+
for block in msg.content:
325+
if isinstance(block, TextBlock):
326+
print(f"User: {block.text}")
327+
elif isinstance(block, ToolResultBlock):
328+
print(f"Tool Result (id: {block.tool_use_id}): {block.content[:100] if block.content else 'None'}...")
329+
330+
elif isinstance(msg, AssistantMessage):
331+
# Assistant messages can contain tool use blocks
332+
for block in msg.content:
333+
if isinstance(block, TextBlock):
334+
print(f"Claude: {block.text}")
335+
elif isinstance(block, ToolUseBlock):
336+
print(f"Tool Use: {block.name} (id: {block.id})")
337+
if block.name == "Bash":
338+
command = block.input.get("command", "")
339+
print(f" Command: {command}")
340+
341+
elif isinstance(msg, ResultMessage):
342+
print("Result ended")
343+
if msg.total_cost_usd:
344+
print(f"Cost: ${msg.total_cost_usd:.4f}")
345+
break
346+
347+
print(f"\nMessage types received: {', '.join(set(message_types))}")
348+
349+
print("\n")
350+
351+
306352
async def example_error_handling():
307353
"""Demonstrate proper error handling."""
308354
print("=== Error Handling Example ===")
@@ -359,6 +405,7 @@ async def main():
359405
"manual_message_handling": example_manual_message_handling,
360406
"with_options": example_with_options,
361407
"async_iterable_prompt": example_async_iterable_prompt,
408+
"bash_command": example_bash_command,
362409
"error_handling": example_error_handling,
363410
}
364411

‎src/claude_code_sdk/_internal/message_parser.py‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,31 @@ def parse_message(data: dict[str, Any]) -> Message:
4545
match message_type:
4646
case "user":
4747
try:
48+
if isinstance(data["message"]["content"], list):
49+
user_content_blocks: list[ContentBlock] = []
50+
for block in data["message"]["content"]:
51+
match block["type"]:
52+
case "text":
53+
user_content_blocks.append(
54+
TextBlock(text=block["text"])
55+
)
56+
case "tool_use":
57+
user_content_blocks.append(
58+
ToolUseBlock(
59+
id=block["id"],
60+
name=block["name"],
61+
input=block["input"],
62+
)
63+
)
64+
case "tool_result":
65+
user_content_blocks.append(
66+
ToolResultBlock(
67+
tool_use_id=block["tool_use_id"],
68+
content=block.get("content"),
69+
is_error=block.get("is_error"),
70+
)
71+
)
72+
return UserMessage(content=user_content_blocks)
4873
return UserMessage(content=data["message"]["content"])
4974
except KeyError as e:
5075
raise MessageParseError(

‎src/claude_code_sdk/types.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -73,7 +73,7 @@ class ToolResultBlock:
7373
class UserMessage:
7474
"""User message."""
7575

76-
content: str
76+
content: str | list[ContentBlock]
7777

7878

7979
@dataclass

‎tests/test_message_parser.py‎

Lines changed: 103 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
ResultMessage,
1010
SystemMessage,
1111
TextBlock,
12+
ToolResultBlock,
1213
ToolUseBlock,
1314
UserMessage,
1415
)
@@ -25,6 +26,108 @@ def test_parse_valid_user_message(self):
2526
}
2627
message = parse_message(data)
2728
assert isinstance(message, UserMessage)
29+
assert len(message.content) == 1
30+
assert isinstance(message.content[0], TextBlock)
31+
assert message.content[0].text == "Hello"
32+
33+
def test_parse_user_message_with_tool_use(self):
34+
"""Test parsing a user message with tool_use block."""
35+
data = {
36+
"type": "user",
37+
"message": {
38+
"content": [
39+
{"type": "text", "text": "Let me read this file"},
40+
{
41+
"type": "tool_use",
42+
"id": "tool_456",
43+
"name": "Read",
44+
"input": {"file_path": "/example.txt"},
45+
},
46+
]
47+
},
48+
}
49+
message = parse_message(data)
50+
assert isinstance(message, UserMessage)
51+
assert len(message.content) == 2
52+
assert isinstance(message.content[0], TextBlock)
53+
assert isinstance(message.content[1], ToolUseBlock)
54+
assert message.content[1].id == "tool_456"
55+
assert message.content[1].name == "Read"
56+
assert message.content[1].input == {"file_path": "/example.txt"}
57+
58+
def test_parse_user_message_with_tool_result(self):
59+
"""Test parsing a user message with tool_result block."""
60+
data = {
61+
"type": "user",
62+
"message": {
63+
"content": [
64+
{
65+
"type": "tool_result",
66+
"tool_use_id": "tool_789",
67+
"content": "File contents here",
68+
}
69+
]
70+
},
71+
}
72+
message = parse_message(data)
73+
assert isinstance(message, UserMessage)
74+
assert len(message.content) == 1
75+
assert isinstance(message.content[0], ToolResultBlock)
76+
assert message.content[0].tool_use_id == "tool_789"
77+
assert message.content[0].content == "File contents here"
78+
79+
def test_parse_user_message_with_tool_result_error(self):
80+
"""Test parsing a user message with error tool_result block."""
81+
data = {
82+
"type": "user",
83+
"message": {
84+
"content": [
85+
{
86+
"type": "tool_result",
87+
"tool_use_id": "tool_error",
88+
"content": "File not found",
89+
"is_error": True,
90+
}
91+
]
92+
},
93+
}
94+
message = parse_message(data)
95+
assert isinstance(message, UserMessage)
96+
assert len(message.content) == 1
97+
assert isinstance(message.content[0], ToolResultBlock)
98+
assert message.content[0].tool_use_id == "tool_error"
99+
assert message.content[0].content == "File not found"
100+
assert message.content[0].is_error is True
101+
102+
def test_parse_user_message_with_mixed_content(self):
103+
"""Test parsing a user message with mixed content blocks."""
104+
data = {
105+
"type": "user",
106+
"message": {
107+
"content": [
108+
{"type": "text", "text": "Here's what I found:"},
109+
{
110+
"type": "tool_use",
111+
"id": "use_1",
112+
"name": "Search",
113+
"input": {"query": "test"},
114+
},
115+
{
116+
"type": "tool_result",
117+
"tool_use_id": "use_1",
118+
"content": "Search results",
119+
},
120+
{"type": "text", "text": "What do you think?"},
121+
]
122+
},
123+
}
124+
message = parse_message(data)
125+
assert isinstance(message, UserMessage)
126+
assert len(message.content) == 4
127+
assert isinstance(message.content[0], TextBlock)
128+
assert isinstance(message.content[1], ToolUseBlock)
129+
assert isinstance(message.content[2], ToolResultBlock)
130+
assert isinstance(message.content[3], TextBlock)
28131

29132
def test_parse_valid_assistant_message(self):
30133
"""Test parsing a valid assistant message."""

0 commit comments

Comments
 (0)