Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions src/claude_code_sdk/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@
ResultMessage,
SystemMessage,
TextBlock,
ThinkingBlock,
ToolResultBlock,
ToolUseBlock,
UserMessage,
Expand All @@ -40,6 +41,7 @@
"Message",
"ClaudeCodeOptions",
"TextBlock",
"ThinkingBlock",
"ToolUseBlock",
"ToolResultBlock",
"ContentBlock",
Expand Down
12 changes: 11 additions & 1 deletion src/claude_code_sdk/_internal/message_parser.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
ResultMessage,
SystemMessage,
TextBlock,
ThinkingBlock,
ToolResultBlock,
ToolUseBlock,
UserMessage,
Expand Down Expand Up @@ -53,6 +54,13 @@ def parse_message(data: dict[str, Any]) -> Message:
user_content_blocks.append(
TextBlock(text=block["text"])
)
case "thinking":
user_content_blocks.append(
ThinkingBlock(
thinking=block["thinking"],
signature=block["signature"],
)
)
case "tool_use":
user_content_blocks.append(
ToolUseBlock(
Expand Down Expand Up @@ -100,7 +108,9 @@ def parse_message(data: dict[str, Any]) -> Message:
)
)

return AssistantMessage(content=content_blocks)
return AssistantMessage(
content=content_blocks, model=data["message"]["model"]
)
except KeyError as e:
raise MessageParseError(
f"Missing required field in assistant message: {e}", data
Expand Down
11 changes: 10 additions & 1 deletion src/claude_code_sdk/types.py
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,14 @@ class TextBlock:
text: str


@dataclass
class ThinkingBlock:
"""Thinking content block."""

thinking: str
signature: str


@dataclass
class ToolUseBlock:
"""Tool use content block."""
Expand All @@ -65,7 +73,7 @@ class ToolResultBlock:
is_error: bool | None = None


ContentBlock = TextBlock | ToolUseBlock | ToolResultBlock
ContentBlock = TextBlock | ThinkingBlock | ToolUseBlock | ToolResultBlock


# Message types
Expand All @@ -81,6 +89,7 @@ class AssistantMessage:
"""Assistant message with content blocks."""

content: list[ContentBlock]
model: str


@dataclass
Expand Down
10 changes: 8 additions & 2 deletions tests/test_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,9 @@ async def _test():
) as mock_process:
# Mock the async generator
async def mock_generator():
yield AssistantMessage(content=[TextBlock(text="4")])
yield AssistantMessage(
content=[TextBlock(text="4")], model="claude-opus-4-1-20250805"
)

mock_process.return_value = mock_generator()

Expand All @@ -43,7 +45,10 @@ async def _test():
) as mock_process:

async def mock_generator():
yield AssistantMessage(content=[TextBlock(text="Hello!")])
yield AssistantMessage(
content=[TextBlock(text="Hello!")],
model="claude-opus-4-1-20250805",
)

mock_process.return_value = mock_generator()

Expand Down Expand Up @@ -83,6 +88,7 @@ async def mock_receive():
"message": {
"role": "assistant",
"content": [{"type": "text", "text": "Done"}],
"model": "claude-opus-4-1-20250805",
},
}
yield {
Expand Down
3 changes: 3 additions & 0 deletions tests/test_integration.py
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,7 @@ async def mock_receive():
"message": {
"role": "assistant",
"content": [{"type": "text", "text": "2 + 2 equals 4"}],
"model": "claude-opus-4-1-20250805",
},
}
yield {
Expand Down Expand Up @@ -103,6 +104,7 @@ async def mock_receive():
"input": {"file_path": "/test.txt"},
},
],
"model": "claude-opus-4-1-20250805",
},
}
yield {
Expand Down Expand Up @@ -179,6 +181,7 @@ async def mock_receive():
"text": "Continuing from previous conversation",
}
],
"model": "claude-opus-4-1-20250805",
},
}

Expand Down
3 changes: 2 additions & 1 deletion tests/test_message_parser.py
Original file line number Diff line number Diff line change
Expand Up @@ -142,7 +142,8 @@ def test_parse_valid_assistant_message(self):
"name": "Read",
"input": {"file_path": "/test.txt"},
},
]
],
"model": "claude-opus-4-1-20250805",
},
}
message = parse_message(data)
Expand Down
6 changes: 6 additions & 0 deletions tests/test_streaming_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -187,6 +187,7 @@ async def mock_receive():
"message": {
"role": "assistant",
"content": [{"type": "text", "text": "Hello!"}],
"model": "claude-opus-4-1-20250805",
},
}
yield {
Expand Down Expand Up @@ -229,6 +230,7 @@ async def mock_receive():
"message": {
"role": "assistant",
"content": [{"type": "text", "text": "Answer"}],
"model": "claude-opus-4-1-20250805",
},
}
yield {
Expand All @@ -250,6 +252,7 @@ async def mock_receive():
{"type": "text", "text": "Should not see this"}
],
},
"model": "claude-opus-4-1-20250805",
}

mock_transport.receive_messages = mock_receive
Expand Down Expand Up @@ -335,6 +338,7 @@ async def mock_receive():
"message": {
"role": "assistant",
"content": [{"type": "text", "text": "Response 1"}],
"model": "claude-opus-4-1-20250805",
},
}
await asyncio.sleep(0.1)
Expand Down Expand Up @@ -531,13 +535,15 @@ async def mock_receive():
"message": {
"role": "assistant",
"content": [{"type": "text", "text": "Hello"}],
"model": "claude-opus-4-1-20250805",
},
}
yield {
"type": "assistant",
"message": {
"role": "assistant",
"content": [{"type": "text", "text": "World"}],
"model": "claude-opus-4-1-20250805",
},
}
yield {
Expand Down
20 changes: 18 additions & 2 deletions tests/test_types.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,13 @@
ClaudeCodeOptions,
ResultMessage,
)
from claude_code_sdk.types import TextBlock, ToolResultBlock, ToolUseBlock, UserMessage
from claude_code_sdk.types import (
TextBlock,
ThinkingBlock,
ToolResultBlock,
ToolUseBlock,
UserMessage,
)


class TestMessageTypes:
Expand All @@ -19,10 +25,20 @@ def test_user_message_creation(self):
def test_assistant_message_with_text(self):
"""Test creating an AssistantMessage with text content."""
text_block = TextBlock(text="Hello, human!")
msg = AssistantMessage(content=[text_block])
msg = AssistantMessage(content=[text_block], model="claude-opus-4-1-20250805")
assert len(msg.content) == 1
assert msg.content[0].text == "Hello, human!"

def test_assistant_message_with_thinking(self):
"""Test creating an AssistantMessage with thinking content."""
thinking_block = ThinkingBlock(thinking="I'm thinking...", signature="sig-123")
msg = AssistantMessage(
content=[thinking_block], model="claude-opus-4-1-20250805"
)
assert len(msg.content) == 1
assert msg.content[0].thinking == "I'm thinking..."
assert msg.content[0].signature == "sig-123"

def test_tool_use_block(self):
"""Test creating a ToolUseBlock."""
block = ToolUseBlock(
Expand Down