python - 如何自定义 FastAPI 请求体文档

Question

我正在使用 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 文档中找到执行此操作的方法。我怎样才能做到这一点？

Answer 1

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

@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 委托(delegate)给请求对象:

data = await request.json()

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

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

这记录在 Custom OpenAPI path operation schema section 中在高级用户指南中。同一部分还包括 Us[ing] the Request object directly ，以及用于处理 Request 正文的各种选项可以在 Starlette Request class documentation 中找到.