138 lines
6.0 KiB
Python
138 lines
6.0 KiB
Python
# 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)
|