diff --git a/litellm/proxy/middleware/prometheus_auth_middleware.py b/litellm/proxy/middleware/prometheus_auth_middleware.py index d0e38f98946..6f33963c4e1 100644 --- a/litellm/proxy/middleware/prometheus_auth_middleware.py +++ b/litellm/proxy/middleware/prometheus_auth_middleware.py @@ -5,7 +5,7 @@ import json from typing import List from fastapi import Request -from starlette.types import ASGIApp, Receive, Scope, Send +from starlette.types import ASGIApp, Message, Receive, Scope, Send import litellm from litellm.proxy._types import SpecialHeaders @@ -43,9 +43,9 @@ class PrometheusAuthMiddleware: # user_api_key_auth reads the request body, which consumes ASGI `receive`. # Buffer those messages and replay them for the inner app; otherwise a # successful auth would forward an exhausted receive and /metrics hangs. - buffered_messages: List[dict] = [] + buffered_messages: List[Message] = [] - async def receive_for_auth() -> dict: + async def receive_for_auth() -> Message: message = await receive() buffered_messages.append(message) return message @@ -81,7 +81,7 @@ class PrometheusAuthMiddleware: replay_idx = 0 - async def receive_replay() -> dict: + async def receive_replay() -> Message: nonlocal replay_idx if replay_idx < len(buffered_messages): msg = buffered_messages[replay_idx]