Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/middleware/budget_reservation_release_middleware.py: 91%

1from collections.abc import Awaitable, Callable, Mapping 

2from typing import Final 

3 

4from starlette.types import ASGIApp, Receive, Scope, Send 

5 

6_SCOPES_AUTH_STAMPS: Final = frozenset({"http", "websocket"}) 

7 

8 

9class BudgetReservationReleaseMiddleware: 

10 """Releases the budget reservation auth made for a request once no callback owns it. 

11 

12 Auth stamps the reservation on the request or socket state; a call that starts 

13 claims it for the cost callbacks, which settle it on success or failure. When the 

14 response has been sent or the socket has closed and the reservation is still 

15 unclaimed, nothing else ever would, so it is released here instead of pinning the 

16 spend counter until its TTL. 

17 """ 

18 

19 def __init__(self, app: ASGIApp, release: Callable[[Mapping[str, object]], Awaitable[None]]) -> None: 

20 self.app = app 

21 self.release = release 

22 

23 async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: 

24 if scope["type"] not in _SCOPES_AUTH_STAMPS: 

25 await self.app(scope, receive, send) 

26 return 

27 try: 

28 await self.app(scope, receive, send) 

29 finally: 

30 state: Final = scope.get("state") 

31 budget_reservation: Final = state.get("budget_reservation") if isinstance(state, Mapping) else None 

32 if isinstance(budget_reservation, Mapping): 32 ↛ 33line 32 didn't jump to line 33 because the condition on line 32 was never true

33 await self.release(budget_reservation)