AI Agent 入门：从概念到第一个可运行闭环

from __future__ import annotations

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

from openai import OpenAI

client = OpenAI()


def get_weather(city: str) -> dict[str, str]:
    sample = {
        "北京": {"condition": "晴", "temperature": "25C"},
        "上海": {"condition": "多云", "temperature": "23C"},
    }
    return sample.get(city, {"condition": "未知", "temperature": "未知"})


TOOLS = [
    {
        "type": "function",
        "function": {
            "name": "get_weather",
            "description": "查询指定城市天气",
            "parameters": {
                "type": "object",
                "properties": {
                    "city": {"type": "string", "description": "城市名称"}
                },
                "required": ["city"],
            },
        },
    }
]


@dataclass
class AgentState:
    user_query: str
    messages: list[dict[str, Any]] = field(default_factory=list)
    steps: int = 0
    finished: bool = False


def run_agent(query: str, max_steps: int = 4) -> str:
    state = AgentState(
        user_query=query,
        messages=[
            {
                "role": "system",
                "content": (
                    "你是一个任务助手。需要实时信息时调用工具。"
                    "如果信息不足，不要猜。"
                ),
            },
            {"role": "user", "content": query},
        ],
    )

    while not state.finished and state.steps < max_steps:
        state.steps += 1

        response = client.chat.completions.create(
            model="gpt-4o-mini",
            messages=state.messages,
            tools=TOOLS,
        )
        message = response.choices[0].message

        if not message.tool_calls:
            state.finished = True
            return message.content or "未生成结果"

        state.messages.append(message.model_dump())

        for tool_call in message.tool_calls:
            if tool_call.function.name != "get_weather":
                continue

            args = json.loads(tool_call.function.arguments)
            result = get_weather(args["city"])

            state.messages.append(
                {
                    "role": "tool",
                    "tool_call_id": tool_call.id,
                    "content": json.dumps(result, ensure_ascii=False),
                }
            )

    return "任务未在限定步数内完成，请转人工处理。"