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

1""" 

2This module is responsible for handling Getting/Setting the proxy server request from cold storage. 

3 

4It allows fetching a dict of the proxy server request from s3 or GCS bucket. 

5""" 

6 

7from typing import Final 

8 

9import litellm 

10from litellm import _custom_logger_compatible_callbacks_literal 

11from litellm.integrations.custom_logger import CustomLogger 

12 

13 

14class ColdStorageHandler: 

15 """ 

16 This class is responsible for handling Getting/Setting the proxy server request from cold storage. 

17 

18 It allows fetching a dict of the proxy server request from s3 or GCS bucket. 

19 

20 The cold storage logger can be injected for testing; when omitted it is 

21 resolved from the configured ``litellm.cold_storage_custom_logger``. 

22 """ 

23 

24 def __init__(self, cold_storage_logger: CustomLogger | None = None): 

25 self._injected_cold_storage_logger = cold_storage_logger 

26 

27 async def get_proxy_server_request_from_cold_storage_with_object_key( 

28 self, 

29 object_key: str, 

30 ) -> dict | None: 

31 """ 

32 Get the proxy server request from cold storage using the object key directly. 

33 

34 Args: 

35 object_key: The S3/GCS object key to retrieve 

36 

37 Returns: 

38 Optional[dict]: The proxy server request dict or None if not found 

39 """ 

40 custom_logger: Final = self._injected_cold_storage_logger or self._resolve_cold_storage_logger() 

41 if custom_logger is None: 

42 return None 

43 

44 return await custom_logger.get_proxy_server_request_from_cold_storage_with_object_key( 

45 object_key=object_key, 

46 ) 

47 

48 def _resolve_cold_storage_logger(self) -> CustomLogger | None: 

49 custom_logger_name: Final = self._select_custom_logger_for_cold_storage() 

50 if custom_logger_name is None: 

51 return None 

52 return litellm.logging_callback_manager.get_active_custom_logger_for_callback_name(custom_logger_name) 

53 

54 def _select_custom_logger_for_cold_storage( 

55 self, 

56 ) -> _custom_logger_compatible_callbacks_literal | None: 

57 cold_storage_custom_logger: Final[_custom_logger_compatible_callbacks_literal | None] = ( 

58 litellm.cold_storage_custom_logger 

59 ) 

60 

61 return cold_storage_custom_logger