chore: Allow pydantic.BaseModel as API handler return schema (#1987)

lablup · Apr 4, 2024 · 652c1aa · 652c1aa
1 parent e15b21e
commit 652c1aa
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/changes/1987.misc.md b/changes/1987.misc.md
@@ -0,0 +1 @@
+Allow `pydantic.BaseModel` as the API handler return schema.
diff --git a/src/ai/backend/manager/api/utils.py b/src/ai/backend/manager/api/utils.py
@@ -214,12 +214,12 @@ async def wrapped(request: web.Request, *args: P.args, **kwargs: P.kwargs) -> TA
 
 
 class BaseResponseModel(BaseModel):
-    status: Annotated[int, Field(strict=True, ge=100, lt=600)] = 200
+    status: Annotated[int, Field(strict=True, exclude=True, ge=100, lt=600)] = 200
 
 
 TParamModel = TypeVar("TParamModel", bound=BaseModel)
 TQueryModel = TypeVar("TQueryModel", bound=BaseModel)
-TResponseModel = TypeVar("TResponseModel", bound=BaseResponseModel)
+TResponseModel = TypeVar("TResponseModel", bound=BaseModel)
 
 TPydanticResponse: TypeAlias = TResponseModel | list
 THandlerFuncWithoutParam: TypeAlias = Callable[
@@ -231,11 +231,13 @@ class BaseResponseModel(BaseModel):
 
 
 def ensure_stream_response_type(
-    response: BaseResponseModel | list[TResponseModel] | web.StreamResponse,
+    response: BaseResponseModel | BaseModel | list[TResponseModel] | web.StreamResponse,
 ) -> web.StreamResponse:
     match response:
         case BaseResponseModel(status=status):
             return web.json_response(response.model_dump(mode="json"), status=status)
+        case BaseModel():
+            return web.json_response(response.model_dump(mode="json"))
         case list():
             return web.json_response(TypeAdapter(type(response)).dump_python(response, mode="json"))
         case web.StreamResponse():