Claude Agent SDK

"""Weave + Claude Agent SDK: tracing queries, tools, and conversations.

Call `weave.init()` to auto-patch the SDK — every query, tool call, and
multi-turn conversation is recorded as a trace you can inspect in the Weave UI.
"""

import anyio

import weave
from claude_agent_sdk import (
    AssistantMessage,
    ClaudeAgentOptions,
    ClaudeSDKClient,
    ResultMessage,
    TextBlock,
    ToolUseBlock,
    create_sdk_mcp_server,
    tool,
)

# --- 1. Define MCP tools ---

@tool("add", "Add two numbers", {"a": float, "b": float})
async def add(args: dict) -> dict:
    return {"content": [{"type": "text", "text": str(args["a"] + args["b"])}]}


@tool("multiply", "Multiply two numbers", {"a": float, "b": float})
async def multiply(args: dict) -> dict:
    return {"content": [{"type": "text", "text": str(args["a"] * args["b"])}]}


math_server = create_sdk_mcp_server(
    name="math", version="1.0.0", tools=[add, multiply],
)


# --- 2. Simple one-shot query ---

async def simple_query():
    from claude_agent_sdk import query

    async for msg in query(prompt="What is 2+2?"):
        if isinstance(msg, AssistantMessage):
            for block in msg.content:
                if isinstance(block, TextBlock):
                    print(block.text, end="")
        elif isinstance(msg, ResultMessage):
            print(f"\ncost=${msg.total_cost_usd:.4f}")


# --- 3. MCP tool use ---

async def tool_query():
    options = ClaudeAgentOptions(
        mcp_servers={"math": math_server},
        allowed_tools=["mcp__math__add", "mcp__math__multiply"],
    )

    async with ClaudeSDKClient(options=options) as client:
        await client.query("Using the math tools, compute (3 + 7) * 2.")
        async for msg in client.receive_response():
            if isinstance(msg, AssistantMessage):
                for block in msg.content:
                    if isinstance(block, ToolUseBlock):
                        print(f"  [tool] {block.name}({block.input})")
                    elif isinstance(block, TextBlock):
                        print(block.text, end="")
            elif isinstance(msg, ResultMessage):
                print(f"\ncost=${msg.total_cost_usd:.4f}")


# --- 4. Multi-turn conversation with thread context ---

async def threaded_conversation():
    with weave.thread("my-conversation") as t:
        print(f"thread_id={t.thread_id}")

        async with ClaudeSDKClient() as client:
            for prompt in [
                "Name a famous sorting algorithm.",
                "What is its time complexity?",
            ]:
                await client.query(prompt)
                reply = ""
                async for msg in client.receive_response():
                    if isinstance(msg, AssistantMessage):
                        reply += "".join(
                            b.text for b in msg.content if isinstance(b, TextBlock)
                        )
                    elif isinstance(msg, ResultMessage):
                        print(f"Q: {prompt}\nA: {reply.strip()[:120]}\n")


# --- Run all examples ---

async def main():
    print("=== Simple Query ===")
    await simple_query()

    print("\n=== MCP Tool Use ===")
    await tool_query()

    print("\n=== Multi-Turn Thread ===")
    await threaded_conversation()


if __name__ == "__main__":
    weave.init("your-team-name/claude-agent-sdk-demo")
    anyio.run(main)