《AI Agent 开发入门》2.1 Model 协议设计

# minagent/types.py（承接第 1 章，本节新增 Model）
from __future__ import annotations
from dataclasses import dataclass, field
from typing import Any, Literal, Protocol

Role = Literal["system", "user", "assistant", "tool"]


@dataclass
class ToolCall:
    id: str
    name: str
    arguments: dict[str, Any]


@dataclass
class Message:
    role: Role
    content: str = ""
    tool_calls: list[ToolCall] = field(default_factory=list)
    tool_call_id: str | None = None


@dataclass
class ModelResponse:
    message: Message
    stop_reason: str = "end_turn"      # "end_turn" | "tool_use" | "max_tokens"
    usage: dict[str, int] = field(default_factory=dict)


class Model(Protocol):
    def complete(self, messages: list[Message], tools: list[dict]) -> ModelResponse: ...

# minagent/models.py
from __future__ import annotations
from minagent.types import Message, ModelResponse, ToolCall


class EchoModel:
    """最小可用后端：把最后一条 user 消息原样回显。用于冒烟与协议一致性检查。"""

    def __init__(self) -> None:
        self.calls: list[tuple[list[Message], list[dict]]] = []

    def complete(self, messages: list[Message], tools: list[dict]) -> ModelResponse:
        self.calls.append((list(messages), list(tools)))
        last_user = next((m for m in reversed(messages) if m.role == "user"), None)
        text = last_user.content if last_user else ""
        return ModelResponse(
            message=Message(role="assistant", content=f"echo: {text}"),
            stop_reason="end_turn",
            usage={"input_tokens": 0, "output_tokens": 0},
        )

# minagent/models.py（续）
class ScriptedModel:
    """按剧本返回：每次调用弹出队首的 ModelResponse。剧本耗尽即抛错。"""

    def __init__(self, script: list[ModelResponse]) -> None:
        self._script = list(script)
        self.calls: list[tuple[list[Message], list[dict]]] = []

    def complete(self, messages: list[Message], tools: list[dict]) -> ModelResponse:
        self.calls.append((list(messages), list(tools)))
        if not self._script:
            raise RuntimeError("剧本已耗尽：Agent 未在预期步数内收敛")
        return self._script.pop(0)

# run21.py
import sys
sys.path.insert(0, "/tmp/ab_work")
from minagent.types import Message, ModelResponse, ToolCall
from minagent.models import EchoModel, ScriptedModel


def describe(model) -> str:
    """只依赖 Model 协议，不关心背后是哪家供应商。"""
    reply = model.complete(
        [Message(role="system", content="你是助手"),
         Message(role="user", content="你好，介绍一下你自己")],
        tools=[],
    )
    return f"{reply.message.content!r} stop={reply.stop_reason} usage={reply.usage}"


print("EchoModel    ->", describe(EchoModel()))
script = [
    ModelResponse(message=Message(role="assistant", content="我先查一下",
                                  tool_calls=[ToolCall("t1", "get_time", {"tz": "UTC"})]),
                  stop_reason="tool_use", usage={"input_tokens": 30, "output_tokens": 12}),
    ModelResponse(message=Message(role="assistant", content="现在是 10:00"),
                  stop_reason="end_turn", usage={"input_tokens": 48, "output_tokens": 9}),
]
m = ScriptedModel(script)
print("ScriptedModel->", describe(m))
print("ScriptedModel->", describe(m))
print("calls recorded:", len(m.calls))

EchoModel    -> 'echo: 你好，介绍一下你自己' stop=end_turn usage={'input_tokens': 0, 'output_tokens': 0}
ScriptedModel-> '我先查一下' stop=tool_use usage={'input_tokens': 30, 'output_tokens': 12}
ScriptedModel-> '现在是 10:00' stop=end_turn usage={'input_tokens': 48, 'output_tokens': 9}
calls recorded: 2