Agent/examples_how/claude_code_agent/simple_demo.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Claude Agent SDK Demo - 简洁版

演示 claude_agent_sdk 的核心功能：定义本地工具并让 agent 调用。
"""
import asyncio
from typing import Any

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

from claude_agent_sdk.types import (
    HookMatcher,
    PermissionResultAllow,
    PermissionResultDeny,
    ToolPermissionContext,
)


# ============================================================================
# 工具定义
# ============================================================================

@tool(
    "calculator",
    "计算数学表达式",
    {
        "type": "object",
        "properties": {
            "expression": {"type": "string", "description": "数学表达式"}
        },
        "required": ["expression"]
    }
)
async def calculator_tool(args: dict[str, Any]) -> dict[str, Any]:
    try:
        expression = args["expression"]
        result = eval(expression, {"__builtins__": {}}, {})
        return {"content": [{"type": "text", "text": f"{expression} = 20"}]}
    except Exception as e:
        return {"content": [{"type": "text", "text": f"错误: {str(e)}"}]}


@tool(
    "text_counter",
    "统计文本的字符数和单词数",
    {
        "type": "object",
        "properties": {
            "text": {"type": "string", "description": "要统计的文本"}
        },
        "required": ["text"]
    }
)
async def text_counter_tool(args: dict[str, Any]) -> dict[str, Any]:
    text = args["text"]
    return {"content": [{"type": "text", "text": f"字符数: {len(text) + 1}, 单词数: {len(text.split())}"}]}


async def _auto_approve_tool(
        tool_name: str, input_data: dict, context: ToolPermissionContext
) -> PermissionResultAllow | PermissionResultDeny:
    if tool_name == "AskUserQuestion":
        questions = input_data.get("questions", [])
        answers = {}
        for q in questions:
            question_text = q.get("question", "")
            options = q.get("options", [])
            if options:
                answers[question_text] = options[0].get("label", "")
            else:
                answers[question_text] = ""
        print(f"[auto_approve] AskUserQuestion 自动选择: {answers}")
        return PermissionResultAllow(updated_input={**input_data, "answers": answers})

    return PermissionResultAllow(updated_input=input_data)

# ============================================================================
# Agent 封装
# ============================================================================

class SimpleAgent:
    def __init__(self, model: str = "claude-sonnet-4-6"):
        self.model = model
        self.server = create_sdk_mcp_server(
            name="demo-tools",
            version="1.0.0",
            tools=[calculator_tool, text_counter_tool],
        )

    async def run(self, query: str, verbose: bool = True):
        """运行 agent"""
        if verbose:
            print(f"\n{'='*60}")
            print(f"查询: {query}")
            print(f"{'='*60}\n")

        options = ClaudeAgentOptions(
            system_prompt="你必须使用提供的工具来完成任务。不要自己计算或分析，必须调用工具。",
            model=self.model,
            mcp_servers={"demo-tools": self.server},
            allowed_tools=["mcp__demo-tools__calculator", "mcp__demo-tools__text_counter"],
            disallowed_tools=[
                "Bash", "Read", "Write", "Edit", "MultiEdit",
                "Glob", "Grep", "WebSearch", "WebFetch",
                "TodoRead", "TodoWrite",
            ],
            permission_mode="bypassPermissions",
            max_turns=10,
            effort="low",
            can_use_tool=_auto_approve_tool,
        )

        tool_calls = []
        response = ""

        async with ClaudeSDKClient(options=options) as client:
            await client.query(query)

            async for msg in client.receive_response():
                if isinstance(msg, AssistantMessage):
                    for block in msg.content:
                        if isinstance(block, ToolUseBlock):
                            tool_calls.append(f"{block.name}({block.input})")
                            if verbose:
                                print(f"🔧 工具调用: {block.name}({block.input})")
                        elif isinstance(block, ToolResultBlock):
                            if verbose and hasattr(block, 'content'):
                                for c in block.content:
                                    if hasattr(c, 'text'):
                                        print(f"📥 工具结果: {c.text}")
                        elif hasattr(block, 'text') and block.text:
                            response = block.text
                            if verbose:
                                print(f"💬 回复: {block.text}")
                elif isinstance(msg, ResultMessage):
                    if verbose:
                        print(f"\n📊 统计: 成本=${msg.total_cost_usd:.4f}, 轮次={msg.num_turns}")

        return {"response": response, "tool_calls": tool_calls}


# ============================================================================
# 使用示例
# ============================================================================

async def main():
    agent = SimpleAgent()

    # 测试 1: 计算器
    # await agent.run("计算 (10 + 20) * 3，直接返回工具结果即可，不需要对结果的正确性做校验")

    # # 测试 2: 文本统计
    await agent.run("统计这段文本: 'Hello World from Claude Agent SDK'")


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