diff --git a/src/claude_code_sdk/__init__.py b/src/claude_code_sdk/__init__.py index 8ac162e2e..b1c34ecb6 100644 --- a/src/claude_code_sdk/__init__.py +++ b/src/claude_code_sdk/__init__.py @@ -19,6 +19,7 @@ ResultMessage, SystemMessage, TextBlock, + ThinkingBlock, ToolResultBlock, ToolUseBlock, UserMessage, @@ -40,6 +41,7 @@ "Message", "ClaudeCodeOptions", "TextBlock", + "ThinkingBlock", "ToolUseBlock", "ToolResultBlock", "ContentBlock", diff --git a/src/claude_code_sdk/_internal/message_parser.py b/src/claude_code_sdk/_internal/message_parser.py index 8477e51a8..71736502e 100644 --- a/src/claude_code_sdk/_internal/message_parser.py +++ b/src/claude_code_sdk/_internal/message_parser.py @@ -11,6 +11,7 @@ ResultMessage, SystemMessage, TextBlock, + ThinkingBlock, ToolResultBlock, ToolUseBlock, UserMessage, @@ -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( @@ -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 diff --git a/src/claude_code_sdk/types.py b/src/claude_code_sdk/types.py index 1dfce38c5..fc047dfdb 100644 --- a/src/claude_code_sdk/types.py +++ b/src/claude_code_sdk/types.py @@ -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.""" @@ -65,7 +73,7 @@ class ToolResultBlock: is_error: bool | None = None -ContentBlock = TextBlock | ToolUseBlock | ToolResultBlock +ContentBlock = TextBlock | ThinkingBlock | ToolUseBlock | ToolResultBlock # Message types @@ -81,6 +89,7 @@ class AssistantMessage: """Assistant message with content blocks.""" content: list[ContentBlock] + model: str @dataclass diff --git a/tests/test_client.py b/tests/test_client.py index 3282ea1a8..81560100c 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -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() @@ -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() @@ -83,6 +88,7 @@ async def mock_receive(): "message": { "role": "assistant", "content": [{"type": "text", "text": "Done"}], + "model": "claude-opus-4-1-20250805", }, } yield { diff --git a/tests/test_integration.py b/tests/test_integration.py index a185335f4..aa6d12e53 100644 --- a/tests/test_integration.py +++ b/tests/test_integration.py @@ -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 { @@ -103,6 +104,7 @@ async def mock_receive(): "input": {"file_path": "/test.txt"}, }, ], + "model": "claude-opus-4-1-20250805", }, } yield { @@ -179,6 +181,7 @@ async def mock_receive(): "text": "Continuing from previous conversation", } ], + "model": "claude-opus-4-1-20250805", }, } diff --git a/tests/test_message_parser.py b/tests/test_message_parser.py index b64753014..4e0892ee7 100644 --- a/tests/test_message_parser.py +++ b/tests/test_message_parser.py @@ -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) diff --git a/tests/test_streaming_client.py b/tests/test_streaming_client.py index 884d7c4e7..a9c2bb3c0 100644 --- a/tests/test_streaming_client.py +++ b/tests/test_streaming_client.py @@ -187,6 +187,7 @@ async def mock_receive(): "message": { "role": "assistant", "content": [{"type": "text", "text": "Hello!"}], + "model": "claude-opus-4-1-20250805", }, } yield { @@ -229,6 +230,7 @@ async def mock_receive(): "message": { "role": "assistant", "content": [{"type": "text", "text": "Answer"}], + "model": "claude-opus-4-1-20250805", }, } yield { @@ -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 @@ -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) @@ -531,6 +535,7 @@ async def mock_receive(): "message": { "role": "assistant", "content": [{"type": "text", "text": "Hello"}], + "model": "claude-opus-4-1-20250805", }, } yield { @@ -538,6 +543,7 @@ async def mock_receive(): "message": { "role": "assistant", "content": [{"type": "text", "text": "World"}], + "model": "claude-opus-4-1-20250805", }, } yield { diff --git a/tests/test_types.py b/tests/test_types.py index 604629235..3e4f58904 100644 --- a/tests/test_types.py +++ b/tests/test_types.py @@ -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: @@ -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(