"""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)
