demo_util.py

"""Shared helpers for the demo scripts: a compact message printer."""

import os
from pathlib import Path

from claude_agent_sdk import (
    AssistantMessage,
    ResultMessage,
    StreamEvent,
    SystemMessage,
    TextBlock,
    ThinkingBlock,
    ToolResultBlock,
    ToolUseBlock,
    UserMessage,
)

_EXAMPLES = Path(__file__).resolve().parent.parent
PRISTINE = _EXAMPLES / "sample-project"            # never modified by the demos
PROJECT = str(_EXAMPLES / "project")               # the working copy the demos run against

if not Path(PROJECT).exists():                     # first run: create the working copy from the pristine project
    import shutil
    shutil.copytree(PRISTINE, PROJECT)


def short(value, n=160):
    s = str(value).replace("\n", "\\n")
    return s if len(s) <= n else s[: n - 3] + "..."


def show(message):
    """Print one SDK message in a compact, human-readable form."""
    if isinstance(message, SystemMessage):
        if message.subtype == "init":
            d = message.data
            print(f"[system:init] session={d.get('session_id')} model={d.get('model')} "
                  f"cli={d.get('claude_code_version')} mode={d.get('permissionMode')}")
            print(f"              tools={d.get('tools')}")
        else:
            print(f"[system:{message.subtype}] {short(message.data)}")
    elif isinstance(message, AssistantMessage):
        for block in message.content:
            if isinstance(block, TextBlock):
                print(f"[assistant:text] {short(block.text, 400)}")
            elif isinstance(block, ThinkingBlock):
                print(f"[assistant:thinking] {short(block.thinking, 120)}")
            elif isinstance(block, ToolUseBlock):
                print(f"[assistant:tool_use] {block.name}({short(block.input, 200)})")
    elif isinstance(message, UserMessage):
        content = message.content
        if isinstance(content, list):
            for block in content:
                if isinstance(block, ToolResultBlock):
                    flag = " (error)" if block.is_error else ""
                    print(f"[user:tool_result]{flag} {short(block.content, 200)}")
        else:
            print(f"[user:text] {short(content)}")
    elif isinstance(message, StreamEvent):
        pass  # handled explicitly by the streaming demo
    elif isinstance(message, ResultMessage):
        cost = f"${message.total_cost_usd:.4f}" if message.total_cost_usd is not None else "n/a"
        print(f"[result] subtype={message.subtype} turns={message.num_turns} "
              f"duration={message.duration_ms}ms cost={cost} session={message.session_id}")
        if message.result:
            print(f"[result:text] {short(message.result, 600)}")
    else:
        print(f"[{type(message).__name__}] {short(message)}")