diff --git a/litellm/caching/caching.py b/litellm/caching/caching.py index 5f5b8210096..1a1de3eeab3 100644 --- a/litellm/caching/caching.py +++ b/litellm/caching/caching.py @@ -109,6 +109,8 @@ class Cache: default_in_redis_ttl: float | None = None, similarity_threshold: float | None = None, supported_call_types: list[CachingSupportedCallTypes] | None = list(DEFAULT_CACHING_SUPPORTED_CALL_TYPES), + excluded_call_types: list[CachingSupportedCallTypes] | None = None, + call_type_ttl: dict[str, float] | None = None, # s3 Bucket, boto3 configuration azure_account_url: str | None = None, azure_blob_container: str | None = None, @@ -314,9 +316,17 @@ class Cache: litellm.logging_callback_manager.add_litellm_success_callback("cache") if "cache" not in litellm._async_success_callback: litellm.logging_callback_manager.add_litellm_async_success_callback("cache") - self.supported_call_types = ( - supported_call_types # default to ["completion", "acompletion", "embedding", "aembedding"] - ) + if excluded_call_types and supported_call_types is not None: + excluded_set: Final = set(excluded_call_types) + self.supported_call_types: list[CachingSupportedCallTypes] | None = [ + ct for ct in supported_call_types if ct not in excluded_set + ] + else: + self.supported_call_types = ( + supported_call_types # default to ["completion", "acompletion", "embedding", "aembedding"] + ) + self.excluded_call_types: list[CachingSupportedCallTypes] | None = excluded_call_types + self.call_type_ttl: dict[str, float] | None = call_type_ttl self.type = type self.namespace = namespace self.redis_flush_size = redis_flush_size @@ -707,6 +717,18 @@ class Cache: print_verbose(f"An exception occurred: {traceback.format_exc()}") return None + def _get_ttl_for_call_type(self, kwargs: Mapping[str, object]) -> float | None: + if self.call_type_ttl: + call_type: Final = kwargs.get("call_type") or ("embedding" if kwargs.get("embedding") is True else None) + if isinstance(call_type, str): + if call_type in self.call_type_ttl: + return self.call_type_ttl[call_type] + if call_type in ("embedding", "aembedding"): + emb_ttl: Final = self.call_type_ttl.get("embedding") or self.call_type_ttl.get("aembedding") + if emb_ttl is not None: + return emb_ttl + return self.ttl + def _add_cache_logic(self, result, **kwargs): """ Common implementation across sync + async add_cache functions @@ -721,8 +743,9 @@ class Cache: result = result.model_dump_json(exclude_unset=True) ## DEFAULT TTL ## - if self.ttl is not None: - kwargs["ttl"] = self.ttl + resolved_ttl: Final = self._get_ttl_for_call_type(kwargs) + if resolved_ttl is not None: + kwargs["ttl"] = resolved_ttl ## Get Cache-Controls ## _cache_kwargs: Final = kwargs.get("cache", None) if isinstance(_cache_kwargs, dict): @@ -911,6 +934,8 @@ class Cache: prompt_tokens_details=prompt_tokens_details, ) + if "call_type" not in kwargs and "embedding" not in kwargs: + kwargs["embedding"] = True cache_key, cached_data, kwargs = self._add_cache_logic( result=embedding_dict, **kwargs, @@ -938,8 +963,9 @@ class Cache: return # set default ttl if not set - if self.ttl is not None: - kwargs["ttl"] = self.ttl + pipeline_ttl: Final = self._get_ttl_for_call_type({**kwargs, "embedding": True}) + if pipeline_ttl is not None: + kwargs["ttl"] = pipeline_ttl cache_list: Final = [] if isinstance(kwargs["input"], list): diff --git a/tests/unit/caching/test_caching.py b/tests/unit/caching/test_caching.py index 91083921130..628245ce2dd 100644 --- a/tests/unit/caching/test_caching.py +++ b/tests/unit/caching/test_caching.py @@ -1566,3 +1566,66 @@ def test_cache_key_separates_model_ids_unless_a_model_group_owns_the_key() -> No assert key_a != bare assert key_a == key_a_again assert grouped_a == grouped_b + + +def test_cache_excluded_call_types_filters_defaults() -> None: + cache: Final = Cache( + type=LiteLLMCacheType.LOCAL, + excluded_call_types=["embedding", "aembedding"], + ) + assert cache.supported_call_types is not None + assert "embedding" not in cache.supported_call_types + assert "aembedding" not in cache.supported_call_types + assert "completion" in cache.supported_call_types + assert "acompletion" in cache.supported_call_types + + +@pytest.mark.asyncio +async def test_cache_call_type_ttl_overrides_default_ttl_for_embeddings() -> None: + from litellm.types.utils import EmbeddingResponse + + cache: Final = Cache( + type=LiteLLMCacheType.LOCAL, + ttl=3600, + call_type_ttl={"embedding": 60}, + ) + response: Final = EmbeddingResponse( + model="text-embedding-3-small", + data=[{"embedding": [0.1, 0.2], "index": 0, "object": "embedding"}], + ) + _, _, updated_kwargs = cache.add_embedding_response_to_cache( + result=response, + input="hello world", + kwargs={"model": "text-embedding-3-small"}, + idx_in_result_data=0, + ) + assert updated_kwargs["ttl"] == 60 + + # Explicit per-request cache={"ttl": ...} still overrides call_type_ttl + _, _, override_kwargs = cache.add_embedding_response_to_cache( + result=response, + input="hello world", + kwargs={"model": "text-embedding-3-small", "cache": {"ttl": 15}}, + idx_in_result_data=0, + ) + assert override_kwargs["ttl"] == 15 + + # Pipeline bulk write also forwards call_type_ttl + mock_backend: Final = AsyncMock() + await cache.async_add_cache_pipeline( + result=response, + dynamic_cache_object=mock_backend, + model="text-embedding-3-small", + input=["hello world"], + ) + mock_backend.async_set_cache_pipeline.assert_awaited_once() + assert mock_backend.async_set_cache_pipeline.call_args.kwargs["ttl"] == 60 + + # Non-embedding calls still use default ttl (3600) + _, _, completion_kwargs = cache._add_cache_logic( + result="cached-completion", + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hi"}], + call_type="completion", + ) + assert completion_kwargs["ttl"] == 3600