# requires: pydantic
"""手写工具调用循环：工具定义（JSON Schema）→ 模型决定调用 → 校验参数 → 执行 → 结果回传 → 直到模型给出答案。

为了离线运行，模型用 ScriptedModel 模拟，它返回的消息格式和 OpenAI Chat Completions 的 tool_calls 一致；
换成真实模型只需要把 ScriptedModel.chat 换成 client.chat.completions.create（见文件末尾注释）。
运行：python tool_loop.py
"""
from __future__ import annotations

import json
from dataclasses import dataclass, field
from typing import Any, Callable

from pydantic import BaseModel, Field, ValidationError


# ---------- 工具：用 pydantic 定义参数，自动生成 JSON Schema ----------
class GetOrderArgs(BaseModel):
    order_id: str = Field(pattern=r"^A\d+$", description="订单号，例如 A1001")


class RefundArgs(BaseModel):
    order_id: str = Field(pattern=r"^A\d+$")
    amount: float = Field(gt=0, le=10000, description="退款金额（元）")


ORDERS = {"A1001": {"user": "u1", "paid": 199.0, "status": "delivered"}}
REFUNDS: dict[str, float] = {}


def get_order(args: GetOrderArgs, ctx: dict) -> dict:
    order = ORDERS.get(args.order_id)
    if not order or order["user"] != ctx["user_id"]:       # 权限在工具里校验，不依赖模型
        return {"error": "订单不存在"}
    return {"order_id": args.order_id, **order}


def refund(args: RefundArgs, ctx: dict) -> dict:
    order = ORDERS.get(args.order_id)
    if not order or order["user"] != ctx["user_id"]:
        return {"error": "订单不存在"}
    if args.amount > order["paid"]:
        return {"error": f"退款金额不能超过实付 {order['paid']}"}
    REFUNDS.setdefault(args.order_id, args.amount)          # 幂等：同一订单只退一次
    return {"ok": True, "refunded": REFUNDS[args.order_id]}


@dataclass
class Tool:
    name: str
    description: str
    args_model: type[BaseModel]
    fn: Callable[[Any, dict], dict]
    needs_confirm: bool = False                              # 有副作用的工具需要用户确认

    def schema(self) -> dict:
        return {"type": "function", "function": {
            "name": self.name, "description": self.description,
            "parameters": self.args_model.model_json_schema()}}


TOOLS = {t.name: t for t in [
    Tool("get_order", "查询当前用户的订单详情", GetOrderArgs, get_order),
    Tool("refund", "为订单发起退款，金额不能超过实付金额", RefundArgs, refund, needs_confirm=True),
]}


# ---------- 模拟模型：按脚本依次返回 assistant 消息 ----------
@dataclass
class ScriptedModel:
    script: list[dict]
    seen: list[list[dict]] = field(default_factory=list)

    def chat(self, messages: list[dict], tools: list[dict]) -> dict:
        self.seen.append([dict(m) for m in messages])
        return self.script.pop(0)


def call(id_: str, name: str, args: dict | str) -> dict:
    return {"id": id_, "type": "function",
            "function": {"name": name, "arguments": args if isinstance(args, str) else json.dumps(args)}}


# ---------- 循环本体 ----------
def run_agent(model, user_msg: str, ctx: dict, confirm: Callable[[str, dict], bool], max_steps: int = 6) -> dict:
    messages: list[dict] = [
        {"role": "system", "content": "你是售后助手。只能通过工具查询和操作订单；信息不足时先询问用户。"},
        {"role": "user", "content": user_msg},
    ]
    trace = []
    for step in range(max_steps):
        msg = model.chat(messages, [t.schema() for t in TOOLS.values()])
        messages.append(msg)
        calls = msg.get("tool_calls") or []
        if not calls:
            return {"answer": msg["content"], "steps": step + 1, "trace": trace}
        for c in calls:                                      # 一次可能返回多个工具调用
            name, raw = c["function"]["name"], c["function"]["arguments"]
            tool = TOOLS.get(name)
            try:
                if tool is None:
                    raise KeyError(f"未知工具 {name}")
                args = tool.args_model.model_validate_json(raw)   # 参数校验：模型给的参数不可信
                if tool.needs_confirm and not confirm(name, args.model_dump()):
                    result = {"error": "用户拒绝执行该操作"}
                else:
                    result = tool.fn(args, ctx)
            except (ValidationError, KeyError, json.JSONDecodeError) as e:
                result = {"error": f"参数错误：{str(e)[:200]}"}     # 错误也作为工具结果返回，让模型自我修正
            trace.append((name, raw, result))
            messages.append({"role": "tool", "tool_call_id": c["id"],
                             "content": json.dumps(result, ensure_ascii=False)})
    return {"answer": "步骤过多，已转人工处理。", "steps": max_steps, "trace": trace}


def demo_script() -> list[dict]:
    return [
        {"role": "assistant", "content": None, "tool_calls": [call("c1", "get_order", {"order_id": "1001"})]},
        {"role": "assistant", "content": None, "tool_calls": [call("c2", "get_order", {"order_id": "A1001"})]},
        {"role": "assistant", "content": None, "tool_calls": [call("c3", "refund", {"order_id": "A1001", "amount": 199})]},
        {"role": "assistant", "content": "已为订单 A1001 退款 199 元。"},
    ]


if __name__ == "__main__":
    model = ScriptedModel(demo_script())
    out = run_agent(model, "订单 1001 我要退款", {"user_id": "u1"}, confirm=lambda n, a: True)
    for name, raw, result in out["trace"]:
        print(f"调用 {name}({raw}) -> {result}")
    print("回答：", out["answer"], "步数：", out["steps"])
    assert "参数错误" in out["trace"][0][2]["error"] and REFUNDS == {"A1001": 199.0}

# 换成真实模型（OpenAI 兼容接口，参数以官方文档为准）：
#   from openai import OpenAI
#   client = OpenAI()
#   resp = client.chat.completions.create(model="...", messages=messages, tools=schemas)
#   msg = resp.choices[0].message.model_dump(exclude_none=True)