From 06e3e32c88326930645e3aabc2e454d7dfc4bfdd Mon Sep 17 00:00:00 2001 From: Ondra Zahradnik Date: Tue, 20 Oct 2026 13:37:46 +0200 Subject: [PATCH] fix(vertex_ai): reuse known context caches instead of listing on every request (#46857) Vertex AI and Gemini context caching now remember each cache they create or find until 10s before its expireTime, so a repeat cache-marked request goes straight to generateContent instead of listing every cachedContents first. Concurrent async requests for one prefix share one lookup or create When generation rejects a remembered cache with a 400, 403 or 404 that names the cached content, the request rebuilds the cache and sends once more. Any other 400, 403 or 404 drops the record without a retry Fixes #36395 --- .../context_caching/known_cached_contents.py | 156 +++++ .../context_caching/rejected_cache_retry.py | 74 ++ .../vertex_ai_context_caching.py | 179 +++-- .../vertex_and_google_ai_studio_gemini.py | 251 ++++--- .../vertex_ai/context_caching/conftest.py | 10 + .../test_known_cached_contents.py | 663 ++++++++++++++++++ 6 files changed, 1173 insertions(+), 160 deletions(-) create mode 100644 litellm/llms/vertex_ai/context_caching/known_cached_contents.py create mode 100644 litellm/llms/vertex_ai/context_caching/rejected_cache_retry.py create mode 100644 tests/unit/llms/vertex_ai/context_caching/conftest.py create mode 100644 tests/unit/llms/vertex_ai/context_caching/test_known_cached_contents.py 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..6b7f098d304 --- /dev/null +++ b/litellm/llms/vertex_ai/context_caching/known_cached_contents.py @@ -0,0 +1,156 @@ +import asyncio +import contextlib +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 + +KNOWN_CACHE_EXPIRY_MARGIN_SECONDS: Final = 10.0 +MAX_KNOWN_CACHED_CONTENTS: Final = 10_000 +_CACHE_REJECTED_STATUS_CODES: Final = frozenset({400, 403, 404}) +_NAMES_CACHED_CONTENT: Final = re.compile(r"cached?\s*content", re.IGNORECASE) +_FORBIDS_FIELDS_BESIDE_A_CACHE: Final = re.compile(r"system[_\s]instruction", re.IGNORECASE) + +_RFC3339_FRACTION: Final = re.compile(r"\.(\d{1,6})\d*") + + +def names_rejected_cache(status_code: int, message: str) -> bool: + return ( + status_code in _CACHE_REJECTED_STATUS_CODES + and _NAMES_CACHED_CONTENT.search(message) is not None + and _FORBIDS_FIELDS_BESIDE_A_CACHE.search(message) is None + ) + + +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 remembered + until ``KNOWN_CACHE_EXPIRY_MARGIN_SECONDS`` before the ``expireTime`` Google reported, + and after that requests list again, as they did before the record existed. + """ + + def __init__(self, clock: Callable[[], float] = time.time) -> None: + self._clock: Final = clock + 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], + vertex_project: str | None, + vertex_location: str | None, + ) -> tuple[str, str]: + """Where a cache name is valid: a Vertex project and location, or the project a Gemini key belongs to. + + A custom ``api_base`` can replace the whole Vertex path, so the project and location are kept beside the URL. + 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, f"{vertex_project}/{vertex_location}") + 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 "") + + def has_not_expired(self, expire_time: object) -> bool: + expires_at: Final = parse_expire_time(expire_time) + return expires_at is None or expires_at > self._clock() + + 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 > self._clock(): + 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 = self._clock() + 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]], + wait_seconds: float | 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). A waiter resolves on its + own when that first attempt fails or has not finished within ``wait_seconds``, and + never cancels it. + """ + 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: + with contextlib.suppress(asyncio.TimeoutError): + if name := await asyncio.wait_for(asyncio.shield(leader), wait_seconds): + 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/rejected_cache_retry.py b/litellm/llms/vertex_ai/context_caching/rejected_cache_retry.py new file mode 100644 index 00000000000..43e4ee8dc84 --- /dev/null +++ b/litellm/llms/vertex_ai/context_caching/rejected_cache_retry.py @@ -0,0 +1,74 @@ +from collections.abc import Awaitable, Callable +from copy import deepcopy +from functools import partial +from typing import Final, TypeVar + +import httpx + +from litellm.litellm_core_utils.core_helpers import independent_snapshot +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.types.llms.openai import AllMessageValues +from litellm.types.llms.vertex_ai import RequestBody +from litellm.utils import is_cached_message + +from .known_cached_contents import names_rejected_cache + +Built: Final = TypeVar("Built") +Sent: Final = TypeVar("Sent") + + +def cache_rebuild( + build: Callable[..., Built], messages: list[AllMessageValues], optional_params: dict +) -> Callable[[], Built] | None: + """Copies the inputs before the first transform, which pops keys from ``optional_params`` and ``messages``.""" + if not any(map(is_cached_message, messages)): + return None + return partial(build, messages=deepcopy(messages), optional_params=independent_snapshot(optional_params)) + + +def _status_and_message(error: httpx.HTTPStatusError | BaseLLMException) -> tuple[int, str]: + match error: + case httpx.HTTPStatusError(): + return error.response.status_code, error.response.text + case BaseLLMException(): + return error.status_code, error.message + + +def _cache_rejected(body: RequestBody, error: httpx.HTTPStatusError | BaseLLMException) -> bool: + return body.get("cachedContent") is not None and names_rejected_cache(*_status_and_message(error)) + + +async def post_with_fresh_cache_async( + body: RequestBody, + rebuild: Callable[[], Awaitable[RequestBody]] | None, + post: Callable[[RequestBody], Awaitable[Sent]], + log_request: Callable[[RequestBody], None], +) -> tuple[RequestBody, Sent]: + try: + return body, await post(body) + except (httpx.HTTPStatusError, BaseLLMException) as error: + if rebuild is None or not _cache_rejected(body, error): + raise + fresh_body: Final = await rebuild() + log_request(fresh_body) + if fresh_body.get("cachedContent") == body.get("cachedContent"): + raise + return fresh_body, await post(fresh_body) + + +def post_with_fresh_cache_sync( + body: RequestBody, + rebuild: Callable[[], RequestBody] | None, + post: Callable[[RequestBody], Sent], + log_request: Callable[[RequestBody], None], +) -> tuple[RequestBody, Sent]: + try: + return body, post(body) + except (httpx.HTTPStatusError, BaseLLMException) as error: + if rebuild is None or not _cache_rejected(body, error): + raise + fresh_body: Final = rebuild() + log_request(fresh_body) + if fresh_body.get("cachedContent") == body.get("cachedContent"): + raise + return fresh_body, post(fresh_body) 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..5bff43828d6 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 @@ -20,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, @@ -35,6 +36,7 @@ MAX_PAGINATION_PAGES: Final = 100 # Reasonable upper bound for pagination 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 +53,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 +70,14 @@ 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 _wait_seconds(timeout: float | httpx.Timeout | None) -> float | None: + match timeout: + case httpx.Timeout(): + return timeout.read + case _: + return timeout + + class ContextCachingEndpoints(VertexBase): """ Covers context caching endpoints for Vertex AI + Google AI Studio @@ -74,8 +85,9 @@ class ContextCachingEndpoints(VertexBase): v0: covers Google AI Studio """ - def __init__(self) -> None: + def __init__(self, known_caches: KnownCachedContents = known_cached_contents) -> None: super().__init__() + self._known_caches: Final = known_caches def _get_token_and_url_context_caching( self, @@ -129,7 +141,7 @@ class ContextCachingEndpoints(VertexBase): self, cache_key: str, client: HTTPHandler, - headers: dict, + headers: dict[str, str], # mutable-ok: passed through to the HTTP client, which takes a dict api_key: str, api_base: str | None, logging_obj: Logging, @@ -159,6 +171,12 @@ class ContextCachingEndpoints(VertexBase): vertex_auth_header=vertex_auth_header, model=model, ) + scope: Final = KnownCachedContents.scope( + base_url, custom_llm_provider, headers, vertex_project=vertex_project, vertex_location=vertex_location + ) + known_name: Final = self._known_caches.get(scope, cache_key) + if known_name is not None: + return known_name page_token: str | None = None @@ -201,8 +219,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 self._known_caches.has_not_expired(expire_time): + continue + name = cached_item.get("name") + self._known_caches.remember(scope, cache_key, name, expire_time) + return name # Check if there are more pages page_token = raw_response.get("nextPageToken") @@ -216,7 +240,7 @@ class ContextCachingEndpoints(VertexBase): self, cache_key: str, client: AsyncHTTPHandler, - headers: dict, + headers: dict[str, str], # mutable-ok: passed through to the HTTP client, which takes a dict api_key: str, api_base: str | None, logging_obj: Logging, @@ -246,6 +270,12 @@ class ContextCachingEndpoints(VertexBase): vertex_auth_header=vertex_auth_header, model=model, ) + scope: Final = KnownCachedContents.scope( + base_url, custom_llm_provider, headers, vertex_project=vertex_project, vertex_location=vertex_location + ) + known_name: Final = self._known_caches.get(scope, cache_key) + if known_name is not None: + return known_name page_token: str | None = None @@ -288,8 +318,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 self._known_caches.has_not_expired(expire_time): + continue + name = cached_item.get("name") + self._known_caches.remember(scope, cache_key, name, expire_time) + return name # Check if there are more pages page_token = raw_response.get("nextPageToken") @@ -451,6 +487,14 @@ class ContextCachingEndpoints(VertexBase): raw_response_cached: Final = _CreatedCachedContentRoot.model_validate(response.json()).root _record_cache_creation_tokens(logging_obj, raw_response_cached) + self._known_caches.remember( + KnownCachedContents.scope( + url, custom_llm_provider, headers, vertex_project=vertex_project, vertex_location=vertex_location + ), + generated_cache_key, + raw_response_cached.get("name"), + raw_response_cached.get("expireTime"), + ) return ( non_cached_messages, optional_params, @@ -548,68 +592,77 @@ 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, headers, vertex_project=vertex_project, vertex_location=vertex_location ) - 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) + self._known_caches.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 self._known_caches.single_flight( + scope, generated_cache_key, lookup_or_create, _wait_seconds(timeout) ) + return non_cached_messages, optional_params, cache_name def get_cache(self): pass 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..114b9cd80ae 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,12 @@ from ..common_utils import ( build_vertex_schema, supports_response_json_schema, ) +from ..context_caching.known_cached_contents import known_cached_contents +from ..context_caching.rejected_cache_retry import ( + cache_rebuild, + post_with_fresh_cache_async, + post_with_fresh_cache_sync, +) from ..vertex_llm_base import VertexBase from .grounding_requests import calculate_grounding_requests from .long_decode import ( @@ -2786,7 +2792,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): litellm_params: dict, api_key: str | dict | None = None, api_base: str | None = None, - ) -> dict: + ) -> dict[str, str]: default_headers: Final = { "Content-Type": "application/json", } @@ -2800,6 +2806,20 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): return default_headers +def _log_generate_request( + logging_obj: LoggingClass, + messages: Sequence[AllMessageValues], + api_base: str, + headers: Mapping[str, str], + body: RequestBody, +) -> None: + logging_obj.pre_call( + input=messages, + api_key="", + additional_args={"complete_input_dict": body, "api_base": api_base, "headers": headers}, + ) + + async def make_call( client: AsyncHTTPHandler | None, # module-level client gemini_client: AsyncHTTPHandler | None, # if passed by user @@ -2832,6 +2852,7 @@ async def make_call( response.raise_for_status() except httpx.HTTPStatusError as e: exception_string: Final = str(await e.response.aread()) + known_cached_contents.forget_rejected(request_body.get("cachedContent"), e.response.status_code) raise VertexAIError( status_code=e.response.status_code, message=VertexGeminiConfig().translate_exception_str(exception_string), @@ -2845,14 +2866,18 @@ async def make_call( ) async def open_hop(body: RequestBody) -> httpx.Response: - return await client.post( - url_for(api_base, body), - headers=dict(await hop_headers()), - data=json.dumps(body), - stream=True, - logging_obj=logging_obj, - timeout=timeout, - ) + try: + return await client.post( + url_for(api_base, body), + headers=dict(await hop_headers()), + data=json.dumps(body), + stream=True, + logging_obj=logging_obj, + timeout=timeout, + ) + except httpx.HTTPStatusError as e: + known_cached_contents.forget_rejected(body.get("cachedContent"), e.response.status_code) + raise completion_stream: Final = ModelResponseIterator( streaming_response=response.aiter_lines(), @@ -2891,14 +2916,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( @@ -2908,14 +2937,18 @@ def make_sync_call( ) def open_hop(body: RequestBody) -> httpx.Response: - return client.post( - url_for(api_base, body), - headers=dict(hop_headers()), - data=json.dumps(body), - stream=True, - logging_obj=logging_obj, - timeout=timeout, - ) + try: + return client.post( + url_for(api_base, body), + headers=dict(hop_headers()), + data=json.dumps(body), + stream=True, + logging_obj=logging_obj, + timeout=timeout, + ) + except httpx.HTTPStatusError as e: + known_cached_contents.forget_rejected(body.get("cachedContent"), e.response.status_code) + raise completion_stream: Final = ModelResponseIterator( streaming_response=response.iter_lines(), @@ -2983,14 +3016,14 @@ class VertexLLM(VertexBase): custom_llm_provider: Literal[ "vertex_ai", "vertex_ai_beta", "gemini" ], # if it's vertex_ai or gemini (google ai studio) - messages: list, + messages: list[AllMessageValues], model_response: ModelResponse, print_verbose: Callable, data: dict, timeout: float | httpx.Timeout | None, encoding: object, - logging_obj, - stream, + logging_obj: LoggingClass, + stream: bool | None, optional_params: dict, litellm_params: dict, logger_fn=None, @@ -3033,6 +3066,17 @@ class VertexLLM(VertexBase): litellm_params=litellm_params, ) + rebuild: Final = cache_rebuild( + partial( + async_transform_request_body, + **data, + vertex_project=vertex_project, + vertex_location=vertex_location, + vertex_auth_header=auth_header, + ), + messages, + optional_params, + ) request_body: Final = await async_transform_request_body( **data, vertex_project=vertex_project, @@ -3041,39 +3085,37 @@ class VertexLLM(VertexBase): ) ## LOGGING - logging_obj.pre_call( - input=messages, - api_key="", - additional_args={ - "complete_input_dict": request_body, - "api_base": api_base, - "headers": headers, - }, + log_request: Final = partial(_log_generate_request, logging_obj, messages, api_base, headers) + log_request(request_body) + + open_stream: Final = partial( + make_call, + gemini_client=(client if client is not None and isinstance(client, AsyncHTTPHandler) else None), + api_base=api_base, + headers=headers, + model=model, + messages=messages, + logging_obj=logging_obj, + hop_headers=partial( + self._fresh_headers_async, + headers, + extra_headers, + vertex_credentials, + vertex_project, + custom_llm_provider, + ), + timeout=timeout, ) - request_body_str: Final = json.dumps(request_body) + async def open_with_fresh_cache(client: AsyncHTTPHandler | None) -> ModelResponseIterator: + async def post(body: RequestBody) -> ModelResponseIterator: + return await open_stream(client=client, data=json.dumps(body), request_body=body) + + return (await post_with_fresh_cache_async(request_body, rebuild, post, log_request))[1] + streaming_response: Final = CustomStreamWrapper( completion_stream=None, - make_call=partial( - make_call, - gemini_client=(client if client is not None and isinstance(client, AsyncHTTPHandler) else None), - api_base=api_base, - headers=headers, - data=request_body_str, - model=model, - messages=messages, - logging_obj=logging_obj, - request_body=request_body, - hop_headers=partial( - self._fresh_headers_async, - headers, - extra_headers, - vertex_credentials, - vertex_project, - custom_llm_provider, - ), - timeout=timeout, - ), + make_call=open_with_fresh_cache, model=model, custom_llm_provider="vertex_ai_beta", logging_obj=logging_obj, @@ -3083,7 +3125,7 @@ class VertexLLM(VertexBase): async def async_completion( self, model: str, - messages: list, + messages: list[AllMessageValues], model_response: ModelResponse, print_verbose: Callable, data: dict, @@ -3093,7 +3135,7 @@ class VertexLLM(VertexBase): timeout: float | httpx.Timeout | None, encoding: object, logging_obj: LoggingClass, - stream, + stream: bool | None, optional_params: dict, litellm_params: dict, logger_fn=None, @@ -3136,6 +3178,17 @@ class VertexLLM(VertexBase): litellm_params=litellm_params, ) + rebuild: Final = cache_rebuild( + partial( + async_transform_request_body, + **data, + vertex_project=vertex_project, + vertex_location=vertex_location, + vertex_auth_header=auth_header, + ), + messages, + optional_params, + ) request_body: Final = await async_transform_request_body( **data, vertex_project=vertex_project, @@ -3151,15 +3204,8 @@ class VertexLLM(VertexBase): else: client = client ## LOGGING - logging_obj.pre_call( - input=messages, - api_key="", - additional_args={ - "complete_input_dict": request_body, - "api_base": api_base, - "headers": headers, - }, - ) + log_request: Final = partial(_log_generate_request, logging_obj, messages, api_base, headers) + log_request(request_body) fresh_headers: Final = partial( self._fresh_headers_async, headers, extra_headers, vertex_credentials, vertex_project, custom_llm_provider @@ -3175,6 +3221,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, @@ -3188,14 +3235,14 @@ class VertexLLM(VertexBase): ) return response - first_response: Final = await post(request_body) + sent_body, first_response = await post_with_fresh_cache_async(request_body, rebuild, post, log_request) return _transform_long_decode_hops( model=model, - hops=await collect_hops_async(parse_body(first_response), request_body, post), + hops=await collect_hops_async(parse_body(first_response), sent_body, post), first_response=first_response, model_response=model_response, logging_obj=logging_obj, - request_body=request_body, + request_body=sent_body, messages=messages, ) @@ -3320,6 +3367,17 @@ class VertexLLM(VertexBase): ) ## TRANSFORMATION ## + rebuild: Final = cache_rebuild( + partial( + sync_transform_request_body, + **transform_request_params, + vertex_project=vertex_project, + vertex_location=vertex_location, + vertex_auth_header=auth_header, + ), + messages, + optional_params, + ) data: Final = sync_transform_request_body( **transform_request_params, vertex_project=vertex_project, @@ -3328,15 +3386,8 @@ class VertexLLM(VertexBase): ) ## LOGGING - logging_obj.pre_call( - input=messages, - api_key="", - additional_args={ - "complete_input_dict": data, - "api_base": url, - "headers": headers, - }, - ) + log_request: Final = partial(_log_generate_request, logging_obj, messages, url, headers) + log_request(data) fresh_headers: Final = partial( self._fresh_headers, headers, extra_headers, vertex_credentials, vertex_project, custom_llm_provider @@ -3344,22 +3395,27 @@ class VertexLLM(VertexBase): ## SYNC STREAMING CALL ## if stream is True: - request_data_str: Final = json.dumps(data) + open_stream: Final = partial( + make_sync_call, + gemini_client=(client if client is not None and isinstance(client, HTTPHandler) else None), + api_base=url, + model=model, + messages=messages, + logging_obj=logging_obj, + headers=headers, + hop_headers=fresh_headers, + timeout=timeout, + ) + + def open_with_fresh_cache(client: HTTPHandler | None) -> ModelResponseIterator: + def post(body: RequestBody) -> ModelResponseIterator: + return open_stream(client=client, data=json.dumps(body), request_body=body) + + return post_with_fresh_cache_sync(data, rebuild, post, log_request)[1] + streaming_response: Final = CustomStreamWrapper( completion_stream=None, - make_call=partial( - make_sync_call, - gemini_client=(client if client is not None and isinstance(client, HTTPHandler) else None), - api_base=url, - data=request_data_str, - model=model, - messages=messages, - logging_obj=logging_obj, - headers=headers, - request_body=data, - hop_headers=fresh_headers, - timeout=timeout, - ), + make_call=open_with_fresh_cache, model=model, custom_llm_provider="vertex_ai_beta", logging_obj=logging_obj, @@ -3385,6 +3441,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, @@ -3398,14 +3455,14 @@ class VertexLLM(VertexBase): ) return response - first_response: Final = post(data) + sent_body, first_response = post_with_fresh_cache_sync(data, rebuild, post, log_request) return _transform_long_decode_hops( model=model, - hops=collect_hops_sync(parse_body(first_response), data, post), + hops=collect_hops_sync(parse_body(first_response), sent_body, post), first_response=first_response, model_response=model_response, logging_obj=logging_obj, - request_body=data, + request_body=sent_body, messages=messages, ) 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..9c80c964608 --- /dev/null +++ b/tests/unit/llms/vertex_ai/context_caching/conftest.py @@ -0,0 +1,10 @@ +import pytest + +from litellm.llms.vertex_ai.context_caching.known_cached_contents 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..0c5c3facb56 --- /dev/null +++ b/tests/unit/llms/vertex_ai/context_caching/test_known_cached_contents.py @@ -0,0 +1,663 @@ +import asyncio +import json +import threading +from datetime import datetime, 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.known_cached_contents import KnownCachedContents, parse_expire_time +from litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching import ContextCachingEndpoints +from litellm.utils import get_prompt_cache_min_tokens + +_REGISTRY_MODULE: Final = "litellm.llms.vertex_ai.context_caching.known_cached_contents" +_SCOPE: Final = ("https://vertex.test/v1/projects/p/locations/us/cachedContents", "") +_NOW: Final = datetime(2026, 10, 14, 18, 0, tzinfo=timezone.utc).timestamp() +_FAR_FUTURE: Final = "2100-01-01T00:00:00.123456789Z" +_GATEWAY: Final = "https://gateway.test/vertex" +_PREFIX_REPEATS: Final = max( + get_prompt_cache_min_tokens(model=f"{provider}/gemini-3.1-pro") for provider in ("vertex_ai", "gemini") +) + + +@pytest.fixture +def clock() -> MagicMock: + return MagicMock(return_value=_NOW) + + +@pytest.fixture +def endpoints(clock: MagicMock) -> ContextCachingEndpoints: + return ContextCachingEndpoints(known_caches=KnownCachedContents(clock=clock)) + + +def _expire_time(seconds_from_now: float) -> str: + expires: Final = datetime.fromtimestamp(_NOW + seconds_from_now, timezone.utc) + return expires.strftime("%Y-%m-%dT%H:%M:%S.%f") + "123Z" + + +def _messages(label: str = "shared") -> list: + text: Final = f"{label} cacheable prefix " * _PREFIX_REPEATS + return [ + {"role": "system", "content": [{"type": "text", "text": text, "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( + endpoints: ContextCachingEndpoints, + client: HTTPHandler, + api_key: str = "k", + custom_llm_provider: str = "vertex_ai", + messages=None, + extra_headers: dict | None = None, + api_base: str | None = None, + vertex_project: str = "p", +): + return endpoints.check_and_create_cache( + extra_headers=extra_headers, + messages=messages or _messages(), + optional_params={}, + api_key=api_key, + api_base=api_base, + model="gemini-3.1-pro", + client=client, + timeout=30.0, + logging_obj=MagicMock(model_call_details={}), + custom_llm_provider=custom_llm_provider, + vertex_project=vertex_project, + vertex_location="us", + vertex_auth_header="token", + ) + + +async def _async_call( + endpoints: ContextCachingEndpoints, + client: AsyncHTTPHandler, + timeout: float | httpx.Timeout = 30.0, + api_base: str | None = None, + vertex_project: str = "p", +): + return await endpoints.async_check_and_create_cache( + messages=_messages(), + optional_params={}, + api_key="k", + api_base=api_base, + model="gemini-3.1-pro", + client=client, + timeout=timeout, + logging_obj=MagicMock(model_call_details={}), + custom_llm_provider="vertex_ai", + vertex_project=vertex_project, + vertex_location="us", + vertex_auth_header="token", + ) + + +def test_created_cache_is_reused_without_listing(endpoints: ContextCachingEndpoints): + client: Final = MagicMock(spec=HTTPHandler) + client.get.return_value = _response({}) + client.post.return_value = _response(_created("cachedContents/1", _expire_time(60))) + + assert _sync_call(endpoints, client)[2] == "cachedContents/1" + assert _sync_call(endpoints, 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(endpoints: ContextCachingEndpoints): + 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(endpoints, client)[2] == "cachedContents/listed" + assert _sync_call(endpoints, client)[2] == "cachedContents/listed" + + assert client.get.call_count == 1 + client.post.assert_not_called() + + +def test_listed_cache_past_its_expire_time_is_not_used(endpoints: ContextCachingEndpoints): + 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(-1)}]} + ) + client.post.return_value = _response(_created("cachedContents/fresh", _expire_time(60))) + + assert _sync_call(endpoints, client)[2] == "cachedContents/fresh" + + +def test_short_lived_cache_is_reused_from_the_listing(endpoints: ContextCachingEndpoints, clock: MagicMock): + 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/1", "displayName": key, "expireTime": _expire_time(8)}]} + ), + ] + client.post.return_value = _response(_created("cachedContents/1", _expire_time(8))) + + assert _sync_call(endpoints, client)[2] == "cachedContents/1" + clock.return_value = _NOW + 1 + assert _sync_call(endpoints, client)[2] == "cachedContents/1" + + assert client.post.call_count == 1 + + +def test_cache_without_expire_time_is_not_remembered(endpoints: ContextCachingEndpoints): + client: Final = MagicMock(spec=HTTPHandler) + client.get.return_value = _response({}) + client.post.return_value = _response(_created("cachedContents/1", None)) + + _sync_call(endpoints, client) + _sync_call(endpoints, client) + + assert client.get.call_count == 2 + + +def test_remembered_cache_is_listed_again_once_its_reuse_window_closes( + endpoints: ContextCachingEndpoints, clock: MagicMock +): + 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(endpoints, client)[2] == "cachedContents/1" + + clock.return_value = _NOW + 55 + assert _sync_call(endpoints, client)[2] == "cachedContents/listed" + + +def test_gemini_caches_are_not_shared_across_api_keys(endpoints: ContextCachingEndpoints): + client: Final = MagicMock(spec=HTTPHandler) + client.get.return_value = _response({}) + client.post.return_value = _response(_created("cachedContents/1", _expire_time(60))) + + _sync_call(endpoints, client, api_key="key-a", custom_llm_provider="gemini") + _sync_call(endpoints, client, api_key="key-b", custom_llm_provider="gemini") + + assert client.post.call_count == 2 + + +def test_gemini_cache_created_with_an_overridden_api_key_is_not_reused_with_the_configured_key( + endpoints: ContextCachingEndpoints, +): + client: Final = MagicMock(spec=HTTPHandler) + client.get.return_value = _response({}) + client.post.side_effect = [ + _response(_created("cachedContents/caller-project", _expire_time(60))), + _response(_created("cachedContents/deployment-project", _expire_time(60))), + ] + + caller: Final = _sync_call( + endpoints, + client, + api_key="deployment-key", + custom_llm_provider="gemini", + extra_headers={"X-Goog-Api-Key": "caller-key"}, + ) + configured: Final = _sync_call(endpoints, client, api_key="deployment-key", custom_llm_provider="gemini") + + assert (caller[2], configured[2]) == ("cachedContents/caller-project", "cachedContents/deployment-project") + + +def test_vertex_caches_behind_one_custom_api_base_are_not_shared_across_projects(endpoints: ContextCachingEndpoints): + client: Final = MagicMock(spec=HTTPHandler) + client.get.return_value = _response({}) + client.post.side_effect = [ + _response(_created("cachedContents/project-a", _expire_time(60))), + _response(_created("cachedContents/project-b", _expire_time(60))), + ] + + sent: Final = [ + _sync_call(endpoints, client, api_base=_GATEWAY, vertex_project=project)[2] + for project in ("project-a", "project-b", "project-a") + ] + + assert sent == ["cachedContents/project-a", "cachedContents/project-b", "cachedContents/project-a"] + + +@pytest.mark.asyncio +async def test_async_vertex_caches_behind_one_custom_api_base_are_not_shared_across_projects( + endpoints: ContextCachingEndpoints, +): + client: Final = MagicMock(spec=AsyncHTTPHandler) + client.get = AsyncMock(return_value=_response({})) + client.post = AsyncMock( + side_effect=[ + _response(_created("cachedContents/project-a", _expire_time(60))), + _response(_created("cachedContents/project-b", _expire_time(60))), + ] + ) + + sent: Final = [ + (await _async_call(endpoints, client, api_base=_GATEWAY, vertex_project=project))[2] + for project in ("project-a", "project-b", "project-a") + ] + + assert sent == ["cachedContents/project-a", "cachedContents/project-b", "cachedContents/project-a"] + + +def test_different_prefixes_do_not_share_a_cache(endpoints: ContextCachingEndpoints): + 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(endpoints, client, messages=_messages("first"))[2] == "cachedContents/a" + assert _sync_call(endpoints, client, messages=_messages("second"))[2] == "cachedContents/b" + + +@pytest.mark.asyncio +async def test_concurrent_requests_for_one_prefix_list_and_create_once(endpoints: ContextCachingEndpoints): + 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(endpoints, 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_async_created_cache_is_reused_by_a_later_request(endpoints: ContextCachingEndpoints): + client: Final = MagicMock(spec=AsyncHTTPHandler) + client.get = AsyncMock(return_value=_response({})) + client.post = AsyncMock(return_value=_response(_created("cachedContents/1", _expire_time(60)))) + + first: Final = await _async_call(endpoints, client) + client.get = AsyncMock(side_effect=AssertionError("a known cache must not be listed again")) + second: Final = await _async_call(endpoints, client) + + assert (first[2], second[2]) == ("cachedContents/1", "cachedContents/1") + + +@pytest.mark.asyncio +async def test_async_listed_cache_past_its_expire_time_is_not_used(endpoints: ContextCachingEndpoints): + client: Final = MagicMock(spec=AsyncHTTPHandler) + 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 = AsyncMock( + return_value=_response( + {"cachedContents": [{"name": "cachedContents/stale", "displayName": key, "expireTime": _expire_time(-1)}]} + ) + ) + client.post = AsyncMock(return_value=_response(_created("cachedContents/fresh", _expire_time(60)))) + + assert (await _async_call(endpoints, client))[2] == "cachedContents/fresh" + + +@pytest.mark.asyncio +async def test_async_create_http_error_surfaces_its_status_code(endpoints: ContextCachingEndpoints): + client: Final = MagicMock(spec=AsyncHTTPHandler) + client.get = AsyncMock(return_value=_response({})) + request: Final = httpx.Request("POST", "https://vertex.test/v1/projects/p/locations/us/cachedContents") + failed: Final = httpx.Response(429, request=request, text="RESOURCE_EXHAUSTED") + client.post = AsyncMock(return_value=failed) + + with pytest.raises(VertexAIError) as raised: + await _async_call(endpoints, client) + + assert raised.value.status_code == 429 + + +@pytest.mark.asyncio +async def test_waiters_resolve_on_their_own_when_the_first_request_fails(endpoints: ContextCachingEndpoints): + 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(endpoints, client), _async_call(endpoints, client), return_exceptions=True + ) + + assert isinstance(results[0], VertexAIError) + assert results[1][2] == "cachedContents/2" + + +@pytest.mark.parametrize("waiter_timeout", [0.05, httpx.Timeout(30.0, read=0.05)]) +@pytest.mark.asyncio +async def test_waiter_resolves_on_its_own_once_its_timeout_passes( + endpoints: ContextCachingEndpoints, waiter_timeout: float | httpx.Timeout +): + client: Final = MagicMock(spec=AsyncHTTPHandler) + + async def slow_list(**_): + await asyncio.sleep(0.5) + return _response({}) + + client.get = AsyncMock(side_effect=slow_list) + client.post = AsyncMock(return_value=_response(_created("cachedContents/1", _expire_time(60)))) + + first, waiter = await asyncio.gather(_async_call(endpoints, client), _async_call(endpoints, client, waiter_timeout)) + + assert (first[2], waiter[2]) == ("cachedContents/1", "cachedContents/1") + assert client.get.await_count == 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_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() + ) + + +def test_cache_expiring_within_the_margin_is_not_remembered(): + known: Final = KnownCachedContents(clock=lambda: _NOW) + known.remember(_SCOPE, "key", "cachedContents/soon", _expire_time(3)) + + assert known.get(_SCOPE, "key") is None + + +def test_registry_is_bounded(): + known: Final = KnownCachedContents(clock=lambda: _NOW) + with patch(f"{_REGISTRY_MODULE}.MAX_KNOWN_CACHED_CONTENTS", 3): + for index in range(5): + known.remember(_SCOPE, f"key-{index}", f"cachedContents/{index}", _expire_time(60)) + + assert known.get(_SCOPE, "key-0") is None + assert known.get(_SCOPE, "key-4") == "cachedContents/4" + + +def _generated(stream: bool) -> httpx.Response: + body: Final = { + "candidates": [{"content": {"role": "model", "parts": [{"text": "ok"}]}, "finishReason": "STOP"}], + "usageMetadata": {"promptTokenCount": 10, "candidatesTokenCount": 1, "totalTokenCount": 11}, + } + if stream: + return httpx.Response(200, text=f"data: {json.dumps(body)}\n\n", headers={"content-type": "text/event-stream"}) + return httpx.Response(200, json=body) + + +def _rejected(status_code: int, message: str) -> httpx.Response: + return httpx.Response(status_code, json={"error": {"code": status_code, "message": message}}) + + +async def _complete(sync_mode: bool, stream: bool, **kwargs) -> str: + request: Final = {"model": "gemini/gemini-3.1-pro", "messages": _messages(), "stream": stream, **kwargs} + if sync_mode: + response = litellm.completion(**request) + if stream: + return "".join(chunk.choices[0].delta.content or "" for chunk in response) + return response.choices[0].message.content + response = await litellm.acompletion(**request) + if stream: + return "".join([chunk.choices[0].delta.content or "" async for chunk in response]) + return response.choices[0].message.content + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.parametrize("stream", [True, False]) +@pytest.mark.asyncio +async def test_rejected_cache_is_replaced_when_a_request_param_cannot_be_deep_copied( + respx_mock: respx.MockRouter, + fake_provider_credentials: None, + monkeypatch: pytest.MonkeyPatch, + sync_mode: bool, + stream: bool, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + respx_mock.get(url__regex=r".*/cachedContents.*").mock(return_value=httpx.Response(200, json={})) + respx_mock.post(url__regex=r".*/cachedContents$").mock( + side_effect=[ + httpx.Response(200, json=_created("cachedContents/deleted", _FAR_FUTURE)), + httpx.Response(200, json=_created("cachedContents/fresh", _FAR_FUTURE)), + ] + ) + generate_route: Final = respx_mock.post(url__regex=r".*:(stream)?[gG]enerateContent.*").mock( + side_effect=[_generated(stream), _rejected(404, "CachedContent not found"), _generated(stream)] + ) + lock: Final = threading.Lock() + + assert await _complete(sync_mode, stream, lock_holder=lock) == "ok" + assert await _complete(sync_mode, stream, lock_holder=lock) == "ok" + + sent: Final = [json.loads(call.request.content)["cachedContent"] for call in generate_route.calls] + assert sent == ["cachedContents/deleted", "cachedContents/deleted", "cachedContents/fresh"] + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.parametrize("stream", [True, False]) +@pytest.mark.parametrize( + ("status_code", "message", "replaced"), + [ + (400, "Invalid resource state for cache content 4484391867902853120.", True), + (403, "CachedContent not found (or permission denied)", True), + (404, "CachedContent not found", True), + (400, "Invalid value at 'generation_config.temperature'", False), + (500, "CachedContent not found", False), + ], +) +@pytest.mark.asyncio +async def test_remembered_cache_rejected_by_generate_content_is_replaced_within_the_same_call( + respx_mock: respx.MockRouter, + fake_provider_credentials: None, + monkeypatch: pytest.MonkeyPatch, + sync_mode: bool, + stream: bool, + status_code: int, + message: str, + replaced: 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( + side_effect=[ + httpx.Response(200, json=_created("cachedContents/deleted", _FAR_FUTURE)), + httpx.Response(200, json=_created("cachedContents/fresh", _FAR_FUTURE)), + ] + ) + generate_route: Final = respx_mock.post(url__regex=r".*:(stream)?[gG]enerateContent.*").mock( + side_effect=[_generated(stream), _rejected(status_code, message), _generated(stream)] + ) + logged: Final[list[dict]] = [] + + assert await _complete(sync_mode, stream) == "ok" + if replaced: + assert await _complete(sync_mode, stream, logger_fn=logged.append) == "ok" + else: + with pytest.raises(openai.APIStatusError): + await _complete(sync_mode, stream, logger_fn=logged.append) + + sent: Final = [json.loads(call.request.content)["cachedContent"] for call in generate_route.calls] + assert sent == ["cachedContents/deleted", "cachedContents/deleted"] + (["cachedContents/fresh"] if replaced else []) + assert list_route.call_count == (2 if replaced else 1) + assert "enerateContent" in logged[-1]["litellm_params"]["api_base"] + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.parametrize("stream", [True, False]) +@pytest.mark.parametrize( + ("status_code", "message", "forgotten"), + [ + (400, "Invalid value at 'generation_config.temperature'", True), + ( + 400, + "Tool config, tools and system instruction should not be set in therequest when using cached content.", + True, + ), + ( + 400, + "CachedContent can not be used with GenerateContent request setting system_instruction, tools or tool_config.", + True, + ), + (404, "Requested entity was not found.", True), + (429, "Resource has been exhausted (e.g. check quota).", False), + (500, "Internal error encountered.", False), + ], +) +@pytest.mark.asyncio +async def test_error_that_does_not_name_a_missing_cache_is_not_retried( + respx_mock: respx.MockRouter, + fake_provider_credentials: None, + monkeypatch: pytest.MonkeyPatch, + sync_mode: bool, + stream: bool, + status_code: int, + message: str, + forgotten: 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( + side_effect=[ + httpx.Response(200, json=_created("cachedContents/first", _FAR_FUTURE)), + httpx.Response(200, json=_created("cachedContents/second", _FAR_FUTURE)), + ] + ) + generate_route: Final = respx_mock.post(url__regex=r".*:(stream)?[gG]enerateContent.*").mock( + side_effect=[_generated(stream), _rejected(status_code, message), _generated(stream)] + ) + + assert await _complete(sync_mode, stream) == "ok" + with pytest.raises(openai.APIStatusError): + await _complete(sync_mode, stream) + assert await _complete(sync_mode, stream) == "ok" + + third: Final = "cachedContents/second" if forgotten else "cachedContents/first" + assert [json.loads(call.request.content)["cachedContent"] for call in generate_route.calls] == [ + "cachedContents/first", + "cachedContents/first", + third, + ] + assert list_route.call_count == (2 if forgotten else 1) + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.parametrize("stream", [True, False]) +@pytest.mark.asyncio +async def test_rejected_caller_supplied_cached_content_is_not_sent_again( + respx_mock: respx.MockRouter, + fake_provider_credentials: None, + monkeypatch: pytest.MonkeyPatch, + sync_mode: bool, + stream: 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={}) + ) + generate_route: Final = respx_mock.post(url__regex=r".*:(stream)?[gG]enerateContent.*").mock( + side_effect=[_rejected(404, "CachedContent not found"), _generated(stream)] + ) + + with pytest.raises(openai.APIStatusError): + await _complete(sync_mode, stream, cached_content="cachedContents/callers") + + assert [json.loads(call.request.content)["cachedContent"] for call in generate_route.calls] == [ + "cachedContents/callers" + ] + assert list_route.call_count == 0 + + +def _cut_off_stream() -> httpx.Response: + body: Final = { + "candidates": [ + { + "content": {"role": "model", "parts": [{"text": "ok"}]}, + "finishReason": "CONTINUATION", + "continuationToken": "dG9rLTE=", + } + ], + "usageMetadata": {"promptTokenCount": 10, "candidatesTokenCount": 1, "totalTokenCount": 11}, + } + return httpx.Response(200, text=f"data: {json.dumps(body)}\n\n", headers={"content-type": "text/event-stream"}) + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +async def test_cache_rejected_on_a_resume_hop_is_not_sent_again( + respx_mock: respx.MockRouter, + fake_provider_credentials: None, + monkeypatch: pytest.MonkeyPatch, + sync_mode: 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( + side_effect=[ + httpx.Response(200, json=_created("cachedContents/deleted", _FAR_FUTURE)), + httpx.Response(200, json=_created("cachedContents/fresh", _FAR_FUTURE)), + ] + ) + generate_route: Final = respx_mock.post(url__regex=r".*:streamGenerateContent.*").mock( + side_effect=[_cut_off_stream(), _rejected(404, "CachedContent not found"), _generated(True)] + ) + + await _complete(sync_mode, True, max_tokens=100_000) + await _complete(sync_mode, True, max_tokens=100_000) + + sent: Final = [json.loads(call.request.content)["cachedContent"] for call in generate_route.calls] + assert sent == ["cachedContents/deleted", "cachedContents/deleted", "cachedContents/fresh"] + assert list_route.call_count == 2