Claude Code SDK - Python Reference Guide

"""
============================================================
  Claude Code SDK - Python Reference Guide
============================================================

This script is a complete reference for using the Claude Code SDK
(claude-agent-sdk) in your Python scripts. It lets you call Claude
programmatically — like an API — from any Python program.

PREREQUISITES:
    1. Install Claude Code CLI and authenticate:
       npm install -g @anthropic-ai/claude-code
       claude  (follow the login prompts)

    2. Install Python dependencies:
       pip install claude-agent-sdk Pillow

    3. That's it! The SDK uses your Claude Code CLI authentication.
       No API key needed if you're on Claude Max plan.

USAGE:
    python claude_sdk_reference.py              # Run all examples
    python claude_sdk_reference.py --example 1  # Run a specific example (1-5)
"""

import asyncio
import argparse
import base64
import io
import time

# ============================================================
#  IMPORTS — Everything the SDK provides
# ============================================================
from claude_agent_sdk import (
    query,                # The main function — sends prompts and streams responses
    ClaudeAgentOptions,   # Configuration object for controlling Claude's behavior
    AssistantMessage,     # Streamed response chunks from Claude
    ResultMessage,        # Final message with usage stats and cost
    TextBlock,            # A block of text in Claude's response
)
# Also available but not used in every example:
# from claude_agent_sdk import UserMessage  # For tracking user messages in the stream


# ============================================================
#  EXAMPLE 1: Simple Text Query
# ============================================================
# The most basic usage — send a prompt, get a response.
# This is your "Hello World" for the SDK.

async def example_simple_query():
    """Send a simple text prompt and print the response."""
    print("\n" + "=" * 60)
    print("  EXAMPLE 1: Simple Text Query")
    print("=" * 60)

    response_text = ""

    # query() is an async generator — you iterate over messages as they stream in.
    # At minimum, you just need a `prompt` string.
    async for message in query(prompt="What is Python? Reply in exactly one sentence."):

        # AssistantMessage = Claude's response (may arrive in multiple chunks)
        if isinstance(message, AssistantMessage):
            for block in message.content:
                if isinstance(block, TextBlock):
                    response_text += block.text

        # ResultMessage = Final message after Claude is done (contains usage stats)
        elif isinstance(message, ResultMessage):
            pass  # We'll cover usage tracking in Example 5

    print(f"\n  Response: {response_text}")
    return response_text


# ============================================================
#  EXAMPLE 2: Query with Options
# ============================================================
# ClaudeAgentOptions lets you control Claude's behavior:
#   - max_turns: How many back-and-forth turns Claude can take (1 = single response)
#   - permission_mode: What Claude is allowed to do
#   - allowed_tools: Which tools Claude can use (Read, Write, Bash, etc.)
#   - cwd: Working directory for file operations

async def example_with_options():
    """Demonstrate the most useful ClaudeAgentOptions settings."""
    print("\n" + "=" * 60)
    print("  EXAMPLE 2: Query with Options")
    print("=" * 60)

    # --- 2a: Limit Claude to a single turn (no tool use, just answer) ---
    print("\n  2a) Single-turn query (max_turns=1):")

    options = ClaudeAgentOptions(
        max_turns=1,  # Claude responds once, no follow-up tool calls
    )

    response = ""
    async for message in query(
        prompt="Write a Python one-liner that reverses a string. Return ONLY the code.",
        options=options,
    ):
        if isinstance(message, AssistantMessage):
            for block in message.content:
                if isinstance(block, TextBlock):
                    response += block.text

    print(f"  Response: {response}")

    # --- 2b: Allow Claude to create files and run code ---
    print("\n  2b) With file and shell permissions:")

    import os
    import tempfile
    work_dir = tempfile.mkdtemp(prefix="claude_sdk_demo_")

    options = ClaudeAgentOptions(
        cwd=work_dir,                         # Set working directory
        permission_mode="acceptEdits",         # Allow file edits without asking
        allowed_tools=["Read", "Write", "Bash"],  # Only these tools are available
    )

    response = ""
    async for message in query(
        prompt=(
            'Create a file called "hello.py" that prints "Hello from Claude SDK!" '
            "and then run it."
        ),
        options=options,
    ):
        if isinstance(message, AssistantMessage):
            for block in message.content:
                if isinstance(block, TextBlock):
                    response += block.text

    print(f"  Response: {response[:200]}...")

    # Verify the file was created
    created_file = os.path.join(work_dir, "hello.py")
    if os.path.exists(created_file):
        with open(created_file) as f:
            print(f"  Created file contents: {f.read().strip()}")

    return response


# ============================================================
#  EXAMPLE 3: Extended Thinking
# ============================================================
# Extended thinking lets Claude "think step by step" before answering.
# Great for complex reasoning, math, code generation, and analysis.
#
# Key options:
#   - thinking.budget_tokens: How many tokens Claude can use for thinking
#   - effort: "low", "medium", or "high" — controls how hard Claude tries

async def example_extended_thinking():
    """Use extended thinking for complex reasoning tasks."""
    print("\n" + "=" * 60)
    print("  EXAMPLE 3: Extended Thinking")
    print("=" * 60)

    options = ClaudeAgentOptions(
        # Enable extended thinking with a token budget
        thinking={"type": "enabled", "budget_tokens": 5000},
        effort="high",       # Try harder (useful for complex tasks)
        max_turns=1,
    )

    response_text = ""
    thinking_used = False

    async for message in query(
        prompt="What are the first 10 prime numbers? Show your reasoning step by step.",
        options=options,
    ):
        if isinstance(message, AssistantMessage):
            for block in message.content:
                # Check if this is a thinking block (Claude's internal reasoning)
                block_type = getattr(block, "type", None)
                if block_type == "thinking":
                    thinking_text = getattr(block, "thinking", "")
                    thinking_used = True
                    print(f"  [Thinking] {len(thinking_text):,} characters of reasoning")
                # Regular text response
                elif isinstance(block, TextBlock):
                    response_text += block.text

    print(f"\n  Thinking used: {thinking_used}")
    print(f"  Response: {response_text[:300]}...")
    return response_text


# ============================================================
#  EXAMPLE 4: Sending Images
# ============================================================
# To send images, you need an async generator that yields a
# properly formatted user message with content blocks.
#
# Content blocks can include:
#   - {"type": "text", "text": "..."} — text content
#   - {"type": "image", "source": {...}} — base64-encoded image
#
# This pattern is how you send any multi-part content to Claude.

async def create_image_message(content_blocks: list):
    """
    Async generator that yields a properly formatted user message.

    This is the REQUIRED pattern for sending images or multi-part
    content to the SDK. You pass this generator to query() instead
    of a plain string prompt.
    """
    yield {
        "type": "user",
        "message": {
            "role": "user",
            "content": content_blocks,
        },
    }


async def example_send_image():
    """Send an image to Claude for analysis."""
    print("\n" + "=" * 60)
    print("  EXAMPLE 4: Sending Images")
    print("=" * 60)

    # --- Step 1: Prepare the image as base64 ---
    # You can load from a file or create one programmatically.
    # Here we create a simple test image with Pillow.
    try:
        from PIL import Image
    except ImportError:
        print("  Pillow not installed. Run: pip install Pillow")
        return ""

    # Create a simple gradient image as a demo
    img = Image.new("RGB", (200, 200), color=(70, 130, 180))  # Steel blue
    buffer = io.BytesIO()
    img.save(buffer, format="JPEG", quality=70)

    # Encode to base64 string
    image_data = base64.standard_b64encode(buffer.getvalue()).decode("utf-8")
    print(f"  Image size: {len(image_data) / 1024:.1f} KB (base64)")

    # --- Step 2: Build content blocks (text + image) ---
    content_blocks = [
        # Text block — your question or instruction
        {"type": "text", "text": "Describe this image in one sentence. What color do you see?"},
        # Image block — the base64-encoded image
        {
            "type": "image",
            "source": {
                "type": "base64",
                "media_type": "image/jpeg",  # or "image/png"
                "data": image_data,
            },
        },
    ]

    # --- Step 3: Send using the async generator pattern ---
    options = ClaudeAgentOptions(max_turns=1)
    response_text = ""

    # NOTE: We pass the async generator, NOT a string prompt
    async for message in query(
        prompt=create_image_message(content_blocks),
        options=options,
    ):
        if isinstance(message, AssistantMessage):
            for block in message.content:
                if isinstance(block, TextBlock):
                    response_text += block.text

    print(f"  Response: {response_text}")
    return response_text


# ============================================================
#  EXAMPLE 5: Usage & Cost Tracking
# ============================================================
# The ResultMessage contains detailed usage information:
#   - usage: input/output tokens, cache stats
#   - total_cost_usd: total cost in USD
#   - model_usage: per-model breakdown (input, output, cost)
#
# On Claude Max plan, the cost reflects what it *would* cost
# on API billing — it's included in your subscription.

async def example_usage_tracking():
    """Track token usage and costs from Claude's response."""
    print("\n" + "=" * 60)
    print("  EXAMPLE 5: Usage & Cost Tracking")
    print("=" * 60)

    response_text = ""
    usage_data = None
    total_cost = None
    model_usage = None

    options = ClaudeAgentOptions(max_turns=1)

    async for message in query(
        prompt="Explain what an API is in 2 sentences.",
        options=options,
    ):
        # --- Collect response text ---
        if isinstance(message, AssistantMessage):
            for block in message.content:
                if isinstance(block, TextBlock):
                    response_text += block.text

            # Each AssistantMessage also has step-level usage
            if hasattr(message, "usage") and message.usage:
                print(f"  [step usage] {message.usage}")

        # --- Extract final usage from ResultMessage ---
        elif isinstance(message, ResultMessage):
            # Cumulative token usage
            if hasattr(message, "usage") and message.usage:
                usage_data = message.usage

            # Total cost in USD
            if hasattr(message, "total_cost_usd"):
                total_cost = message.total_cost_usd

            # Per-model breakdown (useful when Claude uses sub-agents)
            if hasattr(message, "model_usage"):
                model_usage = message.model_usage

    # --- Display results ---
    print(f"\n  Response: {response_text}\n")

    print("  ┌─────────────────────────────────────")
    print("  │ USAGE REPORT")
    print("  ├─────────────────────────────────────")

    if usage_data:
        print(f"  │ Input tokens:  {usage_data.get('input_tokens', 'N/A')}")
        print(f"  │ Output tokens: {usage_data.get('output_tokens', 'N/A')}")
        print(f"  │ Cache read:    {usage_data.get('cache_read_input_tokens', 'N/A')}")
        print(f"  │ Cache create:  {usage_data.get('cache_creation_input_tokens', 'N/A')}")
    else:
        print("  │ (no usage data returned)")

    if total_cost is not None:
        print(f"  │ Total cost:    ${total_cost:.6f}")

    if model_usage:
        print("  │")
        print("  │ Per-model breakdown:")
        for model_name, mu in model_usage.items():
            cost = mu.get("costUSD", mu.get("cost_usd", "N/A"))
            inp = mu.get("inputTokens", mu.get("input_tokens", "N/A"))
            out = mu.get("outputTokens", mu.get("output_tokens", "N/A"))
            print(f"  │   {model_name}: in={inp}, out={out}, cost=${cost}")

    print("  └─────────────────────────────────────")

    if total_cost is not None:
        print("\n  Note: On Claude Max plan, this cost is included in your subscription.")

    return usage_data


# ============================================================
#  QUICK REFERENCE — ClaudeAgentOptions cheat sheet
# ============================================================
#
#  ClaudeAgentOptions(
#      max_turns=1,                    # Limit back-and-forth turns
#      cwd="/path/to/dir",             # Working directory for file ops
#      permission_mode="acceptEdits",  # "default" | "acceptEdits" | "bypassPermissions"
#      allowed_tools=["Read", "Write", "Bash"],  # Restrict available tools
#      thinking={                      # Extended thinking
#          "type": "enabled",
#          "budget_tokens": 15000,     # Token budget for thinking
#      },
#      effort="high",                  # "low" | "medium" | "high"
#      max_buffer_size=50_000_000,     # Max output buffer (bytes) for large responses
#  )
#
#  query() accepts:
#      prompt="string"                 # Simple text prompt
#      prompt=async_generator          # For images/multi-part content
#      options=ClaudeAgentOptions()    # Configuration
#
#  Message types you'll receive:
#      AssistantMessage  — Claude's response (text, thinking blocks)
#      ResultMessage     — Final stats (usage, cost, model_usage)
#      UserMessage       — Echo of user messages (usually ignored)
#


# ============================================================
#  MAIN — Run examples
# ============================================================

async def main():
    parser = argparse.ArgumentParser(description="Claude Code SDK Reference Examples")
    parser.add_argument(
        "--example", type=int, choices=[1, 2, 3, 4, 5],
        help="Run a specific example (1-5). Omit to run all.",
    )
    args = parser.parse_args()

    examples = {
        1: ("Simple Text Query", example_simple_query),
        2: ("Query with Options", example_with_options),
        3: ("Extended Thinking", example_extended_thinking),
        4: ("Sending Images", example_send_image),
        5: ("Usage & Cost Tracking", example_usage_tracking),
    }

    if args.example:
        # Run single example
        label, func = examples[args.example]
        print(f"\nRunning Example {args.example}: {label}")
        await func()
    else:
        # Run all examples
        print("\n" + "=" * 60)
        print("  CLAUDE CODE SDK — Running All Examples")
        print("=" * 60)
        for num, (label, func) in examples.items():
            try:
                await func()
            except Exception as e:
                print(f"\n  Example {num} ({label}) failed: {type(e).__name__}: {e}")
                import traceback
                traceback.print_exc()

    print("\n" + "=" * 60)
    print("  All done! Happy coding.")
    print("=" * 60 + "\n")


if __name__ == "__main__":
    asyncio.run(main())