From 6c25bf7d37cfa976f4620160ecd60f5dbf004d1d Mon Sep 17 00:00:00 2001 From: Dickson Tsai Date: Sat, 19 Jul 2025 13:44:53 -0700 Subject: [PATCH] Fix examples --- examples/streaming_mode.py | 106 ++++++++++------------------- examples/streaming_mode_ipython.py | 67 ++++++++++++------ src/claude_code_sdk/client.py | 45 ++++++------ 3 files changed, 103 insertions(+), 115 deletions(-) diff --git a/examples/streaming_mode.py b/examples/streaming_mode.py index cfa4455..239024d 100644 --- a/examples/streaming_mode.py +++ b/examples/streaming_mode.py @@ -13,6 +13,7 @@ from claude_code_sdk import ( AssistantMessage, ClaudeCodeOptions, ClaudeSDKClient, + CLIConnectionError, ResultMessage, TextBlock, ) @@ -27,18 +28,13 @@ async def example_basic_streaming(): await client.send_message("What is 2+2?") # Receive complete response using the helper method - messages, result = await client.receive_response() - - # Extract text from assistant's response - for msg in messages: + async for msg in client.receive_response(): if isinstance(msg, AssistantMessage): for block in msg.content: if isinstance(block, TextBlock): print(f"Claude: {block.text}") - - # Print cost if available - if result and result.total_cost_usd: - print(f"Cost: ${result.total_cost_usd:.4f}") + elif isinstance(msg, ResultMessage) and msg.total_cost_usd: + print(f"Cost: ${msg.total_cost_usd:.4f}") print("Session ended\n") @@ -52,32 +48,22 @@ async def example_multi_turn_conversation(): print("User: What's the capital of France?") await client.send_message("What's the capital of France?") - messages, _ = await client.receive_response() - # Extract and print response - for msg in messages: - if isinstance(msg, AssistantMessage): - for block in msg.content: - if isinstance(block, TextBlock): - print(f"Claude: {block.text}") + async for msg in client.receive_response(): + content_blocks = getattr(msg, 'content', []) + for block in content_blocks: + if isinstance(block, TextBlock): + print(f"{block.text}") # Second turn - follow-up print("\nUser: What's the population of that city?") await client.send_message("What's the population of that city?") - messages, _ = await client.receive_response() - - for msg in messages: - if isinstance(msg, AssistantMessage): - for block in msg.content: - if isinstance(block, TextBlock): - print(f"Claude: {block.text}") - - for msg in messages: - if isinstance(msg, AssistantMessage): - for block in msg.content: - if isinstance(block, TextBlock): - print(f"Claude: {block.text}") + async for msg in client.receive_response(): + content_blocks = getattr(msg, 'content', []) + for block in content_blocks: + if isinstance(block, TextBlock): + print(f"{block.text}") print("\nConversation ended\n") @@ -102,7 +88,7 @@ async def example_concurrent_responses(): questions = [ "What is 2 + 2?", "What is the square root of 144?", - "What is 15% of 80?", + "What is 10% of 80?", ] for question in questions: @@ -168,9 +154,7 @@ async def example_with_interrupt(): await client.send_message("Never mind, just tell me a quick joke") # Get the joke - messages, result = await client.receive_response() - - for msg in messages: + async for msg in client.receive_response(): if isinstance(msg, AssistantMessage): for block in msg.content: if isinstance(block, TextBlock): @@ -233,10 +217,8 @@ async def example_with_options(): "Create a simple hello.txt file with a greeting message" ) - messages, result = await client.receive_response() - tool_uses = [] - for msg in messages: + async for msg in client.receive_response(): if isinstance(msg, AssistantMessage): for block in msg.content: if isinstance(block, TextBlock): @@ -257,48 +239,34 @@ async def example_error_handling(): client = ClaudeSDKClient() try: - # Connect with custom stream - async def message_stream(): - yield { - "type": "user", - "message": {"role": "user", "content": "Hello"}, - "parent_tool_use_id": None, - "session_id": "error-demo", - } + await client.connect() - await client.connect(message_stream()) + # Send a message that will take time to process + await client.send_message("Run a bash sleep command for 60 seconds") - # Create a background task to consume messages (required for interrupt to work) - consume_task = None - - async def consume_messages(): - """Background message consumer.""" - async for msg in client.receive_messages(): - if isinstance(msg, AssistantMessage): - print("Received response from Claude") - - # Receive messages with timeout + # Try to receive response with a short timeout try: - # Start consuming messages in background - consume_task = asyncio.create_task(consume_messages()) - - # Wait for response with timeout - await asyncio.wait_for(consume_task, timeout=30.0) + messages = [] + async with asyncio.timeout(10.0): + async for msg in client.receive_response(): + messages.append(msg) + if isinstance(msg, AssistantMessage): + for block in msg.content: + if isinstance(block, TextBlock): + print(f"Claude: {block.text[:50]}...") + elif isinstance(msg, ResultMessage): + print("Received complete response") + break except asyncio.TimeoutError: - print("Response timeout - sending interrupt") - # Note: interrupt requires active message consumption - # Since we're already consuming in the background task, interrupt will work - await client.interrupt() + print("\nResponse timeout after 10 seconds - demonstrating graceful handling") + print(f"Received {len(messages)} messages before timeout") - # Cancel the consume task - if consume_task: - consume_task.cancel() - with contextlib.suppress(asyncio.CancelledError): - await consume_task + except CLIConnectionError as e: + print(f"Connection error: {e}") except Exception as e: - print(f"Error: {e}") + print(f"Unexpected error: {e}") finally: # Always disconnect diff --git a/examples/streaming_mode_ipython.py b/examples/streaming_mode_ipython.py index fc2f7cf..6b2b554 100644 --- a/examples/streaming_mode_ipython.py +++ b/examples/streaming_mode_ipython.py @@ -10,17 +10,18 @@ Each example is self-contained and can be run independently. # BASIC STREAMING # ============================================================================ -from claude_code_sdk import ClaudeSDKClient, AssistantMessage, TextBlock +from claude_code_sdk import ClaudeSDKClient, AssistantMessage, TextBlock, ResultMessage async with ClaudeSDKClient() as client: await client.send_message("What is 2+2?") - messages, result = await client.receive_response() - for msg in messages: + async for msg in client.receive_response(): if isinstance(msg, AssistantMessage): for block in msg.content: if isinstance(block, TextBlock): print(f"Claude: {block.text}") + elif isinstance(msg, ResultMessage) and msg.total_cost_usd: + print(f"Cost: ${msg.total_cost_usd:.4f}") # ============================================================================ @@ -31,18 +32,17 @@ import asyncio from claude_code_sdk import ClaudeSDKClient, AssistantMessage, TextBlock async with ClaudeSDKClient() as client: - async def receive_response(): - messages, _ = await client.receive_response() - for msg in messages: + async def send_and_receive(prompt): + await client.send_message(prompt) + async for msg in client.receive_response(): if isinstance(msg, AssistantMessage): for block in msg.content: if isinstance(block, TextBlock): print(f"Claude: {block.text}") - await client.send_message("Tell me a short joke") - await receive_response() - await client.send_message("Now tell me a fun fact") - await receive_response() + await send_and_receive("Tell me a short joke") + print("\n---\n") + await send_and_receive("Now tell me a fun fact") # ============================================================================ @@ -58,8 +58,7 @@ await client.connect() # Helper to get response async def get_response(): - messages, result = await client.receive_response() - for msg in messages: + async for msg in client.receive_response(): if isinstance(msg, AssistantMessage): for block in msg.content: if isinstance(block, TextBlock): @@ -123,9 +122,8 @@ async with ClaudeSDKClient() as client: # Send a new message after interrupt print("\n--- After interrupt, sending new message ---\n") await client.send_message("Just say 'Hello! I was interrupted.'") - messages, result = await client.receive_response() - for msg in messages: + async for msg in client.receive_response(): if isinstance(msg, AssistantMessage): for block in msg.content: if isinstance(block, TextBlock): @@ -142,12 +140,41 @@ try: async with ClaudeSDKClient() as client: await client.send_message("Run a bash sleep command for 60 seconds") - # Timeout after 30 seconds - messages, result = await asyncio.wait_for( - client.receive_response(), timeout=20.0 - ) + # Timeout after 20 seconds + messages = [] + async with asyncio.timeout(20.0): + async for msg in client.receive_response(): + messages.append(msg) + if isinstance(msg, AssistantMessage): + for block in msg.content: + if isinstance(block, TextBlock): + print(f"Claude: {block.text}") except asyncio.TimeoutError: - print("Request timed out") + print("Request timed out after 20 seconds") except Exception as e: - print(f"Error: {e}") \ No newline at end of file + print(f"Error: {e}") + + +# ============================================================================ +# COLLECTING ALL MESSAGES INTO A LIST +# ============================================================================ + +from claude_code_sdk import ClaudeSDKClient, AssistantMessage, TextBlock, ResultMessage + +async with ClaudeSDKClient() as client: + await client.send_message("What are the primary colors?") + + # Collect all messages into a list + messages = [msg async for msg in client.receive_response()] + + # Process them afterwards + for msg in messages: + if isinstance(msg, AssistantMessage): + for block in msg.content: + if isinstance(block, TextBlock): + print(f"Claude: {block.text}") + elif isinstance(msg, ResultMessage): + print(f"Total messages: {len(messages)}") + if msg.total_cost_usd: + print(f"Cost: ${msg.total_cost_usd:.4f}") diff --git a/src/claude_code_sdk/client.py b/src/claude_code_sdk/client.py index a75eece..db7c494 100644 --- a/src/claude_code_sdk/client.py +++ b/src/claude_code_sdk/client.py @@ -146,50 +146,43 @@ class ClaudeSDKClient: raise CLIConnectionError("Not connected. Call connect() first.") await self._transport.interrupt() - async def receive_response(self) -> tuple[list[Message], ResultMessage | None]: + async def receive_response(self) -> AsyncIterator[Message]: """ - Receive a complete response from Claude, collecting all messages until ResultMessage. + Receive messages from Claude until a ResultMessage is received. - Compared to receive_messages(), this is a convenience method that - handles the common pattern of receiving messages until Claude completes - its response. It collects all messages and returns them along with the - final ResultMessage. + This is an async iterator that yields all messages including the final ResultMessage. + It's a convenience method over receive_messages() that automatically stops iteration + after receiving a ResultMessage. - Returns: - tuple: A tuple of (messages, result) where: - - messages: List of all messages received (UserMessage, AssistantMessage, SystemMessage) - - result: The final ResultMessage if received, None if stream ended without result + Yields: + Message: Each message received (UserMessage, AssistantMessage, SystemMessage, ResultMessage) Example: ```python async with ClaudeSDKClient() as client: - # First turn + # Send message and process response await client.send_message("What's the capital of France?") - messages, result = await client.receive_response() - # Extract assistant's response - for msg in messages: + async for msg in client.receive_response(): if isinstance(msg, AssistantMessage): for block in msg.content: if isinstance(block, TextBlock): print(f"Claude: {block.text}") + elif isinstance(msg, ResultMessage): + print(f"Cost: ${msg.total_cost_usd:.4f}") + ``` - # Second turn - await client.send_message("What's the population?") - messages, result = await client.receive_response() - # ... process response + Note: + The iterator will automatically stop after yielding a ResultMessage. + If you need to collect all messages into a list, use: + ```python + messages = [msg async for msg in client.receive_response()] ``` """ - from .types import ResultMessage - - messages = [] async for message in self.receive_messages(): - messages.append(message) + yield message if isinstance(message, ResultMessage): - return messages, message - - # Stream ended without ResultMessage - return messages, None + return async def disconnect(self) -> None: """Disconnect from Claude."""