From 7e2f2a8ffa0464265e2a50d25562ed3e1e2c80e7 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Mon, 2 Mar 2026 19:41:32 +0530 Subject: [PATCH] Fix inflight mypy --- .../in_flight_requests_middleware.py | 27 ++++++++++--------- tests/test_litellm/proxy/test_proxy_server.py | 9 +++++++ 2 files changed, 24 insertions(+), 12 deletions(-) diff --git a/litellm/proxy/middleware/in_flight_requests_middleware.py b/litellm/proxy/middleware/in_flight_requests_middleware.py index d615640d87..3b93e3a399 100644 --- a/litellm/proxy/middleware/in_flight_requests_middleware.py +++ b/litellm/proxy/middleware/in_flight_requests_middleware.py @@ -6,7 +6,7 @@ Prometheus gauge `litellm_in_flight_requests`. """ import os -from typing import Optional +from typing import Any, Optional from starlette.types import ASGIApp, Receive, Scope, Send @@ -27,7 +27,7 @@ class InFlightRequestsMiddleware: """ _in_flight: int = 0 - _gauge: Optional[object] = None + _gauge: Optional[Any] = None _gauge_init_attempted: bool = False def __init__(self, app: ASGIApp) -> None: @@ -41,13 +41,13 @@ class InFlightRequestsMiddleware: InFlightRequestsMiddleware._in_flight += 1 gauge = InFlightRequestsMiddleware._get_gauge() if gauge is not None: - gauge.inc() # type: ignore[union-attr] + gauge.inc() # type: ignore try: await self.app(scope, receive, send) finally: InFlightRequestsMiddleware._in_flight -= 1 if gauge is not None: - gauge.dec() # type: ignore[union-attr] + gauge.dec() # type: ignore @staticmethod def get_count() -> int: @@ -55,22 +55,25 @@ class InFlightRequestsMiddleware: return InFlightRequestsMiddleware._in_flight @staticmethod - def _get_gauge() -> Optional[object]: + def _get_gauge() -> Optional[Any]: if InFlightRequestsMiddleware._gauge_init_attempted: return InFlightRequestsMiddleware._gauge InFlightRequestsMiddleware._gauge_init_attempted = True try: from prometheus_client import Gauge - kwargs = {} if "PROMETHEUS_MULTIPROC_DIR" in os.environ: # livesum aggregates across all worker processes in the scrape response - kwargs["multiprocess_mode"] = "livesum" - InFlightRequestsMiddleware._gauge = Gauge( - "litellm_in_flight_requests", - "Number of HTTP requests currently in-flight on this uvicorn worker", - **kwargs, - ) + InFlightRequestsMiddleware._gauge = Gauge( + "litellm_in_flight_requests", + "Number of HTTP requests currently in-flight on this uvicorn worker", + multiprocess_mode="livesum", + ) + else: + InFlightRequestsMiddleware._gauge = Gauge( + "litellm_in_flight_requests", + "Number of HTTP requests currently in-flight on this uvicorn worker", + ) except Exception: InFlightRequestsMiddleware._gauge = None return InFlightRequestsMiddleware._gauge diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index b993d4c4cf..112a06b173 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -1758,6 +1758,9 @@ class TestPriceDataReloadAPI: } # Mock the database connection with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma: + mock_prisma.db.litellm_config.find_unique = AsyncMock( + return_value=None + ) mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None) response = client_with_auth.post("/reload/model_cost_map") @@ -1813,6 +1816,9 @@ class TestPriceDataReloadAPI: # Mock the database connection with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma: + mock_prisma.db.litellm_config.find_unique = AsyncMock( + return_value=None + ) mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None) response = client_with_auth.post("/reload/model_cost_map") @@ -2008,6 +2014,9 @@ class TestPriceDataReloadIntegration: # Mock the database connection with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma: + mock_prisma.db.litellm_config.find_unique = AsyncMock( + return_value=None + ) mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None) # Test reload endpoint