Pydantic AI 中 Enum 类型使用的注意事项

在 Pydantic AI 项目中，当开发者尝试使用 Enum 类型作为输出类型时，可能会遇到一些意料之外的问题。本文将通过一个具体的案例，深入分析 Enum 类型在 Pydantic AI 中的正确使用方法，帮助开发者避免常见的陷阱。

问题现象

开发者在使用 Pydantic AI 的 Agent 时，发现当输出类型为简单的 str 类型时，代码运行正常：

from pydantic_ai import Agent

agent = Agent('openai:gpt-4o', output_type=str)
await agent.run('I want to buy a hawai pizza')

但当尝试使用 Enum 类型作为输出类型时，却遇到了错误：

from enum import Enum
from pydantic_ai import Agent

class PizzaKind(str, Enum):
    FUNGI: "fungi"
    HAWAI: "hawai"
    PEPPERONI: "pepperoni"
    
agent = Agent('openai:gpt-4o', output_type=PizzaKind)
await agent.run('I want to buy a hawai pizza')

系统报错提示："Invalid schema for function 'final_result': In context=('properties', 'response'), schema must have a 'type' key."

问题分析

这个问题的根源在于 Enum 类的定义方式不正确。在 Python 中，Enum 的成员值应该使用赋值运算符 = 来定义，而不是类型注解 :。错误的定义方式会导致 Pydantic AI 无法正确解析 Enum 的结构，从而在生成 JSON Schema 时出现问题。

正确用法

正确的 Enum 定义应该如下所示：

from enum import Enum

class PizzaKind(str, Enum):
    FUNGI = "fungi"
    HAWAI = "hawai"
    PEPPERONI = "pepperoni"

这种定义方式明确地为每个枚举成员分配了具体的字符串值，使得 Pydantic AI 能够正确识别和处理 Enum 类型。

技术细节

1. Enum 继承 str：通过继承 str，我们确保枚举值可以被序列化为字符串，这对于 JSON 序列化和 API 交互非常重要。

2. 赋值与类型注解的区别：在 Python 中，: 用于类型注解，而 = 用于赋值。Enum 成员需要具体的值，因此必须使用赋值运算符。

3. Pydantic 的 Schema 生成：Pydantic 在生成 JSON Schema 时，会检查类型的定义方式。正确的 Enum 定义能生成包含完整类型信息的 Schema，而错误的定义会导致 Schema 不完整。

最佳实践

1. 始终使用 = 为 Enum 成员赋值
2. 考虑继承 str 或 int 以使枚举值可序列化
3. 在复杂项目中，建议为 Enum 添加文档字符串说明每个值的含义
4. 测试 Enum 类型在不同场景下的行为，包括序列化和反序列化

总结

在 Pydantic AI 项目中使用 Enum 类型时，开发者需要注意正确的定义方式。通过遵循 Python Enum 的标准用法，可以避免 Schema 生成相关的问题，确保代码的稳定性和可维护性。这个小细节虽然简单，但对于保证整个系统的正常运行至关重要。