跳转至

Pydantic 数据模型

演示 Pydantic v2 的核心能力:模型定义、字段校验、自定义校验器、 模型序列化与嵌套模型。Agent 开发中常用于定义工具参数、消息结构等。 对应教程:第2卷《Python 进阶》第4章 FastAPI 与 Pydantic。

文件路径: examples/volume-02-python\pydantic_models.py


源代码

"""Pydantic 数据模型示例。

演示 Pydantic v2 的核心能力:模型定义、字段校验、自定义校验器、
模型序列化与嵌套模型。Agent 开发中常用于定义工具参数、消息结构等。

对应教程:第2卷《Python 进阶》第4章 FastAPI 与 Pydantic。
"""

from __future__ import annotations

from datetime import datetime, timezone
from enum import Enum
from typing import Annotated

from pydantic import BaseModel, Field, field_validator, model_validator


class RoleEnum(str, Enum):
    """消息角色枚举。"""

    SYSTEM = "system"
    USER = "user"
    ASSISTANT = "assistant"
    TOOL = "tool"


class Message(BaseModel):
    """单条聊天消息。

    Attributes:
        role: 消息角色。
        content: 消息内容。
    """

    role: RoleEnum
    content: str = Field(min_length=1, description="消息内容,不能为空")

    @field_validator("content", mode="before")
    @classmethod
    def _strip_content(cls, v: object) -> object:
        # 在字段约束(min_length)校验前去空白,确保纯空格内容被拒绝
        return v.strip() if isinstance(v, str) else v

    model_config = {"use_enum_values": True}


class ToolCallRequest(BaseModel):
    """工具调用请求。

    Attributes:
        name: 工具名称。
        arguments: 工具参数(任意 JSON 对象)。
    """

    name: str = Field(pattern=r"^[a-zA-Z_][a-zA-Z0-9_]*$", description="合法的工具名")
    arguments: dict[str, object] = Field(default_factory=dict)

    @model_validator(mode="after")
    def _check_name_not_empty(self) -> ToolCallRequest:
        if not self.name:
            raise ValueError("工具名不能为空")
        return self


class AgentConfig(BaseModel):
    """Agent 运行配置,演示复合校验。

    Attributes:
        model: 模型名称。
        temperature: 采样温度。
        max_tokens: 最大输出 token。
        top_p: nucleus sampling 参数。
    """

    model: str = "gpt-4o-mini"
    temperature: Annotated[float, Field(ge=0.0, le=2.0)] = 0.7
    max_tokens: int = Field(default=1024, ge=1, le=32_768)
    top_p: float = Field(default=1.0, ge=0.0, le=1.0)

    @model_validator(mode="after")
    def _warn_high_temp_low_top_p(self) -> AgentConfig:
        # 同时高 temperature 与低 top_p 会导致采样不稳定
        if self.temperature > 1.5 and self.top_p < 0.5:
            raise ValueError("高 temperature 与低 top_p 不可同时使用")
        return self


class ChatLog(BaseModel):
    """嵌套模型示例:一次对话日志。

    Attributes:
        session_id: 会话 ID。
        messages: 消息列表。
        created_at: 创建时间。
    """

    session_id: str
    messages: list[Message] = Field(default_factory=list)
    created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))

    def total_content_length(self) -> int:
        """统计所有消息内容的总长度。

        Returns:
            内容字符数总和。
        """
        return sum(len(m.content) for m in self.messages)


def build_default_chat_log() -> ChatLog:
    """构造一个示例对话日志。

    Returns:
        包含系统与用户消息的 ``ChatLog``。
    """
    return ChatLog(
        session_id="demo-001",
        messages=[
            Message(role=RoleEnum.SYSTEM, content="你是一个友好的助手。"),
            Message(role=RoleEnum.USER, content="你好"),
        ],
    )


def main() -> None:
    """演示 Pydantic 模型的使用。"""
    log = build_default_chat_log()
    print(log.model_dump_json(indent=2))
    print(f"消息总长度: {log.total_content_length()}")

    config = AgentConfig(temperature=0.3, max_tokens=512)
    print(f"\n配置: {config.model_dump()}")


if __name__ == "__main__":
    main()