Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/common_utils/proxy_state.py: 90%

1""" 

2This file is used to store the state variables of the proxy server. 

3 

4Example: `spend_logs_row_count` is used to store the number of rows in the `LiteLLM_SpendLogs` table. 

5""" 

6 

7from typing import Any, Literal 

8 

9from litellm.proxy._types import ProxyStateVariables 

10 

11 

12class ProxyState: 

13 """ 

14 Proxy state class has get/set methods for Proxy state variables. 

15 """ 

16 

17 # Note: mypy does not recognize when we fetch ProxyStateVariables.annotations.keys(), so we also need to add the valid keys here 

18 valid_keys_literal = Literal["spend_logs_row_count"] 

19 

20 def __init__(self) -> None: 

21 self.proxy_state_variables: ProxyStateVariables = ProxyStateVariables( 

22 spend_logs_row_count=0, 

23 ) 

24 

25 def get_proxy_state_variable( 

26 self, 

27 variable_name: valid_keys_literal, 

28 ) -> Any: 

29 return self.proxy_state_variables.get(variable_name, None) 

30 

31 def set_proxy_state_variable( 

32 self, 

33 variable_name: valid_keys_literal, 

34 value: Any, 

35 ) -> None: 

36 self.proxy_state_variables[variable_name] = value