tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-02 14:22:16 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fix inflight mypy

Browse Source
This commit is contained in:
Sameer Kankute
2026-03-02 19:41:32 +05:30
parent 3007010f21
commit 7e2f2a8ffa
2 changed files with 24 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/proxy/middleware/in_flight_requests_middleware.py
+15 -12
View File
@@ -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
tests/test_litellm/proxy/test_proxy_server.py
+9
View File
@@ -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