升级FastAPI与Pydantic v2后，openapi_extra子模型未被识别如何修复？

问题背景 #

当前使用版本：

• FastAPI==0.110.1
• Pydantic v2

需求：为FastAPI接口仅提供payload schema，不执行任何Pydantic验证。原有代码在Pydantic v1 + FastAPI==0.98.0环境下可正常运行，但升级后，FastAPI无法识别Cat、Dog、Lizard这些子模型，导致OpenAPI文档中缺失子模型的定义。

原有代码：

from typing import Literal, Annotated
    
import uvicorn
from fastapi import FastAPI
from pydantic import BaseModel, Field, RootModel

app = FastAPI()


def get_openapi_extra(model) -> dict:
    return {
        "requestBody": {
            "content": {"application/json": {"schema": model.schema(ref_template="#/components/schemas/{model}")}}
        }
    }


class Cat(BaseModel):
    pet_type: Literal['cat']
    meows: int


class Dog(BaseModel):
    pet_type: Literal['dog']
    barks: float


class Lizard(BaseModel):
    pet_type: Literal['reptile', 'lizard']
    scales: bool


class Model(RootModel):
    root: Annotated[Cat | Dog | Lizard, Field(..., discriminator='pet_type')]


@app.get("/", openapi_extra=get_openapi_extra(Model))
def index():
    return "ok"

修复方案 #

1. 适配Pydantic v2的Schema生成方法 #

Pydantic v2中BaseModel.schema()已被废弃，需改用model_json_schema()方法，同时ref模板的变量名从{model}改为{name}。

2. 手动注册子模型到OpenAPI组件 #

新版本FastAPI不会自动识别通过openapi_extra传入的模型的关联子模型，需要手动将这些子模型添加到OpenAPI的components/schemas中。

修改后的完整代码

from typing import Literal, Annotated
    
import uvicorn
from fastapi import FastAPI
from pydantic import BaseModel, Field, RootModel

app = FastAPI()


def get_openapi_extra(model, sub_models=None) -> dict:
    # 提前初始化OpenAPI schema结构
    if not hasattr(app, "openapi_schema"):
        app.openapi_schema = {"components": {"schemas": {}}}
    
    # 手动注册子模型到组件
    if sub_models:
        for sub_model in sub_models:
            schema_name = sub_model.__name__
            if schema_name not in app.openapi_schema["components"]["schemas"]:
                app.openapi_schema["components"]["schemas"][schema_name] = sub_model.model_json_schema()
    
    return {
        "requestBody": {
            "content": {"application/json": {"schema": model.model_json_schema(ref_template="#/components/schemas/{name}")}}
        }
    }


class Cat(BaseModel):
    pet_type: Literal['cat']
    meows: int


class Dog(BaseModel):
    pet_type: Literal['dog']
    barks: float


class Lizard(BaseModel):
    pet_type: Literal['reptile', 'lizard']
    scales: bool


class Model(RootModel):
    root: Annotated[Cat | Dog | Lizard, Field(..., discriminator='pet_type')]


# 传入子模型列表，确保它们被注册到OpenAPI
@app.get("/", openapi_extra=get_openapi_extra(Model, sub_models=[Cat, Dog, Lizard]))
def index():
    return "ok"


if __name__ == "__main__":
    uvicorn.run(app)

替代方案：重写OpenAPI生成方法 #

如果需要更统一的管理，可以重写FastAPI的openapi()方法，在生成文档时自动注册所有需要的模型：

from fastapi.openapi.utils import get_openapi

def custom_openapi():
    if app.openapi_schema:
        return app.openapi_schema
    # 生成基础OpenAPI schema
    openapi_schema = get_openapi(
        title="Pet API",
        version="1.0.0",
        routes=app.routes,
    )
    # 注册所有相关模型
    for model in [Cat, Dog, Lizard, Model]:
        schema_name = model.__name__
        openapi_schema["components"]["schemas"][schema_name] = model.model_json_schema()
    app.openapi_schema = openapi_schema
    return app.openapi_schema

app.openapi = custom_openapi

内容的提问来源于stack exchange，提问作者Альберт Александров