diff --git a/litellm/router_utils/prompt_caching_cache.py b/litellm/router_utils/prompt_caching_cache.py index c1ecec048e4..4ec0bc62e7f 100644 --- a/litellm/router_utils/prompt_caching_cache.py +++ b/litellm/router_utils/prompt_caching_cache.py @@ -179,10 +179,16 @@ def _lookback_keys(positions: tuple[PrefixPosition, ...]) -> tuple[str, ...]: return tuple(entry.cache_key for entry in reversed(positions) if entry.position > oldest_probed_position) +def _carrier_cache_control(carrier: Mapping[str, JsonValue]) -> JsonValue | None: + function: Final = carrier.get("function") + nested: Final = function.get("cache_control") if isinstance(function, dict) else None + return carrier.get("cache_control", nested) + + def _breakpoint_pin_ttl_seconds(carrier: JsonValue | Mapping[str, JsonValue]) -> int | None: if not isinstance(carrier, Mapping): return None - cache_control: Final = carrier.get("cache_control") + cache_control: Final = _carrier_cache_control(carrier) if not isinstance(cache_control, dict) or cache_control.get("type") != "ephemeral": return None return PROMPT_CACHE_ONE_HOUR_PIN_TTL_SECONDS if cache_control.get("ttl") == "1h" else PROMPT_CACHE_PIN_TTL_SECONDS diff --git a/tests/unit/router_utils/test_prompt_caching_cache.py b/tests/unit/router_utils/test_prompt_caching_cache.py index 7c8ad04c0fb..46af2454dea 100644 --- a/tests/unit/router_utils/test_prompt_caching_cache.py +++ b/tests/unit/router_utils/test_prompt_caching_cache.py @@ -65,6 +65,11 @@ def _tool(cache_control: dict[str, str] | None) -> ChatCompletionToolParam: return {**tool, "cache_control": cache_control} if cache_control else tool +def _function_marked_tool(cache_control: dict[str, str]) -> ChatCompletionToolParam: + tool: Final = _tool(None) + return {**tool, "function": {**tool["function"], "cache_control": cache_control}} + + @pytest.mark.parametrize( ("messages", "tools", "expected_ttl"), [ @@ -108,6 +113,12 @@ def _tool(cache_control: dict[str, str] | None) -> ChatCompletionToolParam: PROMPT_CACHE_ONE_HOUR_PIN_TTL_SECONDS, id="1h-on-tools-only", ), + pytest.param( + _system_block_messages({"type": "ephemeral"}), + [_function_marked_tool({"type": "ephemeral", "ttl": "1h"})], + PROMPT_CACHE_ONE_HOUR_PIN_TTL_SECONDS, + id="1h-nested-under-function", + ), pytest.param( _system_block_messages({"type": "ephemeral"}), [_tool(None)],