搭建一个 AI 智能体循环：工具、记忆与安全迭代

# 带硬性迭代上限的智能体循环。模型 -> tool_calls -> 执行 ->
# 把结果喂回去 -> 重复，直到模型用文字作答，或我们触顶为止。
# 这个上限，正是「智能体」和「失控账单」之间的分水岭。
from openai import OpenAI
import json

client = OpenAI(
    api_key="sk-brievio-...",
    base_url="https://api.brievio.com/v1",
)

MAX_ITERS = 8   # 多数任务 2-4 轮就能完成；8 已是相当宽裕的余量。

def run_agent(question: str, model: str = "claude-sonnet-4-6") -> str:
    messages = [
        {"role": "system", "content": "You are a helpful research agent. "
         "Use the tools when you need live data. Answer directly when you "
         "already know enough."},
        {"role": "user", "content": question},
    ]

    for step in range(MAX_ITERS):
        resp = client.chat.completions.create(
            model=model,
            messages=messages,
            tools=TOOLS,
            tool_choice="auto",
        )
        msg = resp.choices[0].message

        # 没有请求任何工具 -> 这就是最终答案。结束。
        if not msg.tool_calls:
            return msg.content

        # 原封不动地追加这条 assistant 回合 —— 它带着后续消息
        # 必须引用的 tool_call id。
        messages.append(msg)

        # 执行每一个被请求的调用,并为每个 id 追加一条 tool 消息。
        for call in msg.tool_calls:
            result = dispatch(call)              # 见下一段代码
            messages.append({
                "role": "tool",
                "tool_call_id": call.id,          # 必须与该调用的 id 匹配
                "content": json.dumps(result),
            })
        # 循环:模型现在看到了工具输出,继续往下走。

    # 触顶却仍未解决。要响亮地报错 —— 别默默地永远转下去。
    raise RuntimeError(f"agent did not finish within {MAX_ITERS} iterations")

# 带校验的工具分发。模型负责「提议」一个调用;你的代码负责
# 「裁决」它。每一个参数都是不可信的模型输出 —— 先解析它、
# 确认名字是你注册过的、再校验类型,然后才执行。
def get_weather(city: str, unit: str = "celsius") -> dict:
    if not isinstance(city, str) or not city.strip():
        raise ValueError("city must be a non-empty string")
    if unit not in ("celsius", "fahrenheit"):
        raise ValueError(f"unsupported unit: {unit!r}")
    return {"city": city, "temp": 18, "unit": unit, "sky": "clear"}

# 白名单:模型只能调用你显式注册过的东西。
TOOL_IMPLS = {"get_weather": get_weather}

TOOLS = [{
    "type": "function",
    "function": {
        "name": "get_weather",
        "description": "Get current weather for a city.",
        "parameters": {
            "type": "object",
            "properties": {
                "city": {"type": "string"},
                "unit": {"type": "string",
                         "enum": ["celsius", "fahrenheit"]},
            },
            "required": ["city"],
        },
    },
}]

def dispatch(call) -> dict:
    name = call.function.name
    fn = TOOL_IMPLS.get(name)
    if fn is None:
        # 模型幻想出了一个工具。别让循环崩掉 —— 把错误作为
        # 工具结果交回去,让模型自己纠正。
        return {"error": f"unknown tool: {name}"}

    try:
        args = json.loads(call.function.arguments)   # 永远是一个 JSON 字符串
    except json.JSONDecodeError:
        return {"error": "arguments were not valid JSON"}

    try:
        return fn(**args)
    except (TypeError, ValueError) as e:
        # 参数有误(类型不对、缺字段、超范围)。把消息喂回去;
        # 模型通常会用修正后的调用重试。
        return {"error": str(e)}

# 单次运行的成本/token 预算护栏。每一轮循环都会重新发送一份
# 不断增长的上下文(历史 + 工具输出),所以成本逐轮攀升。每次调用后
# 读取诚实的 usage 对象,给它定价,一旦本次运行超出预算就停下来
# —— 这与迭代上限相互独立。
from decimal import Decimal

# Brievio 公开费率,美元/每百万 token(约比官方参考价低 15%)。
RATES = {
    "claude-sonnet-4-6": {"in": Decimal("2.55"), "out": Decimal("12.75")},
    "claude-haiku-4-5":  {"in": Decimal("0.85"), "out": Decimal("4.25")},
}

def call_cost(model: str, usage) -> Decimal:
    r = RATES[model]
    m = Decimal("1000000")
    return usage.prompt_tokens * r["in"] / m + usage.completion_tokens * r["out"] / m

RUN_BUDGET = Decimal("0.10")   # 每次智能体运行 10 美分,硬性上限。

def run_agent_budgeted(question: str, model: str = "claude-sonnet-4-6") -> str:
    messages = [{"role": "user", "content": question}]
    spent = Decimal("0")

    for step in range(MAX_ITERS):
        resp = client.chat.completions.create(
            model=model, messages=messages, tools=TOOLS, tool_choice="auto",
        )
        spent += call_cost(model, resp.usage)   # 每一轮都累计真实 token
        if spent > RUN_BUDGET:
            raise RuntimeError(f"run exceeded ${RUN_BUDGET} (spent ${spent:.4f})")

        msg = resp.choices[0].message
        if not msg.tool_calls:
            return msg.content

        messages.append(msg)
        for call in msg.tool_calls:
            messages.append({"role": "tool", "tool_call_id": call.id,
                             "content": json.dumps(dispatch(call))})

    raise RuntimeError(f"agent did not finish within {MAX_ITERS} iterations")