Initial commit
This commit is contained in:
Zhongwei Li
@@ -0,0 +1,436 @@
|
||||
"""
|
||||
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
|
||||
|
||||
Reference in New Issue
Block a user