如何自定义 FastAPI 请求体文档

How to customize FastAPI request body documentation

我正在使用 FastAPI 来提供 ML 模型。我的端点接收和发送 JSON 形式的数据:

[
  {"id": 1, "data": [{"code": "foo", "value": 0.1}, {"code": "bar", "value": 0.2}, ...]},
  {"id": 2, "data": [{"code": "baz", "value": 0.3}, {"code": "foo", "value": 0.4}, ...]},
  ...
]

我的模型和应用如下所示:

from typing import Dict, List
 
from fastapi import Body
from fastapi.responses import JSONResponse
from pydantic import BaseModel
import pandas as pd


class Item(BaseModel):
    code: str
    value: float


class Sample(BaseModel):
    id: int
    data: List[Item]


app = FastAPI()


@app.post("/score", response_model=List[Sample])  # correct response documentation
def score(input_data: List[Sample] = Body(...)):  # 1. conversion dict -> Pydantic models, slow
    input_df: pd.DataFrame = models_to_df(input_data)  # 2. conversion Pydantic models -> df

    output_df: pd.DataFrame = predict(input_df)

    output_data: Dict = df_to_dict(output_df)  # direct conversion df -> dict, fast
    return JSONResponse(output_data)

一切正常,自动文档看起来不错,但性能很差。由于数据可能非常大,Pydantic 的转换和验证可能会花费很多时间。

这可以通过在 JSON 数据和数据帧之间编写直接转换函数轻松解决,跳过 Pydantic 模型的中间表示。这就是我为响应所做的,实现了 10 倍的加速,同时保留了带有 response_model=List[Sample] 参数的自动 API 文档。

我想通过请求实现相同的目的:能够使用自定义 JSON 输入解析,同时使用 Pydantic 模型保留 API 文档。遗憾的是,我在 FastAPI 文档中找不到执行此操作的方法。我怎样才能做到这一点?

您始终可以接受原始请求,将 request.body() 数据加载为 bytes 并进行您自己的解码。然后应该使用 @app.post() 装饰器的 openapi_extra 参数将请求主体的模式记录为 (partial) raw OpenAPI Operation structure

@app.post(
    "/score",
    response_model=List[Sample],
    openapi_extra={
        "requestBody": {
            "content": {
                "application/json": {
                    "schema": {
                        "type": "array",
                        "items": Sample.schema(ref_template="#/components/schemas/{model}"),
                    }
                }
            }
        }
    },
)
async def score(request: Request):
    raw_body = await request.body()
    # parse the `raw_body` request data (bytes) into your DF directly.

openapi_extra结构合并到其他组件(如response_model)生成的操作结构中。我在此处使用您现有的 Sample 模型来为数组项提供架构,但您也可以手动映射整个架构。

除了使用正文的原始字节,您还可以将解析作为 JSON 委托给请求对象:

data = await request.json()

如果有办法将数据解析为流(将块推送到解析器),您可以通过将请求视为异步循环中的流来避免一次加载整个主体的内存开销:

parser = ...  # something that can be fed chunks of data
async for chunk in request.stream():
    parser.feed(chunk)

这在 Custom OpenAPI path operation schema section in the Advanced User Guide. The same section also covers Us[ing] the Request object directly, and the various options for handling the Request body can be found in the Starlette Request class documentation 中有记录。