diff --git a/litellm/caching/caching_handler.py b/litellm/caching/caching_handler.py index 063faa70f2f..c5533136910 100644 --- a/litellm/caching/caching_handler.py +++ b/litellm/caching/caching_handler.py @@ -17,6 +17,7 @@ In each method it will call the appropriate method from caching.py import asyncio import datetime import inspect +import json import time from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Generator, Mapping from typing import TYPE_CHECKING, Any, Final, Optional, TypeVar @@ -214,6 +215,21 @@ async def _complete_cache_write_despite_cancellation(write_factory: Callable[[], raise +def is_json_serializable(value: object) -> bool: + """ + Whether `value` can be round-tripped through json.dumps. + + Used to guard cache writes against provider-specific response wrappers + (e.g. synthetic streaming iterators) that aren't a recognized response + type but also aren't safe to hand to the cache backend as-is. + """ + try: + json.dumps(value) + except TypeError: + return False + return True + + def create_cache_write_task(write_factory: Callable[[], Awaitable[None]]) -> "asyncio.Task[None]": with post_response_phase(): task: Final = asyncio.create_task(_complete_cache_write_despite_cancellation(write_factory)) @@ -1196,7 +1212,7 @@ class LLMCachingHandler: **new_kwargs, ) ) - else: + elif is_json_serializable(result): create_cache_write_task(lambda: cache.async_add_cache(result, **new_kwargs)) def sync_set_cache( diff --git a/tests/unit/caching/test_caching_handler.py b/tests/unit/caching/test_caching_handler.py index 7d6b1e15fc6..8f6398b25e4 100644 --- a/tests/unit/caching/test_caching_handler.py +++ b/tests/unit/caching/test_caching_handler.py @@ -24,6 +24,7 @@ from litellm.caching.caching_handler import ( CachingHandlerResponse, _is_chat_completion_cached_dict, _should_defer_streaming_cache_hit_callbacks, + is_json_serializable, ) from litellm.caching import DualCache, InMemoryCache from litellm.caching.caching import LiteLLMCacheType @@ -379,6 +380,53 @@ async def test_async_set_get_cache(response): assert cached_response.cached_result.id == result.id +def test_is_json_serializable(): + """Plain values round-trip; arbitrary objects don't.""" + + class _NotSerializable: + pass + + assert is_json_serializable({"a": 1}) is True + assert is_json_serializable("hello") is True + assert is_json_serializable(None) is True + assert is_json_serializable(_NotSerializable()) is False + + +@pytest.mark.asyncio +async def test_async_set_cache_skips_non_serializable_result(): + """ + Regression test: a response object that is neither one of the + recognized response types (ModelResponse, EmbeddingResponse, etc.) nor + JSON-serializable (e.g. a provider-specific synthetic streaming + iterator) must be skipped instead of being handed to the cache backend, + which would raise `TypeError: Object of type ... is not JSON + serializable`. + """ + setup_cache() + caching_handler = LLMCachingHandler( + original_function=completion, request_kwargs={}, start_time=datetime.now() + ) + + class _UnserializableStreamWrapper: + """Stand-in for a provider-specific object that can't be cached.""" + + pass + + result = _UnserializableStreamWrapper() + messages = [{"role": "user", "content": f"Unique message {datetime.now()}"}] + + with patch( + "litellm.caching.caching_handler.create_cache_write_task" + ) as mock_create_cache_write_task: + await caching_handler.async_set_cache( + result=result, + original_function=litellm.acompletion, + kwargs={"messages": messages}, + ) + + mock_create_cache_write_task.assert_not_called() + + @pytest.mark.asyncio async def test_async_log_cache_hit_on_callbacks(): """