Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/spend_tracking/carried_budget_state.py: 76%

1"""Pins the budget state auth resolved onto ``UserAPIKeyAuth`` and emits it as request metadata.""" 

2 

3from collections.abc import Mapping 

4from types import MappingProxyType 

5from typing import Final 

6 

7from litellm.models.organization import LiteLLM_OrganizationTable 

8from litellm.models.team import LiteLLM_TeamTable 

9from litellm.models.user import LiteLLM_UserTable 

10from litellm.proxy._types import UserAPIKeyAuth 

11from litellm.types.proxy.carried_budget_state import ( 

12 OrgBudgetSnapshot, 

13 TeamBudgetSnapshot, 

14 UserBudgetSnapshot, 

15) 

16 

17 

18def carry_team_and_user_budget_state( 

19 valid_token: UserAPIKeyAuth, 

20 team_object: LiteLLM_TeamTable | None, 

21 user_object: LiteLLM_UserTable | None, 

22) -> None: 

23 if team_object is not None: 23 ↛ 24line 23 didn't jump to line 24 because the condition on line 23 was never true

24 valid_token.team_budget_snapshot = TeamBudgetSnapshot( # rebind-ok: the request credential is pinned in place 

25 budget_reset_at=team_object.budget_reset_at, 

26 max_budget=team_object.max_budget, 

27 ) 

28 valid_token.team_model_max_budget = team_object.model_max_budget # rebind-ok: caller keeps this object 

29 if user_object is not None: 

30 valid_token.user_budget_snapshot = UserBudgetSnapshot( # rebind-ok: same object the caller keeps using 

31 budget_reset_at=user_object.budget_reset_at, 

32 max_budget=user_object.max_budget, 

33 user_alias=user_object.user_alias, 

34 ) 

35 

36 

37def carry_organization_budget_state(valid_token: UserAPIKeyAuth, org_table: LiteLLM_OrganizationTable) -> None: 

38 budget_table: Final = org_table.litellm_budget_table 

39 valid_token.organization_alias = org_table.organization_alias 

40 valid_token.org_budget_snapshot = OrgBudgetSnapshot( # rebind-ok: same object the caller keeps using 

41 spend=org_table.spend, 

42 max_budget=budget_table.max_budget if budget_table is not None else None, 

43 ) 

44 

45 

46def carried_budget_metadata(valid_token: UserAPIKeyAuth) -> Mapping[str, object]: 

47 snapshots: Final = ( 

48 valid_token.team_budget_snapshot, 

49 valid_token.user_budget_snapshot, 

50 valid_token.org_budget_snapshot, 

51 ) 

52 return MappingProxyType( 

53 { 

54 key: value 

55 for snapshot in snapshots 

56 if snapshot is not None 

57 for key, value in snapshot.metadata_entries().items() 

58 } 

59 )