diff --git a/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py b/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py index 06b8f3a3fbb..65116b407be 100644 --- a/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py +++ b/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py @@ -1,4 +1,11 @@ -from collections.abc import Sequence +import asyncio +import hashlib +import re +import threading +import time +import weakref +from collections.abc import Awaitable, Callable, Sequence +from datetime import datetime, timezone from typing import Final, Literal import httpx @@ -30,11 +37,19 @@ local_cache_obj: Final = Cache(type=LiteLLMCacheType.LOCAL) # only used for cal MAX_PAGINATION_PAGES: Final = 100 # Reasonable upper bound for pagination +# A remembered cache stops being reused this long before Vertex expires it, so a request +# resolved just before expiry does not reach generateContent with an expired cache. +KNOWN_CACHE_EXPIRY_MARGIN_SECONDS: Final = 10.0 +MAX_KNOWN_CACHED_CONTENTS: Final = 10_000 + +_RFC3339_FRACTION: Final = re.compile(r"\.(\d{1,6})\d*") + @with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True)) class _CachedContentItem(TypedDict): displayName: ReadOnly[NotRequired[object]] name: ReadOnly[NotRequired[str | None]] + expireTime: ReadOnly[NotRequired[str | None]] @with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True)) @@ -51,6 +66,7 @@ class _CachedContentUsageMetadata(TypedDict): @with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True)) class _CreatedCachedContent(TypedDict): name: ReadOnly[NotRequired[str | None]] + expireTime: ReadOnly[NotRequired[str | None]] usageMetadata: ReadOnly[NotRequired[_CachedContentUsageMetadata | None]] @@ -67,6 +83,118 @@ def _record_cache_creation_tokens(logging_obj: Logging, created: _CreatedCachedC logging_obj.vertex_ai_cache_creation_input_tokens = max(tokens, 0) # rebind-ok: set provider count +def _parse_expire_time(expire_time: object) -> float | None: + """Parse a Vertex RFC 3339 ``expireTime`` (nanosecond precision, ``Z`` suffix) to epoch seconds.""" + if not isinstance(expire_time, str) or not expire_time: + return None + normalized: Final = _RFC3339_FRACTION.sub(lambda m: "." + m.group(1), expire_time.replace("Z", "+00:00"), count=1) + try: + parsed: Final = datetime.fromisoformat(normalized) + except ValueError: + return None + if parsed.tzinfo is None: + return parsed.replace(tzinfo=timezone.utc).timestamp() + return parsed.timestamp() + + +class _KnownCachedContents: + """Process-local record of the cachedContents this process created or found. + + Vertex meters every cachedContents list page against a per-minute, per-location quota, + so listing on every cache-marked request exhausts it under load. A cache is only + remembered with the ``expireTime`` Vertex reported, and is never reused within + ``KNOWN_CACHE_EXPIRY_MARGIN_SECONDS`` of it. + """ + + def __init__(self) -> None: + self._entries: dict[tuple[str, str], tuple[str, float]] = {} # mutable-ok: private, guarded by _lock + self._lock: Final = threading.Lock() + self._inflight: weakref.WeakKeyDictionary[ # mutable-ok: private, guarded by _lock + asyncio.AbstractEventLoop, dict[tuple[str, str], asyncio.Future[str | None]] + ] = weakref.WeakKeyDictionary() + + @staticmethod + def scope( + base_url: str, + custom_llm_provider: Literal["vertex_ai", "vertex_ai_beta", "gemini"], + api_key: str | None, + ) -> str: + """Where a cache name is valid: the Vertex URL names project and location, a Gemini key names the project.""" + if custom_llm_provider == "gemini": + return f"{base_url}#{hashlib.sha256((api_key or '').encode()).hexdigest()}" + return base_url + + @staticmethod + def is_usable(expire_time: object) -> bool: + expires_at: Final = _parse_expire_time(expire_time) + return expires_at is None or expires_at - KNOWN_CACHE_EXPIRY_MARGIN_SECONDS > time.time() + + def get(self, scope: str, cache_key: str) -> str | None: + with self._lock: + entry: Final = self._entries.get((scope, cache_key)) + if entry is None: + return None + name, reuse_until = entry + if reuse_until > time.time(): + return name + del self._entries[(scope, cache_key)] + return None + + def remember(self, scope: str, cache_key: str, name: str | None, expire_time: object) -> None: + expires_at: Final = _parse_expire_time(expire_time) + if not name or expires_at is None: + return + reuse_until: Final = expires_at - KNOWN_CACHE_EXPIRY_MARGIN_SECONDS + now: Final = time.time() + if reuse_until <= now: + return + with self._lock: + if len(self._entries) >= MAX_KNOWN_CACHED_CONTENTS: + self._entries = {k: v for k, v in self._entries.items() if v[1] > now} + while len(self._entries) >= MAX_KNOWN_CACHED_CONTENTS: + del self._entries[next(iter(self._entries))] + self._entries[(scope, cache_key)] = (name, reuse_until) + + async def single_flight( + self, scope: str, cache_key: str, resolve: Callable[[], Awaitable[str | None]] + ) -> str | None: + """Run ``resolve`` once per (scope, cache_key) at a time within this event loop. + + Concurrent requests for the same prefix wait for the first one's lookup-or-create + instead of each listing (and creating a duplicate cache). If that first attempt + fails, each waiter resolves on its own. + """ + loop: Final = asyncio.get_running_loop() + own: Final[asyncio.Future[str | None]] = loop.create_future() + with self._lock: + inflight: Final = self._inflight.setdefault(loop, {}) + leader: Final = inflight.setdefault((scope, cache_key), own) + if leader is not own: + name: Final = await asyncio.shield(leader) + if name: + return name + return await resolve() + try: + resolved: Final = await resolve() + except BaseException: + own.set_result(None) + raise + else: + own.set_result(resolved) + return resolved + finally: + with self._lock: + inflight.pop((scope, cache_key), None) + + def clear(self) -> None: + with self._lock: + self._entries.clear() + self._inflight.clear() + + +known_cached_contents: Final = _KnownCachedContents() + + class ContextCachingEndpoints(VertexBase): """ Covers context caching endpoints for Vertex AI + Google AI Studio @@ -159,6 +287,10 @@ class ContextCachingEndpoints(VertexBase): vertex_auth_header=vertex_auth_header, model=model, ) + scope: Final = _KnownCachedContents.scope(base_url, custom_llm_provider, api_key) + known_name: Final = known_cached_contents.get(scope, cache_key) + if known_name is not None: + return known_name page_token: str | None = None @@ -201,8 +333,14 @@ class ContextCachingEndpoints(VertexBase): # Check current page for matching cache_key for cached_item in raw_response["cachedContents"]: display_name = cached_item.get("displayName") - if display_name is not None and display_name == cache_key: - return cached_item.get("name") + if display_name is None or display_name != cache_key: + continue + expire_time = cached_item.get("expireTime") + if not _KnownCachedContents.is_usable(expire_time): + continue + name = cached_item.get("name") + known_cached_contents.remember(scope, cache_key, name, expire_time) + return name # Check if there are more pages page_token = raw_response.get("nextPageToken") @@ -246,6 +384,10 @@ class ContextCachingEndpoints(VertexBase): vertex_auth_header=vertex_auth_header, model=model, ) + scope: Final = _KnownCachedContents.scope(base_url, custom_llm_provider, api_key) + known_name: Final = known_cached_contents.get(scope, cache_key) + if known_name is not None: + return known_name page_token: str | None = None @@ -288,8 +430,14 @@ class ContextCachingEndpoints(VertexBase): # Check current page for matching cache_key for cached_item in raw_response["cachedContents"]: display_name = cached_item.get("displayName") - if display_name is not None and display_name == cache_key: - return cached_item.get("name") + if display_name is None or display_name != cache_key: + continue + expire_time = cached_item.get("expireTime") + if not _KnownCachedContents.is_usable(expire_time): + continue + name = cached_item.get("name") + known_cached_contents.remember(scope, cache_key, name, expire_time) + return name # Check if there are more pages page_token = raw_response.get("nextPageToken") @@ -451,6 +599,12 @@ class ContextCachingEndpoints(VertexBase): raw_response_cached: Final = _CreatedCachedContentRoot.model_validate(response.json()).root _record_cache_creation_tokens(logging_obj, raw_response_cached) + known_cached_contents.remember( + _KnownCachedContents.scope(url, custom_llm_provider, api_key), + generated_cache_key, + raw_response_cached.get("name"), + raw_response_cached.get("expireTime"), + ) return ( non_cached_messages, optional_params, @@ -548,68 +702,73 @@ class ContextCachingEndpoints(VertexBase): generated_cache_key: Final = local_cache_obj.get_cache_key( messages=cached_messages, tools=tools, tool_choice=tool_choice, model=model ) - google_cache_name: Final = await self.async_check_cache( - cache_key=generated_cache_key, - client=client, - headers=headers, - api_key=api_key, - api_base=api_base, - logging_obj=logging_obj, - custom_llm_provider=custom_llm_provider, - vertex_project=vertex_project, - vertex_location=vertex_location, - vertex_auth_header=vertex_auth_header, - model=model, - ) + scope: Final = _KnownCachedContents.scope(url, custom_llm_provider, api_key) - if google_cache_name: - return non_cached_messages, optional_params, google_cache_name - - ## TRANSFORM REQUEST - cached_content_request_body: Final = transform_openai_messages_to_gemini_context_caching( - model=model, - messages=cached_messages, - cache_key=generated_cache_key, - custom_llm_provider=custom_llm_provider, - vertex_project=vertex_project, - vertex_location=vertex_location, - ) - - cached_content_request_body["tools"] = tools - if tool_choice is not None: - cached_content_request_body["toolConfig"] = tool_choice - - ## LOGGING - logging_obj.pre_call( - input=messages, - api_key="", - additional_args={ - "complete_input_dict": cached_content_request_body, - "api_base": url, - "headers": headers, - }, - ) - - try: - response: Final = await client.post( - url=url, + async def lookup_or_create() -> str | None: + google_cache_name: Final = await self.async_check_cache( + cache_key=generated_cache_key, + client=client, headers=headers, - json=cached_content_request_body, + api_key=api_key, + api_base=api_base, + logging_obj=logging_obj, + custom_llm_provider=custom_llm_provider, + vertex_project=vertex_project, + vertex_location=vertex_location, + vertex_auth_header=vertex_auth_header, + model=model, ) - response.raise_for_status() - except httpx.HTTPStatusError as err: - error_code: Final = err.response.status_code - raise VertexAIError(status_code=error_code, message=err.response.text) - except httpx.TimeoutException: - raise VertexAIError(status_code=408, message="Timeout error occurred.") - raw_response_cached: Final = _CreatedCachedContentRoot.model_validate(response.json()).root - _record_cache_creation_tokens(logging_obj, raw_response_cached) - return ( - non_cached_messages, - optional_params, - raw_response_cached.get("name"), - ) + if google_cache_name: + return google_cache_name + + ## TRANSFORM REQUEST + cached_content_request_body: Final = transform_openai_messages_to_gemini_context_caching( + model=model, + messages=cached_messages, + cache_key=generated_cache_key, + custom_llm_provider=custom_llm_provider, + vertex_project=vertex_project, + vertex_location=vertex_location, + ) + + cached_content_request_body["tools"] = tools + if tool_choice is not None: + cached_content_request_body["toolConfig"] = tool_choice + + ## LOGGING + logging_obj.pre_call( + input=messages, + api_key="", + additional_args={ + "complete_input_dict": cached_content_request_body, + "api_base": url, + "headers": headers, + }, + ) + + try: + response: Final = await client.post( + url=url, + headers=headers, + json=cached_content_request_body, + ) + response.raise_for_status() + except httpx.HTTPStatusError as err: + error_code: Final = err.response.status_code + raise VertexAIError(status_code=error_code, message=err.response.text) + except httpx.TimeoutException: + raise VertexAIError(status_code=408, message="Timeout error occurred.") + + raw_response_cached: Final = _CreatedCachedContentRoot.model_validate(response.json()).root + _record_cache_creation_tokens(logging_obj, raw_response_cached) + known_cached_contents.remember( + scope, generated_cache_key, raw_response_cached.get("name"), raw_response_cached.get("expireTime") + ) + return raw_response_cached.get("name") + + cache_name: Final = await known_cached_contents.single_flight(scope, generated_cache_key, lookup_or_create) + return non_cached_messages, optional_params, cache_name def get_cache(self): pass diff --git a/tests/unit/llms/vertex_ai/context_caching/conftest.py b/tests/unit/llms/vertex_ai/context_caching/conftest.py new file mode 100644 index 00000000000..75bcb5cde8a --- /dev/null +++ b/tests/unit/llms/vertex_ai/context_caching/conftest.py @@ -0,0 +1,12 @@ +import pytest + +from litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching import ( + known_cached_contents, +) + + +@pytest.fixture(autouse=True) +def _forget_known_cached_contents(): + known_cached_contents.clear() + yield + known_cached_contents.clear() diff --git a/tests/unit/llms/vertex_ai/context_caching/test_known_cached_contents.py b/tests/unit/llms/vertex_ai/context_caching/test_known_cached_contents.py new file mode 100644 index 00000000000..089ed9cb69b --- /dev/null +++ b/tests/unit/llms/vertex_ai/context_caching/test_known_cached_contents.py @@ -0,0 +1,238 @@ +import asyncio +from datetime import datetime, timedelta, timezone +from typing import Final +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler +from litellm.llms.vertex_ai.common_utils import VertexAIError +from litellm.llms.vertex_ai.context_caching import vertex_ai_context_caching +from litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching import ( + ContextCachingEndpoints, + _parse_expire_time, + known_cached_contents, +) + +_MODULE: Final = "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching" + + +@pytest.fixture(autouse=True) +def _cacheable_prompt_and_fixed_url(): + with ( + patch(f"{_MODULE}.is_prompt_caching_valid_prompt", return_value=True), + patch.object( + ContextCachingEndpoints, + "_get_token_and_url_context_caching", + return_value=("token", "https://vertex.test/v1/projects/p/locations/us/cachedContents"), + ), + ): + yield + + +def _expire_time(seconds_from_now: float) -> str: + expires: Final = datetime.now(timezone.utc) + timedelta(seconds=seconds_from_now) + return expires.strftime("%Y-%m-%dT%H:%M:%S.%f") + "123Z" + + +def _messages(prefix: str = "cacheable prefix ") -> list: + return [ + {"role": "system", "content": [{"type": "text", "text": prefix * 50, "cache_control": {"type": "ephemeral"}}]}, + {"role": "user", "content": "question"}, + ] + + +def _response(body: dict) -> MagicMock: + return MagicMock(json=lambda: body) + + +def _created(name: str, expire_time: str | None) -> dict: + body: Final = {"name": name, "usageMetadata": {"totalTokenCount": 2048}} + return body if expire_time is None else {**body, "expireTime": expire_time} + + +def _sync_call(client: HTTPHandler, api_key: str = "k", custom_llm_provider: str = "vertex_ai", messages=None): + return ContextCachingEndpoints().check_and_create_cache( + messages=messages or _messages(), + optional_params={}, + api_key=api_key, + api_base=None, + model="gemini-3.1-pro", + client=client, + timeout=30.0, + logging_obj=MagicMock(model_call_details={}), + custom_llm_provider=custom_llm_provider, + vertex_project="p", + vertex_location="us", + vertex_auth_header=None, + ) + + +async def _async_call(client: AsyncHTTPHandler): + return await ContextCachingEndpoints().async_check_and_create_cache( + messages=_messages(), + optional_params={}, + api_key="k", + api_base=None, + model="gemini-3.1-pro", + client=client, + timeout=30.0, + logging_obj=MagicMock(model_call_details={}), + custom_llm_provider="vertex_ai", + vertex_project="p", + vertex_location="us", + vertex_auth_header=None, + ) + + +def test_created_cache_is_reused_without_listing(): + client: Final = MagicMock(spec=HTTPHandler) + client.get.return_value = _response({}) + client.post.return_value = _response(_created("cachedContents/1", _expire_time(60))) + + assert _sync_call(client)[2] == "cachedContents/1" + assert _sync_call(client)[2] == "cachedContents/1" + + assert client.get.call_count == 1 + assert client.post.call_count == 1 + + +def test_listed_cache_is_reused_without_listing_again(): + client: Final = MagicMock(spec=HTTPHandler) + key: Final = vertex_ai_context_caching.local_cache_obj.get_cache_key( + messages=_messages()[:1], tools=None, tool_choice=None, model="gemini-3.1-pro" + ) + client.get.return_value = _response( + {"cachedContents": [{"name": "cachedContents/listed", "displayName": key, "expireTime": _expire_time(60)}]} + ) + + assert _sync_call(client)[2] == "cachedContents/listed" + assert _sync_call(client)[2] == "cachedContents/listed" + + assert client.get.call_count == 1 + client.post.assert_not_called() + + +def test_listed_cache_about_to_expire_is_not_used(): + client: Final = MagicMock(spec=HTTPHandler) + key: Final = vertex_ai_context_caching.local_cache_obj.get_cache_key( + messages=_messages()[:1], tools=None, tool_choice=None, model="gemini-3.1-pro" + ) + client.get.return_value = _response( + {"cachedContents": [{"name": "cachedContents/stale", "displayName": key, "expireTime": _expire_time(3)}]} + ) + client.post.return_value = _response(_created("cachedContents/fresh", _expire_time(60))) + + assert _sync_call(client)[2] == "cachedContents/fresh" + + +def test_cache_without_expire_time_is_not_remembered(): + client: Final = MagicMock(spec=HTTPHandler) + client.get.return_value = _response({}) + client.post.return_value = _response(_created("cachedContents/1", None)) + + _sync_call(client) + _sync_call(client) + + assert client.get.call_count == 2 + + +def test_remembered_cache_is_listed_again_once_its_reuse_window_closes(): + client: Final = MagicMock(spec=HTTPHandler) + key: Final = vertex_ai_context_caching.local_cache_obj.get_cache_key( + messages=_messages()[:1], tools=None, tool_choice=None, model="gemini-3.1-pro" + ) + client.get.side_effect = [ + _response({}), + _response( + {"cachedContents": [{"name": "cachedContents/listed", "displayName": key, "expireTime": _expire_time(600)}]} + ), + ] + client.post.return_value = _response(_created("cachedContents/1", _expire_time(60))) + assert _sync_call(client)[2] == "cachedContents/1" + + with patch(f"{_MODULE}.time.time", return_value=datetime.now(timezone.utc).timestamp() + 55): + assert _sync_call(client)[2] == "cachedContents/listed" + + +def test_gemini_caches_are_not_shared_across_api_keys(): + client: Final = MagicMock(spec=HTTPHandler) + client.get.return_value = _response({}) + client.post.return_value = _response(_created("cachedContents/1", _expire_time(60))) + + _sync_call(client, api_key="key-a", custom_llm_provider="gemini") + _sync_call(client, api_key="key-b", custom_llm_provider="gemini") + + assert client.post.call_count == 2 + + +def test_different_prefixes_do_not_share_a_cache(): + client: Final = MagicMock(spec=HTTPHandler) + client.get.return_value = _response({}) + client.post.side_effect = [ + _response(_created("cachedContents/a", _expire_time(60))), + _response(_created("cachedContents/b", _expire_time(60))), + ] + + assert _sync_call(client, messages=_messages("first "))[2] == "cachedContents/a" + assert _sync_call(client, messages=_messages("second "))[2] == "cachedContents/b" + + +@pytest.mark.asyncio +async def test_concurrent_requests_for_one_prefix_list_and_create_once(): + client: Final = MagicMock(spec=AsyncHTTPHandler) + + async def slow_list(**_): + await asyncio.sleep(0.01) + return _response({}) + + client.get = AsyncMock(side_effect=slow_list) + client.post = AsyncMock(return_value=_response(_created("cachedContents/1", _expire_time(60)))) + + results: Final = await asyncio.gather(*(_async_call(client) for _ in range(5))) + + assert {result[2] for result in results} == {"cachedContents/1"} + assert client.get.await_count == 1 + assert client.post.await_count == 1 + + +@pytest.mark.asyncio +async def test_waiters_resolve_on_their_own_when_the_first_request_fails(): + client: Final = MagicMock(spec=AsyncHTTPHandler) + + async def slow_list(**_): + await asyncio.sleep(0.01) + return _response({}) + + client.get = AsyncMock(side_effect=slow_list) + client.post = AsyncMock( + side_effect=[ + VertexAIError(status_code=500, message="create failed"), + _response(_created("cachedContents/2", _expire_time(60))), + ] + ) + + results: Final = await asyncio.gather(_async_call(client), _async_call(client), return_exceptions=True) + + assert isinstance(results[0], VertexAIError) + assert results[1][2] == "cachedContents/2" + + +def test_parse_expire_time_accepts_nanosecond_precision(): + assert _parse_expire_time("2026-10-14T18:06:21.123456789Z") == pytest.approx( + datetime(2026, 10, 14, 18, 6, 21, 123456, tzinfo=timezone.utc).timestamp() + ) + assert ( + _parse_expire_time("2026-10-14T18:06:21Z") == datetime(2026, 10, 14, 18, 6, 21, tzinfo=timezone.utc).timestamp() + ) + assert _parse_expire_time("not a time") is None + assert _parse_expire_time(None) is None + + +def test_registry_is_bounded(): + with patch(f"{_MODULE}.MAX_KNOWN_CACHED_CONTENTS", 3): + for index in range(5): + known_cached_contents.remember("scope", f"key-{index}", f"cachedContents/{index}", _expire_time(60)) + + assert known_cached_contents.get("scope", "key-0") is None + assert known_cached_contents.get("scope", "key-4") == "cachedContents/4"