Skip to content

Commit

Permalink
Revert "remove middleware"
Browse files Browse the repository at this point in the history
This reverts commit 1dd2c33.
  • Loading branch information
fregataa committed Jan 2, 2024
1 parent 1dd2c33 commit c35a829
Show file tree
Hide file tree
Showing 2 changed files with 10 additions and 3 deletions.
9 changes: 8 additions & 1 deletion src/ai/backend/manager/api/service.py
Original file line number Diff line number Diff line change
Expand Up @@ -61,6 +61,7 @@
from .types import CORSOptions, WebMiddleware
from .utils import (
check_api_params_v2,
convert_response,
get_access_key_scopes,
get_user_uuid_scopes,
undefined,
Expand Down Expand Up @@ -912,6 +913,12 @@ async def clear_error(request: web.Request) -> web.Response:
return web.Response(status=204)


@web.middleware
async def response_middleware(request: web.Request, handler) -> web.StreamResponse:
result = await handler(request)
return convert_response(result)


@attrs.define(slots=True, auto_attribs=True, init=False)
class PrivateContext:
database_ptask_group: aiotools.PersistentTaskGroup
Expand Down Expand Up @@ -950,4 +957,4 @@ def create_app(
cors.add(add_route("PUT", "/{service_id}/routings/{route_id}", update_route))
cors.add(add_route("DELETE", "/{service_id}/routings/{route_id}", delete_route))
cors.add(add_route("POST", "/{service_id}/token", generate_token))
return app, []
return app, [response_middleware]
4 changes: 2 additions & 2 deletions src/ai/backend/manager/api/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -217,12 +217,12 @@ async def wrapped(request: web.Request, *args: P.args, **kwargs: P.kwargs) -> TA

def convert_response(response: TResponseModel | list | TAnyResponse) -> web.StreamResponse:
match response:
case web_response.StreamResponse():
return response
case BaseModel():
return web.json_response(response.model_dump())
case list():
return web.json_response(TypeAdapter(list[TResponseModel]).dump_python(response))
case web_response.StreamResponse():
return response
case _:
raise RuntimeError(f"Unsupported response type ({type(response)})")

Expand Down

0 comments on commit c35a829

Please sign in to comment.