编程助手

"""
编程助手:Claude 读文件、写文件、跑测试。
依赖:pip install anthropic
"""
import os, subprocess, pathlib
from anthropic import Anthropic

REPO_ROOT = pathlib.Path(".").resolve()

client = Anthropic(
    base_url="https://buzzai.cc",
    api_key=os.environ["BUZZ_API_KEY"],
)

SYSTEM_PROMPT = """你是用户仓库内的资深软件工程师。
写之前先读。改完做有意义的编辑后跑测试。diff 越小越好。"""

TOOLS = [
    {
        "name": "read_file",
        "description": "读取项目内 UTF-8 文本文件。",
        "input_schema": {
            "type": "object",
            "properties": {
                "path": {"type": "string"},
                "start_line": {"type": "integer", "minimum": 1},
                "end_line": {"type": "integer", "minimum": 1},
            },
            "required": ["path"],
        },
    },
    {
        "name": "write_file",
        "description": "覆盖写 UTF-8 文件,返回字节数。",
        "input_schema": {
            "type": "object",
            "properties": {"path": {"type": "string"}, "content": {"type": "string"}},
            "required": ["path", "content"],
        },
    },
    {
        "name": "run_tests",
        "description": "跑项目测试,可选 pattern 缩小范围。",
        "input_schema": {
            "type": "object",
            "properties": {"pattern": {"type": "string"}},
        },
    },
]


def safe_path(p: str) -> pathlib.Path:
    full = (REPO_ROOT / p).resolve()
    if REPO_ROOT not in full.parents and full != REPO_ROOT:
        raise ValueError(f"路径越界: {p}")
    return full


def execute_tool(name: str, args: dict) -> str:
    if name == "read_file":
        text = safe_path(args["path"]).read_text()
        if "start_line" in args or "end_line" in args:
            lines = text.splitlines()
            s = args.get("start_line", 1) - 1
            e = args.get("end_line", len(lines))
            text = "\n".join(lines[s:e])
        return text
    if name == "write_file":
        path = safe_path(args["path"])
        path.parent.mkdir(parents=True, exist_ok=True)
        path.write_text(args["content"])
        return f"wrote {path.stat().st_size} bytes"
    if name == "run_tests":
        cmd = ["pytest", "-x", "-q"]
        if args.get("pattern"):
            cmd.append(args["pattern"])
        out = subprocess.run(cmd, capture_output=True, text=True, cwd=REPO_ROOT)
        return f"exit={out.returncode}\n{out.stdout}\n{out.stderr}"
    raise ValueError(f"未知工具: {name}")


def build_system():
    readme = (REPO_ROOT / "README.md").read_text() if (REPO_ROOT / "README.md").exists() else ""
    return [
        {"type": "text", "text": SYSTEM_PROMPT},
        {
            "type": "text",
            "text": f"<PROJECT_README>\n{readme}\n</PROJECT_README>",
            "cache_control": {"type": "ephemeral"},
        },
    ]


def chat(user_request: str, max_turns: int = 12):
    messages = [{"role": "user", "content": user_request}]
    for turn in range(max_turns):
        resp = client.messages.create(
            model="claude-sonnet-4-6",
            max_tokens=4096,
            system=build_system(),
            tools=TOOLS,
            messages=messages,
        )
        messages.append({"role": "assistant", "content": resp.content})

        if resp.stop_reason != "tool_use":
            for block in resp.content:
                if block.type == "text":
                    print(block.text)
            return

        tool_results = []
        for block in resp.content:
            if block.type == "tool_use":
                try:
                    result = execute_tool(block.name, block.input)
                    tool_results.append({
                        "type": "tool_result",
                        "tool_use_id": block.id,
                        "content": result,
                    })
                except Exception as e:
                    tool_results.append({
                        "type": "tool_result",
                        "tool_use_id": block.id,
                        "content": f"ERROR: {e}",
                        "is_error": True,
                    })
        messages.append({"role": "user", "content": tool_results})
    print("到达 max_turns 仍未完成")


if __name__ == "__main__":
    chat("给 app/server.py 加一个 /healthz 路由,补一个单元测试。")