437 lines
12 KiB
Python
437 lines
12 KiB
Python
"""
|
|
Claude Code SDK - The SDK Way
|
|
|
|
This module demonstrates the idiomatic way to use the Claude Code Python SDK
|
|
for programmatic agent interactions. It focuses on clean, type-safe patterns
|
|
using the SDK's native abstractions.
|
|
|
|
Key Concepts:
|
|
- Use `query()` for one-shot operations
|
|
- Use `ClaudeSDKClient` for interactive sessions
|
|
- Work directly with SDK message types
|
|
- Leverage async/await for clean concurrency
|
|
- Configure options for your use case
|
|
|
|
Example Usage:
|
|
# Simple query
|
|
async for message in query(prompt="What is 2 + 2?"):
|
|
if isinstance(message, AssistantMessage):
|
|
print(extract_text(message))
|
|
|
|
# With options
|
|
options = ClaudeCodeOptions(
|
|
model="claude-sonnet-4-20250514",
|
|
allowed_tools=["Read", "Write"],
|
|
permission_mode="bypassPermissions"
|
|
)
|
|
async for message in query(prompt="Create hello.py", options=options):
|
|
process_message(message)
|
|
|
|
# Interactive session
|
|
async with create_session() as client:
|
|
await client.query("Debug this error")
|
|
async for msg in client.receive_response():
|
|
handle_message(msg)
|
|
"""
|
|
|
|
import logging
|
|
from pathlib import Path
|
|
from typing import AsyncIterator, Optional, List
|
|
from contextlib import asynccontextmanager
|
|
|
|
# Import all SDK components we'll use
|
|
from claude_code_sdk import (
|
|
# Main functions
|
|
query,
|
|
ClaudeSDKClient,
|
|
|
|
# Configuration
|
|
ClaudeCodeOptions,
|
|
PermissionMode,
|
|
|
|
# Message types
|
|
Message,
|
|
AssistantMessage,
|
|
UserMessage,
|
|
SystemMessage,
|
|
ResultMessage,
|
|
|
|
# Content blocks
|
|
ContentBlock,
|
|
TextBlock,
|
|
ToolUseBlock,
|
|
ToolResultBlock,
|
|
|
|
# Errors
|
|
ClaudeSDKError,
|
|
CLIConnectionError,
|
|
CLINotFoundError,
|
|
ProcessError,
|
|
)
|
|
|
|
# Set up logging
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
# ============================================================================
|
|
# UTILITY FUNCTIONS
|
|
# ============================================================================
|
|
|
|
def extract_text(message: AssistantMessage) -> str:
|
|
"""Extract all text content from an assistant message.
|
|
|
|
The SDK way: Work directly with typed message objects.
|
|
|
|
Args:
|
|
message: AssistantMessage with content blocks
|
|
|
|
Returns:
|
|
Concatenated text from all text blocks
|
|
"""
|
|
texts = []
|
|
for block in message.content:
|
|
if isinstance(block, TextBlock):
|
|
texts.append(block.text)
|
|
return "\n".join(texts)
|
|
|
|
|
|
def extract_tool_uses(message: AssistantMessage) -> List[ToolUseBlock]:
|
|
"""Extract all tool use blocks from an assistant message.
|
|
|
|
Args:
|
|
message: AssistantMessage with content blocks
|
|
|
|
Returns:
|
|
List of ToolUseBlock objects
|
|
"""
|
|
return [
|
|
block for block in message.content
|
|
if isinstance(block, ToolUseBlock)
|
|
]
|
|
|
|
|
|
def get_result_text(messages: List[Message]) -> str:
|
|
"""Extract final result text from a list of messages.
|
|
|
|
Args:
|
|
messages: List of messages from a query
|
|
|
|
Returns:
|
|
Result text or assistant responses
|
|
"""
|
|
# First check for ResultMessage
|
|
for msg in reversed(messages):
|
|
if isinstance(msg, ResultMessage) and msg.result:
|
|
return msg.result
|
|
|
|
# Otherwise collect assistant text
|
|
texts = []
|
|
for msg in messages:
|
|
if isinstance(msg, AssistantMessage):
|
|
text = extract_text(msg)
|
|
if text:
|
|
texts.append(text)
|
|
|
|
return "\n".join(texts)
|
|
|
|
|
|
# ============================================================================
|
|
# ONE-SHOT QUERIES (The Simple SDK Way)
|
|
# ============================================================================
|
|
|
|
async def simple_query(prompt: str, model: str = "claude-sonnet-4-5-20250929") -> str:
|
|
"""Simple one-shot query with text response.
|
|
|
|
The SDK way: Direct use of query() with minimal setup.
|
|
|
|
Args:
|
|
prompt: What to ask Claude
|
|
model: Which model to use
|
|
|
|
Returns:
|
|
Text response from Claude
|
|
|
|
Example:
|
|
response = await simple_query("What is 2 + 2?")
|
|
print(response) # "4" or "2 + 2 equals 4"
|
|
"""
|
|
options = ClaudeCodeOptions(model=model)
|
|
|
|
texts = []
|
|
async for message in query(prompt=prompt, options=options):
|
|
if isinstance(message, AssistantMessage):
|
|
text = extract_text(message)
|
|
if text:
|
|
texts.append(text)
|
|
|
|
return "\n".join(texts) if texts else "No response"
|
|
|
|
|
|
async def query_with_tools(
|
|
prompt: str,
|
|
allowed_tools: List[str],
|
|
working_dir: Optional[Path] = None
|
|
) -> AsyncIterator[Message]:
|
|
"""Query with specific tools enabled.
|
|
|
|
The SDK way: Configure options for your use case.
|
|
|
|
Args:
|
|
prompt: What to ask Claude
|
|
allowed_tools: List of tool names to allow
|
|
working_dir: Optional working directory
|
|
|
|
Yields:
|
|
SDK message objects
|
|
|
|
Example:
|
|
async for msg in query_with_tools(
|
|
"Create a Python script",
|
|
allowed_tools=["Write", "Read"]
|
|
):
|
|
if isinstance(msg, AssistantMessage):
|
|
for block in msg.content:
|
|
if isinstance(block, ToolUseBlock):
|
|
print(f"Using tool: {block.name}")
|
|
"""
|
|
options = ClaudeCodeOptions(
|
|
allowed_tools=allowed_tools,
|
|
cwd=str(working_dir) if working_dir else None,
|
|
permission_mode="bypassPermissions" # For automated workflows
|
|
)
|
|
|
|
async for message in query(prompt=prompt, options=options):
|
|
yield message
|
|
|
|
|
|
async def collect_query_response(
|
|
prompt: str,
|
|
options: Optional[ClaudeCodeOptions] = None
|
|
) -> tuple[List[Message], Optional[ResultMessage]]:
|
|
"""Collect all messages from a query.
|
|
|
|
The SDK way: Async iteration with type checking.
|
|
|
|
Args:
|
|
prompt: What to ask Claude
|
|
options: Optional configuration
|
|
|
|
Returns:
|
|
Tuple of (all_messages, result_message)
|
|
|
|
Example:
|
|
messages, result = await collect_query_response("List files")
|
|
if result and not result.is_error:
|
|
print("Success!")
|
|
for msg in messages:
|
|
process_message(msg)
|
|
"""
|
|
if options is None:
|
|
options = ClaudeCodeOptions()
|
|
|
|
messages = []
|
|
result = None
|
|
|
|
async for message in query(prompt=prompt, options=options):
|
|
messages.append(message)
|
|
if isinstance(message, ResultMessage):
|
|
result = message
|
|
|
|
return messages, result
|
|
|
|
|
|
# ============================================================================
|
|
# INTERACTIVE SESSIONS (The SDK Client Way)
|
|
# ============================================================================
|
|
|
|
@asynccontextmanager
|
|
async def create_session(
|
|
model: str = "claude-sonnet-4-5-20250929",
|
|
working_dir: Optional[Path] = None
|
|
):
|
|
"""Create an interactive session with Claude.
|
|
|
|
The SDK way: Use context managers for resource management.
|
|
|
|
Args:
|
|
model: Which model to use
|
|
working_dir: Optional working directory
|
|
|
|
Yields:
|
|
Connected ClaudeSDKClient
|
|
|
|
Example:
|
|
async with create_session() as client:
|
|
await client.query("Hello")
|
|
async for msg in client.receive_response():
|
|
print(msg)
|
|
"""
|
|
options = ClaudeCodeOptions(
|
|
model=model,
|
|
cwd=str(working_dir) if working_dir else None,
|
|
permission_mode="bypassPermissions"
|
|
)
|
|
|
|
client = ClaudeSDKClient(options=options)
|
|
await client.connect()
|
|
|
|
try:
|
|
yield client
|
|
finally:
|
|
await client.disconnect()
|
|
|
|
|
|
async def interactive_conversation(prompts: List[str]) -> List[Message]:
|
|
"""Have an interactive conversation with Claude.
|
|
|
|
The SDK way: Bidirectional communication with the client.
|
|
|
|
Args:
|
|
prompts: List of prompts to send in sequence
|
|
|
|
Returns:
|
|
All messages from the conversation
|
|
|
|
Example:
|
|
messages = await interactive_conversation([
|
|
"What's the weather like?",
|
|
"Tell me more about clouds",
|
|
"How do they form?"
|
|
])
|
|
"""
|
|
all_messages = []
|
|
|
|
async with create_session() as client:
|
|
for prompt in prompts:
|
|
# Send prompt
|
|
await client.query(prompt)
|
|
|
|
# Collect response
|
|
async for msg in client.receive_response():
|
|
all_messages.append(msg)
|
|
if isinstance(msg, ResultMessage):
|
|
break
|
|
|
|
return all_messages
|
|
|
|
|
|
# ============================================================================
|
|
# ERROR HANDLING (The SDK Way)
|
|
# ============================================================================
|
|
|
|
async def safe_query(prompt: str) -> tuple[Optional[str], Optional[str]]:
|
|
"""Query with comprehensive error handling.
|
|
|
|
The SDK way: Handle specific SDK exceptions.
|
|
|
|
Args:
|
|
prompt: What to ask Claude
|
|
|
|
Returns:
|
|
Tuple of (response_text, error_message)
|
|
|
|
Example:
|
|
response, error = await safe_query("Help me debug this")
|
|
if error:
|
|
print(f"Error: {error}")
|
|
else:
|
|
print(f"Response: {response}")
|
|
"""
|
|
try:
|
|
response = await simple_query(prompt)
|
|
return response, None
|
|
|
|
except CLINotFoundError:
|
|
return None, "Claude Code CLI not found. Install with: npm install -g @anthropic-ai/claude-code"
|
|
|
|
except CLIConnectionError as e:
|
|
return None, f"Connection error: {str(e)}"
|
|
|
|
except ProcessError as e:
|
|
return None, f"Process error (exit code {e.exit_code}): {str(e)}"
|
|
|
|
except ClaudeSDKError as e:
|
|
return None, f"SDK error: {str(e)}"
|
|
|
|
except Exception as e:
|
|
return None, f"Unexpected error: {str(e)}"
|
|
|
|
|
|
# ============================================================================
|
|
# ADVANCED PATTERNS (The SDK Way)
|
|
# ============================================================================
|
|
|
|
async def stream_with_progress(
|
|
prompt: str,
|
|
on_text: Optional[callable] = None,
|
|
on_tool: Optional[callable] = None
|
|
) -> ResultMessage:
|
|
"""Stream query with progress callbacks.
|
|
|
|
The SDK way: Process messages as they arrive.
|
|
|
|
Args:
|
|
prompt: What to ask Claude
|
|
on_text: Callback for text blocks (optional)
|
|
on_tool: Callback for tool use blocks (optional)
|
|
|
|
Returns:
|
|
Final ResultMessage
|
|
|
|
Example:
|
|
result = await stream_with_progress(
|
|
"Analyze this codebase",
|
|
on_text=lambda text: print(f"Claude: {text}"),
|
|
on_tool=lambda tool: print(f"Using: {tool.name}")
|
|
)
|
|
print(f"Cost: ${result.total_cost_usd:.4f}")
|
|
"""
|
|
result = None
|
|
|
|
async for message in query(prompt=prompt):
|
|
if isinstance(message, AssistantMessage):
|
|
for block in message.content:
|
|
if isinstance(block, TextBlock) and on_text:
|
|
on_text(block.text)
|
|
elif isinstance(block, ToolUseBlock) and on_tool:
|
|
on_tool(block)
|
|
|
|
elif isinstance(message, ResultMessage):
|
|
result = message
|
|
|
|
return result
|
|
|
|
|
|
async def query_with_timeout(prompt: str, timeout_seconds: float = 30) -> Optional[str]:
|
|
"""Query with timeout protection.
|
|
|
|
The SDK way: Use asyncio for timeout control.
|
|
|
|
Args:
|
|
prompt: What to ask Claude
|
|
timeout_seconds: Maximum time to wait
|
|
|
|
Returns:
|
|
Response text or None if timeout
|
|
|
|
Example:
|
|
response = await query_with_timeout("Complex analysis", timeout_seconds=60)
|
|
if response is None:
|
|
print("Query timed out")
|
|
"""
|
|
import asyncio
|
|
|
|
try:
|
|
# Create the query task
|
|
async def _query():
|
|
return await simple_query(prompt)
|
|
|
|
# Run with timeout
|
|
response = await asyncio.wait_for(_query(), timeout=timeout_seconds)
|
|
return response
|
|
|
|
except asyncio.TimeoutError:
|
|
logger.warning(f"Query timed out after {timeout_seconds} seconds")
|
|
return None
|
|
|