从「能聊」到「能办事」，中间隔着权限与可观测性

"""智能体工具调用循环。

设计要点：
    1. 工具分级：只读工具直接执行，写操作工具需人工审批
    2. 幂等键：每次写操作带幂等键，防止重试造成重复执行
    3. 步数上限：防止模型陷入死循环烧钱
    4. 全链路审计：每步的输入输出都落库
"""

from __future__ import annotations

import json
import uuid
from typing import Any, Callable

import boto3

BEDROCK = boto3.client("bedrock-runtime", region_name="ap-northeast-1")
MODEL = "anthropic.claude-3-5-sonnet-20241022-v2:0"
MAX_TURNS = 8

# ---------- 工具注册表 ----------

TOOL_SPECS = [
    {
        "toolSpec": {
            "name": "query_monthly_cost",
            "description": "查询指定月份、指定维度的 AWS 成本。只读操作。",
            "inputSchema": {"json": {
                "type": "object",
                "properties": {
                    "month": {"type": "string", "description": "格式 YYYY-MM"},
                    "group_by": {"type": "string", "enum": ["SERVICE", "ACCOUNT", "TAG"]},
                },
                "required": ["month", "group_by"],
            }},
        }
    },
    {
        "toolSpec": {
            "name": "create_optimization_ticket",
            "description": "创建一张成本优化工单。写操作，需人工审批。",
            "inputSchema": {"json": {
                "type": "object",
                "properties": {
                    "title": {"type": "string"},
                    "detail": {"type": "string"},
                    "estimated_saving": {"type": "number"},
                    "owner": {"type": "string"},
                },
                "required": ["title", "detail", "owner"],
            }},
        }
    },
]

READ_ONLY_TOOLS = {"query_monthly_cost"}


def execute_tool(
    name: str,
    args: dict[str, Any],
    registry: dict[str, Callable],
    approve: Callable[[str, dict], bool],
    idem_key: str,
) -> dict[str, Any]:
    """执行工具，写操作先过审批。"""
    if name not in registry:
        return {"error": f"未注册的工具：{name}"}

    if name not in READ_ONLY_TOOLS:
        if not approve(name, args):
            return {"error": "人工审批未通过，操作已取消", "requires_approval": True}
        args = {**args, "idempotency_key": idem_key}

    try:
        return {"result": registry[name](**args)}
    except Exception as exc:  # noqa: BLE001 - 工具异常需回传给模型
        return {"error": f"{type(exc).__name__}: {exc}"}


def run_agent(
    user_input: str,
    registry: dict[str, Callable],
    approve: Callable[[str, dict], bool],
    audit: Callable[[dict], None],
) -> dict[str, Any]:
    """主循环。返回最终答复与完整调用轨迹。"""
    messages: list[dict[str, Any]] = [
        {"role": "user", "content": [{"text": user_input}]}
    ]
    trace: list[dict[str, Any]] = []
    session_id = str(uuid.uuid4())

    for turn in range(MAX_TURNS):
        resp = BEDROCK.converse(
            modelId=MODEL,
            messages=messages,
            toolConfig={"tools": TOOL_SPECS},
            inferenceConfig={"maxTokens": 2000, "temperature": 0.2},
            system=[{"text": (
                "你是云成本运营助手。先查数据再给结论，"
                "涉及创建工单等写操作时，必须先说明理由和预估收益。"
            )}],
        )

        msg = resp["output"]["message"]
        messages.append(msg)
        stop = resp["stopReason"]

        audit({
            "session": session_id, "turn": turn,
            "stop_reason": stop, "usage": resp["usage"],
        })

        if stop != "tool_use":
            text =".join(b.get("text",") for b in msg["content"])
            return {"answer": text, "trace": trace, "turns": turn + 1}

        tool_results = []
        for block in msg["content"]:
            if "toolUse" not in block:
                continue
            tu = block["toolUse"]
            idem = f"{session_id}:{turn}:{tu['name']}"
            out = execute_tool(tu["name"], tu["input"], registry, approve, idem)
            trace.append({"tool": tu["name"], "input": tu["input"], "output": out})
            tool_results.append({"toolResult": {
                "toolUseId": tu["toolUseId"],
                "content": [{"json": out}],
                "status": "error" if "error" in out else "success",
            }})

        messages.append({"role": "user", "content": tool_results})

    return {"answer": "达到最大步数上限，任务未完成。", "trace": trace, "turns": MAX_TURNS}