scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-05 08:07:05 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 137

chore: create shared TimedeltaJSONEncoder in json_utils.py to avoid code duplication

Browse source
This commit is contained in:
swarnabhasinha 2025-09-20 15:09:32 +05:30
parent 47efc7f846
commit 84132fa1ca
No known key found for this signature in database
GPG key ID: 9B0C3EA24F21722E
6 changed files with 84 additions and 55 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

23
litellm/caching/azure_blob_cache.py
View file

@ -11,20 +11,11 @@ Has 4 methods:
import asyncio
import json
from contextlib import suppress
from datetime import timedelta
from litellm._logging import print_verbose, verbose_logger
from .base_cache import BaseCache
class TimedeltaJSONEncoder(json.JSONEncoder):
"""Custom JSON encoder that handles timedelta objects by converting them to seconds."""
def default(self, obj):
if isinstance(obj, timedelta):
return obj.total_seconds()
return super().default(obj)
from .json_utils import TimedeltaJSONEncoder
class AzureBlobCache(BaseCache):
@ -32,7 +23,9 @@ class AzureBlobCache(BaseCache):
from azure.storage.blob import BlobServiceClient
from azure.core.exceptions import ResourceExistsError
from azure.identity import DefaultAzureCredential
from azure.identity.aio import DefaultAzureCredential as AsyncDefaultAzureCredential
from azure.identity.aio import (
DefaultAzureCredential as AsyncDefaultAzureCredential,
)
from azure.storage.blob.aio import BlobServiceClient as AsyncBlobServiceClient
self.container_client = BlobServiceClient(
@ -60,14 +53,16 @@ class AzureBlobCache(BaseCache):
print_verbose(f"LiteLLM SET Cache - Azure Blob. Key={key}. Value={value}")
serialized_value = json.dumps(value, cls=TimedeltaJSONEncoder)
try:
await self.async_container_client.upload_blob(key, serialized_value, overwrite=True)
await self.async_container_client.upload_blob(
key, serialized_value, overwrite=True
)
except Exception as e:
# NON blocking - notify users Azure Blob is throwing an exception
print_verbose(f"LiteLLM set_cache() - Got exception from Azure Blob: {e}")
def get_cache(self, key, **kwargs):
from azure.core.exceptions import ResourceNotFoundError
try:
print_verbose(f"Get Azure Blob Cache: key: {key}")
as_bytes = self.container_client.download_blob(key).readall()
@ -84,7 +79,7 @@ class AzureBlobCache(BaseCache):
async def async_get_cache(self, key, **kwargs):
from azure.core.exceptions import ResourceNotFoundError
try:
print_verbose(f"Get Azure Blob Cache: key: {key}")
blob = await self.async_container_client.download_blob(key)

39
litellm/caching/gcs_cache.py
View file

@ -3,7 +3,6 @@ Supports syncing responses to Google Cloud Storage Buckets using HTTP requests.
"""
import json
import asyncio
from datetime import timedelta
from typing import Optional
from litellm._logging import print_verbose, verbose_logger
@ -14,25 +13,27 @@ from litellm.llms.custom_httpx.http_handler import (
httpxSpecialProvider,
)
from .base_cache import BaseCache
class TimedeltaJSONEncoder(json.JSONEncoder):
"""Custom JSON encoder that handles timedelta objects by converting them to seconds."""
def default(self, obj):
if isinstance(obj, timedelta):
return obj.total_seconds()
return super().default(obj)
from .json_utils import TimedeltaJSONEncoder
class GCSCache(BaseCache):
def __init__(self, bucket_name: Optional[str] = None, path_service_account: Optional[str] = None, gcs_path: Optional[str] = None) -> None:
def __init__(
self,
bucket_name: Optional[str] = None,
path_service_account: Optional[str] = None,
gcs_path: Optional[str] = None,
) -> None:
super().__init__()
self.bucket_name = bucket_name or GCSBucketBase(bucket_name=None).BUCKET_NAME
self.path_service_account = path_service_account or GCSBucketBase(bucket_name=None).path_service_account_json
self.path_service_account = (
path_service_account
or GCSBucketBase(bucket_name=None).path_service_account_json
)
self.key_prefix = gcs_path.rstrip("/") + "/" if gcs_path else ""
# create httpx clients
self.async_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback)
self.async_client = get_async_httpx_client(
llm_provider=httpxSpecialProvider.LoggingCallback
)
self.sync_client = _get_httpx_client()
def _construct_headers(self) -> dict:
@ -62,7 +63,9 @@ class GCSCache(BaseCache):
data = json.dumps(value, cls=TimedeltaJSONEncoder)
await self.async_client.post(url=url, data=data, headers=headers)
except Exception as e:
print_verbose(f"GCS Caching: async_set_cache() - Got exception from GCS: {e}")
print_verbose(
f"GCS Caching: async_set_cache() - Got exception from GCS: {e}"
)
def get_cache(self, key, **kwargs):
try:
@ -79,7 +82,9 @@ class GCSCache(BaseCache):
return cached_response
return None
except Exception as e:
verbose_logger.error(f"GCS Caching: get_cache() - Got exception from GCS: {e}")
verbose_logger.error(
f"GCS Caching: get_cache() - Got exception from GCS: {e}"
)
async def async_get_cache(self, key, **kwargs):
try:
@ -92,7 +97,9 @@ class GCSCache(BaseCache):
return json.loads(response.text)
return None
except Exception as e:
verbose_logger.error(f"GCS Caching: async_get_cache() - Got exception from GCS: {e}")
verbose_logger.error(
f"GCS Caching: async_get_cache() - Got exception from GCS: {e}"
)
def flush_cache(self):
pass

40
litellm/caching/json_utils.py Normal file
View file

@ -0,0 +1,40 @@
"""
JSON utilities for caching implementations.
This module provides shared JSON encoding functionality across all cache implementations.
"""
import json
from datetime import timedelta
from typing import Any
class TimedeltaJSONEncoder(json.JSONEncoder):
"""
Custom JSON encoder that handles timedelta objects by converting them to seconds.
This encoder is used across all cache implementations (Redis, S3, GCS, Azure Blob)
to prevent 'Object of type timedelta is not JSON serializable' errors when
caching metrics that contain timedelta objects.
Example:
>>> import json
>>> from datetime import timedelta
>>> data = {"latency": [timedelta(seconds=1.5)]}
>>> json.dumps(data, cls=TimedeltaJSONEncoder)
'{"latency": [1.5]}'
"""
def default(self, obj: Any) -> Any:
"""
Convert timedelta objects to seconds (float) for JSON serialization.
Args:
obj: Object to serialize
Returns:
Serializable representation of the object
"""
if isinstance(obj, timedelta):
return obj.total_seconds()
return super().default(obj)

12
litellm/caching/redis_cache.py
View file

@ -23,15 +23,7 @@ from litellm.types.caching import RedisPipelineIncrementOperation
from litellm.types.services import ServiceTypes
from .base_cache import BaseCache
class TimedeltaJSONEncoder(json.JSONEncoder):
"""Custom JSON encoder that handles timedelta objects by converting them to seconds."""
def default(self, obj):
if isinstance(obj, timedelta):
return obj.total_seconds()
return super().default(obj)
from .json_utils import TimedeltaJSONEncoder
if TYPE_CHECKING:
@ -225,7 +217,7 @@ class RedisCache(BaseCache):
try:
start_time = time.time()
# Convert value to JSON string to handle complex objects like timedelta
if isinstance(value, (dict, list)) or hasattr(value, '__dict__'):
if isinstance(value, (dict, list)) or hasattr(value, "__dict__"):
serialized_value = json.dumps(value, cls=TimedeltaJSONEncoder)
else:
serialized_value = str(value)

16
litellm/caching/s3_cache.py
View file

@ -18,15 +18,7 @@ from datetime import datetime, timezone, timedelta
from litellm._logging import print_verbose, verbose_logger
from .base_cache import BaseCache
class TimedeltaJSONEncoder(json.JSONEncoder):
"""Custom JSON encoder that handles timedelta objects by converting them to seconds."""
def default(self, obj):
if isinstance(obj, timedelta):
return obj.total_seconds()
return super().default(obj)
from .json_utils import TimedeltaJSONEncoder
class S3Cache(BaseCache):
@ -119,7 +111,9 @@ class S3Cache(BaseCache):
func = partial(self.set_cache, key, value, **kwargs)
await loop.run_in_executor(None, func)
except Exception as e:
verbose_logger.error(f"S3 Caching: async_set_cache() - Got exception from S3: {e}")
verbose_logger.error(
f"S3 Caching: async_set_cache() - Got exception from S3: {e}"
)
def get_cache(self, key, **kwargs):
import botocore
@ -135,7 +129,7 @@ class S3Cache(BaseCache):
if cached_response is not None:
if "Expires" in cached_response:
expires_time = cached_response['Expires']
expires_time = cached_response["Expires"]
current_time = datetime.now(expires_time.tzinfo)
if current_time > expires_time:

9
tests/test_litellm/caching/test_timedelta_serialization.py
View file

@ -12,10 +12,11 @@ from unittest.mock import Mock, patch
import pytest
from litellm.caching.redis_cache import RedisCache, TimedeltaJSONEncoder
from litellm.caching.s3_cache import S3Cache, TimedeltaJSONEncoder as S3TimedeltaJSONEncoder
from litellm.caching.gcs_cache import GCSCache, TimedeltaJSONEncoder as GCSTimedeltaJSONEncoder
from litellm.caching.azure_blob_cache import AzureBlobCache, TimedeltaJSONEncoder as AzureTimedeltaJSONEncoder
from litellm.caching.redis_cache import RedisCache
from litellm.caching.s3_cache import S3Cache
from litellm.caching.gcs_cache import GCSCache
from litellm.caching.azure_blob_cache import AzureBlobCache
from litellm.caching.json_utils import TimedeltaJSONEncoder
class TestTimedeltaJSONEncoder(unittest.TestCase):