第 4 课：自定义工具 MCP

"""待办管理工具"""

import json
import datetime
from claude_agent_sdk import tool

# 简单存储（后面第 8 课换成 SQLite）
TODO_LIST: list[dict] = []
_id_counter = 0


@tool(
    "add_todo",
    "添加一条待办任务",
    {
        "content": {"type": "string", "description": "任务内容"},
        "priority": {"type": "string", "description": "优先级: high/medium/low，默认 medium"},
        "deadline": {"type": "string", "description": "截止日期，如 2025-03-20，可选"},
    }
)
async def add_todo(content: str, priority: str = "medium", deadline: str = "") -> str:
    """添加待办"""
    global _id_counter
    _id_counter += 1

    todo = {
        "id": _id_counter,
        "content": content,
        "priority": priority,
        "deadline": deadline or None,
        "done": False,
        "created_at": datetime.datetime.now().strftime("%Y-%m-%d %H:%M"),
    }
    TODO_LIST.append(todo)

    emoji = {"high": "🔴", "medium": "🟡", "low": "🟢"}
    result = f"✅ 已添加任务 #{_id_counter}: {content}"
    result += f" [{emoji.get(priority, '⚪')}{priority}]"
    if deadline:
        result += f" 截止: {deadline}"
    return result


@tool(
    "list_todos",
    "列出待办任务。可以按状态筛选（all/pending/done）",
    {"status": {"type": "string", "description": "筛选: all/pending/done，默认 pending"}}
)
async def list_todos(status: str = "pending") -> str:
    """列出待办"""
    if not TODO_LIST:
        return "📋 待办清单是空的，用 add_todo 添加任务吧！"

    if status == "pending":
        items = [t for t in TODO_LIST if not t["done"]]
    elif status == "done":
        items = [t for t in TODO_LIST if t["done"]]
    else:
        items = TODO_LIST

    if not items:
        return f"没有 {status} 状态的任务"

    # 按优先级排序
    order = {"high": 0, "medium": 1, "low": 2}
    items.sort(key=lambda x: order.get(x["priority"], 9))

    emoji = {"high": "🔴", "medium": "🟡", "low": "🟢"}
    lines = [f"📋 待办清单（{len(items)} 项）:\n"]
    for t in items:
        check = "✅" if t["done"] else "☐"
        e = emoji.get(t["priority"], "⚪")
        line = f"  {check} #{t['id']} {e} {t['content']}"
        if t.get("deadline"):
            line += f" (截止: {t['deadline']})"
        lines.append(line)

    return "\n".join(lines)


@tool(
    "complete_todo",
    "标记一个待办任务为已完成",
    {"todo_id": {"type": "number", "description": "任务 ID"}}
)
async def complete_todo(todo_id: int) -> str:
    """完成任务"""
    for t in TODO_LIST:
        if t["id"] == todo_id:
            if t["done"]:
                return f"任务 #{todo_id} 已经完成过了"
            t["done"] = True
            return f"✅ 任务 #{todo_id} 已完成: {t['content']}"
    return f"未找到任务 #{todo_id}"


@tool(
    "delete_todo",
    "删除一个待办任务",
    {"todo_id": {"type": "number", "description": "任务 ID"}}
)
async def delete_todo(todo_id: int) -> str:
    """删除任务"""
    for i, t in enumerate(TODO_LIST):
        if t["id"] == todo_id:
            removed = TODO_LIST.pop(i)
            return f"🗑️ 已删除任务 #{todo_id}: {removed['content']}"
    return f"未找到任务 #{todo_id}"

"""笔记管理工具"""

import datetime
from claude_agent_sdk import tool

NOTES: list[dict] = []
_note_id = 0


@tool(
    "add_note",
    "保存一条笔记",
    {
        "content": {"type": "string", "description": "笔记内容"},
        "tags": {"type": "string", "description": "标签，用逗号分隔，如: python,学习"},
    }
)
async def add_note(content: str, tags: str = "") -> str:
    """保存笔记"""
    global _note_id
    _note_id += 1

    tag_list = [t.strip() for t in tags.split(",") if t.strip()] if tags else []

    note = {
        "id": _note_id,
        "content": content,
        "tags": tag_list,
        "created_at": datetime.datetime.now().strftime("%Y-%m-%d %H:%M"),
    }
    NOTES.append(note)

    result = f"📝 已保存笔记 #{_note_id}"
    if tag_list:
        result += f" 标签: {', '.join('#' + t for t in tag_list)}"
    return result


@tool(
    "search_notes",
    "搜索笔记。可以按关键词或标签搜索",
    {
        "keyword": {"type": "string", "description": "搜索关键词"},
        "tag": {"type": "string", "description": "按标签筛选（可选）"},
    }
)
async def search_notes(keyword: str = "", tag: str = "") -> str:
    """搜索笔记"""
    if not NOTES:
        return "📓 笔记本是空的"

    results = NOTES
    if keyword:
        results = [n for n in results if keyword.lower() in n["content"].lower()]
    if tag:
        results = [n for n in results if tag.lower() in [t.lower() for t in n["tags"]]]

    if not results:
        return f"没有找到匹配的笔记"

    lines = [f"📓 找到 {len(results)} 条笔记:\n"]
    for n in results:
        tags_str = " ".join(f"#{t}" for t in n["tags"]) if n["tags"] else ""
        lines.append(f"  [{n['created_at']}] #{n['id']}")
        lines.append(f"  {n['content'][:100]}")
        if tags_str:
            lines.append(f"  {tags_str}")
        lines.append("")

    return "\n".join(lines)


@tool(
    "list_recent_notes",
    "列出最近的笔记",
    {"count": {"type": "number", "description": "列出几条，默认 5"}}
)
async def list_recent_notes(count: int = 5) -> str:
    """最近的笔记"""
    if not NOTES:
        return "📓 笔记本是空的"

    recent = NOTES[-count:][::-1]
    lines = [f"📓 最近 {len(recent)} 条笔记:\n"]
    for n in recent:
        tags_str = " ".join(f"#{t}" for t in n["tags"]) if n["tags"] else ""
        lines.append(f"  [{n['created_at']}] {n['content'][:80]} {tags_str}")

    return "\n".join(lines)

"""文件操作工具"""

import os
from pathlib import Path
from claude_agent_sdk import tool


@tool(
    "list_files",
    "列出指定目录下的文件和子目录",
    {
        "path": {"type": "string", "description": "目录路径"},
        "pattern": {"type": "string", "description": "文件名过滤（如 *.py），可选"},
    }
)
async def list_files(path: str = ".", pattern: str = "") -> str:
    """列出文件"""
    p = Path(path).expanduser()
    if not p.exists():
        return f"目录不存在: {path}"
    if not p.is_dir():
        return f"不是目录: {path}"

    try:
        if pattern:
            files = sorted(p.glob(pattern))
        else:
            files = sorted(p.iterdir())

        if not files:
            return f"目录 {path} 是空的" + (f"（过滤: {pattern}）" if pattern else "")

        lines = [f"📁 {path} ({len(files)} 项):\n"]
        for f in files[:50]:  # 最多显示 50 个
            if f.is_dir():
                lines.append(f"  📁 {f.name}/")
            else:
                size = f.stat().st_size
                if size > 1024 * 1024:
                    size_str = f"{size / 1024 / 1024:.1f}MB"
                elif size > 1024:
                    size_str = f"{size / 1024:.1f}KB"
                else:
                    size_str = f"{size}B"
                lines.append(f"  📄 {f.name} ({size_str})")

        if len(files) > 50:
            lines.append(f"  ... 还有 {len(files) - 50} 个文件")

        return "\n".join(lines)
    except PermissionError:
        return f"无权限访问: {path}"


@tool(
    "find_large_files",
    "找出指定目录下超过指定大小的文件",
    {
        "path": {"type": "string", "description": "目录路径"},
        "min_size_mb": {"type": "number", "description": "最小大小（MB），默认 10"},
    }
)
async def find_large_files(path: str = ".", min_size_mb: float = 10) -> str:
    """找大文件"""
    p = Path(path).expanduser()
    if not p.exists():
        return f"目录不存在: {path}"

    min_bytes = int(min_size_mb * 1024 * 1024)
    large_files = []

    try:
        for f in p.rglob("*"):
            if f.is_file():
                try:
                    size = f.stat().st_size
                    if size >= min_bytes:
                        large_files.append((f, size))
                except (PermissionError, OSError):
                    continue
    except PermissionError:
        return f"无权限访问: {path}"

    if not large_files:
        return f"没有找到超过 {min_size_mb}MB 的文件"

    large_files.sort(key=lambda x: -x[1])

    lines = [f"🔍 找到 {len(large_files)} 个大文件（>{min_size_mb}MB）:\n"]
    for f, size in large_files[:20]:
        lines.append(f"  {size / 1024 / 1024:.1f}MB  {f.relative_to(p)}")

    return "\n".join(lines)

from claude_agent_sdk import create_sdk_mcp_server

# 把所有工具注册到 MCP 服务器
productivity_tools = create_sdk_mcp_server(
    name="productivity",
    version="1.0.0",
    tools=[
        # 待办
        add_todo, list_todos, complete_todo, delete_todo,
        # 笔记
        add_note, search_notes, list_recent_notes,
        # 文件
        list_files, find_large_files,
    ],
)

import anyio
from claude_agent_sdk import (
    query, ClaudeAgentOptions,
    AssistantMessage, TextBlock, ToolUseBlock,
    tool, create_sdk_mcp_server,
)

# （假设上面的工具都已定义）

SYSTEM_PROMPT = """你是私人助理 MiniClaw。你有以下工具：

待办管理: add_todo, list_todos, complete_todo, delete_todo
笔记: add_note, search_notes, list_recent_notes
文件: list_files, find_large_files

规则：
- 用户说"加个任务"→ 用 add_todo
- 用户说"记一下"→ 用 add_note
- 用户说"看看有什么要做的"→ 用 list_todos
- 主动为笔记加合适的标签
- 用中文回复，简洁友好"""

async def main():
    options = ClaudeAgentOptions(
        system_prompt=SYSTEM_PROMPT,
        mcp_servers={"productivity": productivity_tools},
        allowed_tools=[
            "mcp__productivity__add_todo",
            "mcp__productivity__list_todos",
            "mcp__productivity__complete_todo",
            "mcp__productivity__delete_todo",
            "mcp__productivity__add_note",
            "mcp__productivity__search_notes",
            "mcp__productivity__list_recent_notes",
            "mcp__productivity__list_files",
            "mcp__productivity__find_large_files",
        ],
    )

    tasks = [
        "加个任务：下周五之前提交项目报告，优先级高",
        "记一下：Python 3.12 支持了 type 语句定义类型别名",
        "再加个任务：周末整理一下书签",
        "看看有什么要做的",
    ]

    for task in tasks:
        print(f"\n你: {task}")
        async for msg in query(prompt=task, options=options):
            if isinstance(msg, AssistantMessage):
                for block in msg.content:
                    if isinstance(block, TextBlock):
                        print(f"🤖: {block.text}")
                    elif isinstance(block, ToolUseBlock):
                        print(f"  🔧 {block.name}")
        print()

anyio.run(main)

mcp__{服务器名}__{工具名}

例如:
  服务器: productivity
  工具: add_todo
  → mcp__productivity__add_todo

# ❌ 程序员视角
@tool("add_todo", "INSERT INTO todos VALUES (...)", ...)

# ✅ AI 视角
@tool("add_todo", "添加一条待办任务", ...)

# ❌ 太简单
return "ok"

# ✅ 有信息量
return f"✅ 已添加任务 #{id}: {content} [高优先级] 截止: 3月20日"

# ❌ 抛异常
raise FileNotFoundError(path)

# ✅ 返回说明
return f"目录不存在: {path}，请确认路径是否正确"