07_streaming.py

"""07: Token-level streaming with include_partial_messages."""

import asyncio
import sys

from claude_agent_sdk import ClaudeAgentOptions, ResultMessage, StreamEvent, query
from demo_util import PROJECT


async def main():
    options = ClaudeAgentOptions(
        cwd=PROJECT,
        include_partial_messages=True,
        allowed_tools=["Read", "Glob"],
        max_turns=4,
    )
    in_tool = False
    event_counts: dict[str, int] = {}

    async for message in query(prompt="Read README.md and summarize it in one sentence.", options=options):
        if isinstance(message, StreamEvent):
            event = message.event
            etype = event.get("type")
            event_counts[etype] = event_counts.get(etype, 0) + 1

            if etype == "content_block_start" and event.get("content_block", {}).get("type") == "tool_use":
                sys.stdout.write(f"\n[Using {event['content_block']['name']}...]")
                in_tool = True
            elif etype == "content_block_delta":
                delta = event.get("delta", {})
                if delta.get("type") == "text_delta" and not in_tool:
                    sys.stdout.write(delta.get("text", ""))
                    sys.stdout.flush()
                elif delta.get("type") == "input_json_delta":
                    sys.stdout.write(".")  # one dot per chunk of tool input
            elif etype == "content_block_stop" and in_tool:
                sys.stdout.write(" done\n")
                in_tool = False
        elif isinstance(message, ResultMessage):
            print(f"\n\n--- Complete ({message.subtype}) ---")

    print("Stream event counts:", event_counts)


asyncio.run(main())