diff --git a/litellm/llms/vertex_ai/context_caching/known_cached_contents.py b/litellm/llms/vertex_ai/context_caching/known_cached_contents.py new file mode 100644 index 00000000000..daef9076349 --- /dev/null +++ b/litellm/llms/vertex_ai/context_caching/known_cached_contents.py @@ -0,0 +1,139 @@ +import asyncio +import re +import threading +import time +import weakref +from collections.abc import Awaitable, Callable, Mapping +from datetime import datetime, timezone +from typing import Final, Literal + +# 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 +CACHE_REJECTED_STATUS_CODES: Final = frozenset({400, 403, 404}) + +_RFC3339_FRACTION: Final = re.compile(r"\.(\d{1,6})\d*") + + +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[tuple[str, str], str], tuple[str, float]] = {} # mutable-ok: guarded by _lock + self._lock: Final = threading.Lock() + self._inflight: weakref.WeakKeyDictionary[ # mutable-ok: private, guarded by _lock + asyncio.AbstractEventLoop, dict[tuple[tuple[str, str], str], asyncio.Future[str | None]] + ] = weakref.WeakKeyDictionary() + + @staticmethod + def scope( + base_url: str, + custom_llm_provider: Literal["vertex_ai", "vertex_ai_beta", "gemini"], + headers: Mapping[str, str], + ) -> tuple[str, str]: + """Where a cache name is valid: the Vertex URL names project and location, a Gemini key names the project. + + The Gemini key is read from the headers actually sent, so an ``extra_headers`` override scopes the + cache to the caller's project. It is kept as-is in this in-memory scope and never logged or persisted. + """ + if custom_llm_provider != "gemini": + return (base_url, "") + sent_keys: Final = [value for name, value in headers.items() if name.lower() == "x-goog-api-key"] + return (base_url, sent_keys[-1] if sent_keys else "") + + @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: tuple[str, 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: tuple[str, 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) + + def forget_rejected(self, cached_content: str | None, status_code: int) -> None: + if cached_content is None or status_code not in CACHE_REJECTED_STATUS_CODES: + return + with self._lock: + self._entries = {key: entry for key, entry in self._entries.items() if entry[0] != cached_content} + + async def single_flight( + self, scope: tuple[str, 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() 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 c423b0ccaf3..326872a367f 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,10 +1,4 @@ -import asyncio -import re -import threading -import time -import weakref -from collections.abc import Awaitable, Callable, Mapping, Sequence -from datetime import datetime, timezone +from collections.abc import Sequence from typing import Final, Literal import httpx @@ -26,6 +20,7 @@ from litellm.utils import is_prompt_caching_valid_prompt from ..common_utils import VertexAIError, get_vertex_base_url from ..vertex_llm_base import VertexBase +from .known_cached_contents import KnownCachedContents, known_cached_contents from .transformation import ( cached_messages_end_on_supported_turn, separate_cached_messages, @@ -36,13 +31,6 @@ 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): @@ -82,123 +70,6 @@ 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[tuple[str, str], str], tuple[str, float]] = {} # mutable-ok: guarded by _lock - self._lock: Final = threading.Lock() - self._inflight: weakref.WeakKeyDictionary[ # mutable-ok: private, guarded by _lock - asyncio.AbstractEventLoop, dict[tuple[tuple[str, str], str], asyncio.Future[str | None]] - ] = weakref.WeakKeyDictionary() - - @staticmethod - def scope( - base_url: str, - custom_llm_provider: Literal["vertex_ai", "vertex_ai_beta", "gemini"], - headers: Mapping[str, str], - ) -> tuple[str, str]: - """Where a cache name is valid: the Vertex URL names project and location, a Gemini key names the project. - - The Gemini key is read from the headers actually sent, so an ``extra_headers`` override scopes the - cache to the caller's project. It is kept as-is in this in-memory scope and never logged or persisted. - """ - if custom_llm_provider != "gemini": - return (base_url, "") - sent_keys: Final = [value for name, value in headers.items() if name.lower() == "x-goog-api-key"] - return (base_url, sent_keys[-1] if sent_keys else "") - - @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: tuple[str, 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: tuple[str, 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: tuple[str, 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 @@ -291,7 +162,7 @@ class ContextCachingEndpoints(VertexBase): vertex_auth_header=vertex_auth_header, model=model, ) - scope: Final = _KnownCachedContents.scope(base_url, custom_llm_provider, headers) + scope: Final = KnownCachedContents.scope(base_url, custom_llm_provider, headers) known_name: Final = known_cached_contents.get(scope, cache_key) if known_name is not None: return known_name @@ -340,7 +211,7 @@ class ContextCachingEndpoints(VertexBase): if display_name is None or display_name != cache_key: continue expire_time = cached_item.get("expireTime") - if not _KnownCachedContents.is_usable(expire_time): + if not KnownCachedContents.is_usable(expire_time): continue name = cached_item.get("name") known_cached_contents.remember(scope, cache_key, name, expire_time) @@ -388,7 +259,7 @@ class ContextCachingEndpoints(VertexBase): vertex_auth_header=vertex_auth_header, model=model, ) - scope: Final = _KnownCachedContents.scope(base_url, custom_llm_provider, headers) + scope: Final = KnownCachedContents.scope(base_url, custom_llm_provider, headers) known_name: Final = known_cached_contents.get(scope, cache_key) if known_name is not None: return known_name @@ -437,7 +308,7 @@ class ContextCachingEndpoints(VertexBase): if display_name is None or display_name != cache_key: continue expire_time = cached_item.get("expireTime") - if not _KnownCachedContents.is_usable(expire_time): + if not KnownCachedContents.is_usable(expire_time): continue name = cached_item.get("name") known_cached_contents.remember(scope, cache_key, name, expire_time) @@ -604,7 +475,7 @@ 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, headers), + KnownCachedContents.scope(url, custom_llm_provider, headers), generated_cache_key, raw_response_cached.get("name"), raw_response_cached.get("expireTime"), @@ -706,7 +577,7 @@ class ContextCachingEndpoints(VertexBase): generated_cache_key: Final = local_cache_obj.get_cache_key( messages=cached_messages, tools=tools, tool_choice=tool_choice, model=model ) - scope: Final = _KnownCachedContents.scope(url, custom_llm_provider, headers) + scope: Final = KnownCachedContents.scope(url, custom_llm_provider, headers) async def lookup_or_create() -> str | None: google_cache_name: Final = await self.async_check_cache( diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index 257fafbb6b8..f30e8c69324 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -96,6 +96,7 @@ from ..common_utils import ( build_vertex_schema, supports_response_json_schema, ) +from ..context_caching.known_cached_contents import known_cached_contents from ..vertex_llm_base import VertexBase from .grounding_requests import calculate_grounding_requests from .long_decode import ( @@ -2831,6 +2832,7 @@ async def make_call( ) response.raise_for_status() except httpx.HTTPStatusError as e: + known_cached_contents.forget_rejected(request_body.get("cachedContent"), e.response.status_code) exception_string: Final = str(await e.response.aread()) raise VertexAIError( status_code=e.response.status_code, @@ -2891,14 +2893,18 @@ def make_sync_call( if client is None: client = HTTPHandler() # Create a new client if none provided - response: Final = client.post( - url_for(api_base, request_body), - headers=headers, - data=data, - stream=True, - logging_obj=logging_obj, - timeout=timeout, - ) + try: + response: Final = client.post( + url_for(api_base, request_body), + headers=headers, + data=data, + stream=True, + logging_obj=logging_obj, + timeout=timeout, + ) + except httpx.HTTPStatusError as e: + known_cached_contents.forget_rejected(request_body.get("cachedContent"), e.response.status_code) + raise if response.status_code != 200 and response.status_code != 201: raise VertexAIError( @@ -3175,6 +3181,7 @@ class VertexLLM(VertexBase): ) response.raise_for_status() except httpx.HTTPStatusError as err: + known_cached_contents.forget_rejected(body.get("cachedContent"), err.response.status_code) raise VertexAIError( status_code=err.response.status_code, message=err.response.text, @@ -3385,6 +3392,7 @@ class VertexLLM(VertexBase): ) response.raise_for_status() except httpx.HTTPStatusError as err: + known_cached_contents.forget_rejected(body.get("cachedContent"), err.response.status_code) raise VertexAIError( status_code=err.response.status_code, message=err.response.text, diff --git a/tests/unit/llms/vertex_ai/context_caching/conftest.py b/tests/unit/llms/vertex_ai/context_caching/conftest.py index 75bcb5cde8a..9c80c964608 100644 --- a/tests/unit/llms/vertex_ai/context_caching/conftest.py +++ b/tests/unit/llms/vertex_ai/context_caching/conftest.py @@ -1,8 +1,6 @@ import pytest -from litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching import ( - known_cached_contents, -) +from litellm.llms.vertex_ai.context_caching.known_cached_contents import known_cached_contents @pytest.fixture(autouse=True) 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 index 33052341756..8f9c4ff05d9 100644 --- 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 @@ -1,21 +1,24 @@ import asyncio +import json from datetime import datetime, timedelta, timezone from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import httpx +import openai import pytest +import respx + +import litellm 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, -) +from litellm.llms.vertex_ai.context_caching.known_cached_contents import known_cached_contents, parse_expire_time +from litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching import ContextCachingEndpoints _MODULE: Final = "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching" +_REGISTRY_MODULE: Final = "litellm.llms.vertex_ai.context_caching.known_cached_contents" _SCOPE: Final = ("https://vertex.test/v1/projects/p/locations/us/cachedContents", "") @@ -163,7 +166,7 @@ def test_remembered_cache_is_listed_again_once_its_reuse_window_closes(): 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): + with patch(f"{_REGISTRY_MODULE}.time.time", return_value=datetime.now(timezone.utc).timestamp() + 55): assert _sync_call(client)[2] == "cachedContents/listed" @@ -290,18 +293,18 @@ async def test_waiters_resolve_on_their_own_when_the_first_request_fails(): def test_parse_expire_time_accepts_nanosecond_precision(): - assert _parse_expire_time("2026-10-14T18:06:21.123456789Z") == pytest.approx( + 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() + 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 + assert parse_expire_time("not a time") is None + assert parse_expire_time(None) is None def test_parse_expire_time_reads_a_missing_offset_as_utc(): - assert _parse_expire_time("2026-10-14T18:06:21") == datetime(2026, 10, 14, 18, 6, 21, tzinfo=timezone.utc).timestamp() + assert parse_expire_time("2026-10-14T18:06:21") == datetime(2026, 10, 14, 18, 6, 21, tzinfo=timezone.utc).timestamp() def test_cache_expiring_within_the_margin_is_not_remembered(): @@ -311,9 +314,59 @@ def test_cache_expiring_within_the_margin_is_not_remembered(): def test_registry_is_bounded(): - with patch(f"{_MODULE}.MAX_KNOWN_CACHED_CONTENTS", 3): + with patch(f"{_REGISTRY_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" + + +async def _complete(sync_mode: bool, stream: bool) -> None: + if sync_mode: + response = litellm.completion(model="gemini/gemini-3.1-pro", messages=_messages(), stream=stream) + for _chunk in response if stream else (): + pass + return + response = await litellm.acompletion(model="gemini/gemini-3.1-pro", messages=_messages(), stream=stream) + if stream: + async for _chunk in response: + pass + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.parametrize("stream", [True, False]) +@pytest.mark.parametrize( + ("status_code", "relisted"), + [(400, True), (403, True), (404, True), (429, False), (500, False)], +) +@pytest.mark.asyncio +async def test_cache_rejected_by_generate_content_is_looked_up_again( + respx_mock: respx.MockRouter, + fake_provider_credentials: None, + monkeypatch: pytest.MonkeyPatch, + sync_mode: bool, + stream: bool, + status_code: int, + relisted: bool, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + list_route: Final = respx_mock.get(url__regex=r".*/cachedContents.*").mock( + return_value=httpx.Response(200, json={}) + ) + respx_mock.post(url__regex=r".*/cachedContents$").mock( + return_value=httpx.Response(200, json=_created("cachedContents/deleted", _expire_time(300))) + ) + generate_route: Final = respx_mock.post(url__regex=r".*:(stream)?[gG]enerateContent.*").mock( + return_value=httpx.Response(status_code, json={"error": {"code": status_code, "message": "rejected"}}) + ) + + for _ in range(2): + with pytest.raises(openai.APIStatusError): + await _complete(sync_mode, stream) + + assert [json.loads(call.request.content)["cachedContent"] for call in generate_route.calls] == [ + "cachedContents/deleted", + "cachedContents/deleted", + ] + assert list_route.call_count == (2 if relisted else 1)