From 184ec7b89449394d00cb7d7c50a7afee96d4e927 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 29 Oct 2026 03:46:35 -0700 Subject: [PATCH] fix(vertex_ai): build the cachedContents collection URL under a custom api_base (#44783) * fix(gemini): use collection endpoint for custom context cache * fix(vertex_ai): build the cachedContents collection URL under a custom api_base * fix(vertex_ai): keep a gateway prefix that contains /publishers/ * test(integration): expect the cachedContents collection URL on a custom Gemini api_base * fix(vertex_ai): cut cache collection URLs at the base's location, graft a mount prefix * fix(vertex_ai): cache under a publisher model api_base in that model's location A Vertex AI api_base that names an endpoint or no model keeps its old cache URL, so no cache is created that the chat call cannot use. The cache body takes its project and location from the api_base model, matching the collection URL. * fix(vertex_ai): cache for the api_base's model and skip caching on a Gemini v1 base A Vertex api_base naming a publisher model now gives the cache its model as well as its project and location, so the cache matches the model generateContent calls. A Gemini api_base ending in /v1 skips context caching, since the v1 API has no cachedContents and rejects the cachedContent field. * fix(vertex_ai): check the cache token minimum against the api_base's model The cache is created for the model a publisher-model api_base names, so the minimum prompt size now comes from that model too. A prompt below its minimum is sent uncached instead of failing at cache create. * fix(vertex_ai): use the stricter cache minimum and send a Gemini base's 404 uncached --------- Co-authored-by: paulo037 Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../vertex_ai_context_caching.py | 238 +++--- litellm/llms/vertex_ai/vertex_llm_base.py | 51 +- litellm/types/llms/vertex_ai.py | 2 +- .../cost_calculation/cases/gemini.py | 4 +- ...test_gemini_messages_cache_control_wire.py | 8 +- ...test_cache_control_tool_call_marks_wire.py | 4 +- ...ni_context_caching_custom_api_base_wire.py | 130 +++ .../test_vertex_ai_context_caching.py | 802 ++++++++++-------- 8 files changed, 767 insertions(+), 472 deletions(-) create mode 100644 tests/integration/providers/test_gemini_context_caching_custom_api_base_wire.py 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 5bff43828d6..884c82fe331 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,5 @@ from collections.abc import Sequence +from dataclasses import dataclass from typing import Final, Literal import httpx @@ -8,7 +9,6 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger from litellm.caching.caching import Cache, LiteLLMCacheType -from litellm.constants import MINIMUM_PROMPT_CACHE_TOKEN_COUNT from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, @@ -16,10 +16,11 @@ from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, ) from litellm.llms.openai.openai import AllMessageValues -from litellm.utils import is_prompt_caching_valid_prompt +from litellm.types.llms.openai import ChatCompletionToolParam +from litellm.utils import get_prompt_cache_min_tokens, is_prompt_caching_valid_prompt from ..common_utils import VertexAIError, get_vertex_base_url -from ..vertex_llm_base import VertexBase +from ..vertex_llm_base import VertexBase, publisher_model_location from .known_cached_contents import KnownCachedContents, known_cached_contents from .transformation import ( cached_messages_end_on_supported_turn, @@ -65,6 +66,19 @@ class _CreatedCachedContentRoot(RootModel[_CreatedCachedContent]): model_config = ConfigDict(frozen=True, hide_input_in_errors=True) +_TOOL_PARAMS: Final = ("tools", "tool_choice") + + +def _sync_cache_client(timeout: float | httpx.Timeout | None) -> HTTPHandler: + if timeout is None: + return HTTPHandler() + return HTTPHandler(timeout=httpx.Timeout(timeout) if isinstance(timeout, (float, int)) else timeout) + + +class _CachedContentsNotServedError(VertexAIError): + pass + + def _record_cache_creation_tokens(logging_obj: Logging, created: _CreatedCachedContent) -> None: tokens: Final = (created.get("usageMetadata") or {}).get("totalTokenCount") or 0 logging_obj.vertex_ai_cache_creation_input_tokens = max(tokens, 0) # rebind-ok: set provider count @@ -78,6 +92,61 @@ def _wait_seconds(timeout: float | httpx.Timeout | None) -> float | None: return timeout +@dataclass(frozen=True, slots=True) +class _CacheTarget: + model: str + project: str | None + location: str | None + + +def _cache_target( + custom_llm_provider: Literal["vertex_ai", "vertex_ai_beta", "gemini"], + api_base: str | None, + model: str, + vertex_project: str | None, + vertex_location: str | None, +) -> _CacheTarget: + model_location: Final = ( + None if custom_llm_provider == "gemini" or api_base is None else publisher_model_location(api_base) + ) + if model_location is None: + return _CacheTarget(model=model, project=vertex_project, location=vertex_location) + return _CacheTarget(model=model_location.model, project=model_location.project, location=model_location.location) + + +def _prompt_cache_min_tokens(models: Sequence[str], custom_llm_provider: str) -> int: + return max( + get_prompt_cache_min_tokens( + model if model.startswith(custom_llm_provider) else f"{custom_llm_provider}/{model}" + ) + for model in models + ) + + +def _context_caching_skip_reason( + cached_messages: Sequence[AllMessageValues], + tools: Sequence[ChatCompletionToolParam] | None, + deployment_model: str, + cache_model: str, + custom_llm_provider: Literal["vertex_ai", "vertex_ai_beta", "gemini"], +) -> str | None: + if not cached_messages_end_on_supported_turn(cached_messages): + return ( + "Vertex AI context caching: cached message block ends on a model turn once " + "system messages are extracted, which the cachedContents API rejects." + ) + min_tokens: Final = _prompt_cache_min_tokens((deployment_model, cache_model), custom_llm_provider) + if not is_prompt_caching_valid_prompt( + model=cache_model, + messages=list(cached_messages), + tools=None if tools is None else list(tools), + custom_llm_provider=custom_llm_provider, + min_token_count=min_tokens, + ): + return f"Vertex AI context caching: cached content is below minimum token count ({min_tokens})." + return None + + class ContextCachingEndpoints(VertexBase): """ Covers context caching endpoints for Vertex AI + Google AI Studio @@ -135,6 +204,7 @@ class ContextCachingEndpoints(VertexBase): vertex_project=vertex_project, vertex_location=vertex_location, vertex_api_version=("v1beta1" if custom_llm_provider == "vertex_ai_beta" else "v1"), + collection_endpoint=True, ) def check_cache( @@ -206,6 +276,8 @@ class ContextCachingEndpoints(VertexBase): except httpx.HTTPStatusError as e: if e.response.status_code == 403: return None + if e.response.status_code == 404 and custom_llm_provider == "gemini": + raise _CachedContentsNotServedError(status_code=404, message=e.response.text) raise VertexAIError(status_code=e.response.status_code, message=e.response.text) except Exception as e: raise VertexAIError(status_code=500, message=str(e)) @@ -305,6 +377,8 @@ class ContextCachingEndpoints(VertexBase): except httpx.HTTPStatusError as e: if e.response.status_code == 403: return None + if e.response.status_code == 404 and custom_llm_provider == "gemini": + raise _CachedContentsNotServedError(status_code=404, message=e.response.text) raise VertexAIError(status_code=e.response.status_code, message=e.response.text) except Exception as e: raise VertexAIError(status_code=500, message=str(e)) @@ -370,41 +444,25 @@ class ContextCachingEndpoints(VertexBase): if len(cached_messages) == 0: return messages, optional_params, None - if not cached_messages_end_on_supported_turn(cached_messages): - verbose_logger.debug( - "Vertex AI context caching: cached message block ends on a model turn once " - "system messages are extracted, which the cachedContents API rejects. " - "Skipping context caching." - ) + target: Final = _cache_target(custom_llm_provider, api_base, model, vertex_project, vertex_location) + tools: Final = optional_params.get("tools") + tool_choice: Final = optional_params.get("tool_choice") + skip_reason: Final = _context_caching_skip_reason( + cached_messages, tools, model, target.model, custom_llm_provider + ) + if skip_reason is not None: + verbose_logger.debug("%s Skipping context caching.", skip_reason) return messages, optional_params, None - # Gemini requires a minimum of 1024 tokens for context caching. - # Skip caching if the cached content is too small to avoid API errors. - if not is_prompt_caching_valid_prompt( - model=model, - messages=cached_messages, - tools=optional_params.get("tools"), - custom_llm_provider=custom_llm_provider, - ): - verbose_logger.debug( - "Vertex AI context caching: cached content is below minimum token " - "count (%d). Skipping context caching.", - MINIMUM_PROMPT_CACHE_TOKEN_COUNT, - ) - return messages, optional_params, None - - tools: Final = optional_params.pop("tools", None) - tool_choice: Final = optional_params.pop("tool_choice", None) - ## AUTHORIZATION ## token, url = self._get_token_and_url_context_caching( gemini_api_key=api_key, custom_llm_provider=custom_llm_provider, api_base=api_base, - vertex_project=vertex_project, - vertex_location=vertex_location, + vertex_project=target.project, + vertex_location=target.location, vertex_auth_header=vertex_auth_header, - model=model, + model=target.model, ) headers: Final = { @@ -417,44 +475,42 @@ class ContextCachingEndpoints(VertexBase): if extra_headers is not None: headers.update(extra_headers) - if client is None or not isinstance(client, HTTPHandler): - _params: Final = {} - if timeout is not None: - if isinstance(timeout, float) or isinstance(timeout, int): - timeout = httpx.Timeout(timeout) - _params["timeout"] = timeout - client = HTTPHandler(**_params) - else: - client = client + http_client: Final = client if isinstance(client, HTTPHandler) else _sync_cache_client(timeout) ## CHECK IF CACHED ALREADY 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 = self.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, + messages=cached_messages, tools=tools, tool_choice=tool_choice, model=target.model ) + try: + google_cache_name: Final = self.check_cache( + cache_key=generated_cache_key, + client=http_client, + headers=headers, + api_key=api_key, + api_base=api_base, + logging_obj=logging_obj, + custom_llm_provider=custom_llm_provider, + vertex_project=target.project, + vertex_location=target.location, + vertex_auth_header=vertex_auth_header, + model=target.model, + ) + except _CachedContentsNotServedError: + verbose_logger.debug("Gemini context caching: %s answers 404. Skipping context caching.", url) + return messages, optional_params, None + for tool_param in _TOOL_PARAMS: + optional_params.pop(tool_param, None) 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, + model=target.model, messages=cached_messages, cache_key=generated_cache_key, custom_llm_provider=custom_llm_provider, - vertex_project=vertex_project, - vertex_location=vertex_location, + vertex_project=target.project, + vertex_location=target.location, ) cached_content_request_body["tools"] = tools @@ -473,7 +529,7 @@ class ContextCachingEndpoints(VertexBase): ) try: - response: Final = client.post( + response: Final = http_client.post( url=url, headers=headers, json=cached_content_request_body, @@ -489,7 +545,7 @@ class ContextCachingEndpoints(VertexBase): _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 + url, custom_llm_provider, headers, vertex_project=target.project, vertex_location=target.location ), generated_cache_key, raw_response_cached.get("name"), @@ -536,41 +592,25 @@ class ContextCachingEndpoints(VertexBase): if len(cached_messages) == 0: return messages, optional_params, None - if not cached_messages_end_on_supported_turn(cached_messages): - verbose_logger.debug( - "Vertex AI context caching: cached message block ends on a model turn once " - "system messages are extracted, which the cachedContents API rejects. " - "Skipping context caching." - ) + target: Final = _cache_target(custom_llm_provider, api_base, model, vertex_project, vertex_location) + tools: Final = optional_params.get("tools") + tool_choice: Final = optional_params.get("tool_choice") + skip_reason: Final = _context_caching_skip_reason( + cached_messages, tools, model, target.model, custom_llm_provider + ) + if skip_reason is not None: + verbose_logger.debug("%s Skipping context caching.", skip_reason) return messages, optional_params, None - # Gemini requires a minimum of 1024 tokens for context caching. - # Skip caching if the cached content is too small to avoid API errors. - if not is_prompt_caching_valid_prompt( - model=model, - messages=cached_messages, - tools=optional_params.get("tools"), - custom_llm_provider=custom_llm_provider, - ): - verbose_logger.debug( - "Vertex AI context caching: cached content is below minimum token " - "count (%d). Skipping context caching.", - MINIMUM_PROMPT_CACHE_TOKEN_COUNT, - ) - return messages, optional_params, None - - tools: Final = optional_params.pop("tools", None) - tool_choice: Final = optional_params.pop("tool_choice", None) - ## AUTHORIZATION ## token, url = self._get_token_and_url_context_caching( gemini_api_key=api_key, custom_llm_provider=custom_llm_provider, api_base=api_base, - vertex_project=vertex_project, - vertex_location=vertex_location, + vertex_project=target.project, + vertex_location=target.location, vertex_auth_header=vertex_auth_header, - model=model, + model=target.model, ) headers: Final = { @@ -590,10 +630,10 @@ class ContextCachingEndpoints(VertexBase): ## CHECK IF CACHED ALREADY generated_cache_key: Final = local_cache_obj.get_cache_key( - messages=cached_messages, tools=tools, tool_choice=tool_choice, model=model + messages=cached_messages, tools=tools, tool_choice=tool_choice, model=target.model ) scope: Final = KnownCachedContents.scope( - url, custom_llm_provider, headers, vertex_project=vertex_project, vertex_location=vertex_location + url, custom_llm_provider, headers, vertex_project=target.project, vertex_location=target.location ) async def lookup_or_create() -> str | None: @@ -605,10 +645,10 @@ class ContextCachingEndpoints(VertexBase): api_base=api_base, logging_obj=logging_obj, custom_llm_provider=custom_llm_provider, - vertex_project=vertex_project, - vertex_location=vertex_location, + vertex_project=target.project, + vertex_location=target.location, vertex_auth_header=vertex_auth_header, - model=model, + model=target.model, ) if google_cache_name: @@ -616,12 +656,12 @@ class ContextCachingEndpoints(VertexBase): ## TRANSFORM REQUEST cached_content_request_body: Final = transform_openai_messages_to_gemini_context_caching( - model=model, + model=target.model, messages=cached_messages, cache_key=generated_cache_key, custom_llm_provider=custom_llm_provider, - vertex_project=vertex_project, - vertex_location=vertex_location, + vertex_project=target.project, + vertex_location=target.location, ) cached_content_request_body["tools"] = tools @@ -659,9 +699,15 @@ class ContextCachingEndpoints(VertexBase): ) 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) - ) + try: + cache_name: Final = await self._known_caches.single_flight( + scope, generated_cache_key, lookup_or_create, _wait_seconds(timeout) + ) + except _CachedContentsNotServedError: + verbose_logger.debug("Gemini context caching: %s answers 404. Skipping context caching.", url) + return messages, optional_params, None + for tool_param in _TOOL_PARAMS: + optional_params.pop(tool_param, None) return non_cached_messages, optional_params, cache_name def get_cache(self): diff --git a/litellm/llms/vertex_ai/vertex_llm_base.py b/litellm/llms/vertex_ai/vertex_llm_base.py index 0cdd92bb071..7dd98212f76 100644 --- a/litellm/llms/vertex_ai/vertex_llm_base.py +++ b/litellm/llms/vertex_ai/vertex_llm_base.py @@ -6,9 +6,11 @@ Handles Authentication and generating request urls for Vertex AI and Google AI S import asyncio import json +import re import threading from collections import OrderedDict from collections.abc import Mapping +from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Final, Literal, Protocol from urllib.parse import urlparse @@ -46,6 +48,31 @@ def graft_default_vertex_path(api_base: str, default_url: str) -> str: _graft_default_vertex_path = graft_default_vertex_path + +_PUBLISHER_MODEL_LOCATION: Final = re.compile(r".*/projects/[^/]+/locations/[^/]+(?=/publishers/[^/]+/models/[^/]+/?$)") + + +@dataclass(frozen=True, slots=True) +class PublisherModelLocation: + parent_url: str + project: str + location: str + model: str + + +def publisher_model_location(api_base: str) -> PublisherModelLocation | None: + match: Final = _PUBLISHER_MODEL_LOCATION.match(api_base) + if match is None: + return None + _, _, project, _, location = match.group().rsplit("/", 4) + return PublisherModelLocation( + parent_url=match.group(), + project=project, + location=location, + model=api_base.rstrip("/").rsplit("/", 1)[-1], + ) + + GOOGLE_IMPORT_ERROR_MESSAGE: Final = ( "Google Cloud SDK not found. Install it with: pip install 'litellm[google]' or pip install google-cloud-aiplatform" ) @@ -637,15 +664,18 @@ class VertexBase: vertex_location: str | None = None, vertex_api_version: Literal["v1", "v1beta1"] | None = None, use_psc_endpoint_format: bool = False, + collection_endpoint: bool = False, ) -> tuple[str | None, str]: """ for cloudflare ai gateway - https://github.com/BerriAI/litellm/issues/4317 Handles custom api_base for: - 1. Gemini (Google AI Studio) - constructs /models/{model}:{endpoint} + 1. Gemini (Google AI Studio) - constructs /models/{model}:{endpoint}, or {api_base}/{endpoint} + for a collection endpoint 2. Vertex AI with standard proxies - grafts the default vertex URL path onto the api_base when its path is empty or only an API version (/v1, /v1beta1); - otherwise constructs {api_base}:{endpoint} + otherwise constructs {api_base}:{endpoint}, or, for a collection endpoint under an api_base + naming a publisher model, the collection under that model's projects/{project}/locations/{location} 3. Vertex AI with PSC endpoints - constructs full path structure {api_base}/v1/projects/{project}/locations/{location}/endpoints/{model}:{endpoint} (only when use_psc_endpoint_format=True) @@ -653,22 +683,22 @@ class VertexBase: Args: use_psc_endpoint_format: If True, constructs PSC endpoint URL format. If False (default), uses api_base as-is and appends :{endpoint} + collection_endpoint: True when `endpoint` names a collection resource such as + cachedContents rather than a model action such as generateContent ## Returns - (auth_header, url) - Tuple[Optional[str], str] """ if api_base: if custom_llm_provider == "gemini": - # For Gemini (Google AI Studio), construct the full path like other providers - if model is None: + if model is None and not collection_endpoint: raise ValueError("Model parameter is required for Gemini custom API base URLs") - url = join_api_base(api_base, f"models/{model}:{endpoint}") + url = join_api_base(api_base, endpoint if collection_endpoint else f"models/{model}:{endpoint}") if gemini_api_key is None: raise ValueError( "Missing Gemini API key. Set the GEMINI_API_KEY or GOOGLE_API_KEY environment variable." ) - if gemini_api_key is not None: - auth_header = {"x-goog-api-key": gemini_api_key} + auth_header = {"x-goog-api-key": gemini_api_key} else: # For Vertex AI if use_psc_endpoint_format: @@ -692,7 +722,12 @@ class VertexBase: elif urlparse(api_base).path.rstrip("/") in ("/v1", "/v1beta1") and "/projects/" in urlparse(url).path: url = graft_default_vertex_path(api_base=api_base, default_url=url) else: - url = f"{api_base.rstrip('/')}:{endpoint}" + model_location: Final = publisher_model_location(api_base) if collection_endpoint else None + url = ( + f"{api_base.rstrip('/')}:{endpoint}" + if model_location is None + else join_api_base(model_location.parent_url, endpoint) + ) if stream is True: parsed_stream_url: Final = urlparse(url) stream_query: Final = f"{parsed_stream_url.query}&alt=sse" if parsed_stream_url.query else "alt=sse" diff --git a/litellm/types/llms/vertex_ai.py b/litellm/types/llms/vertex_ai.py index f1176f777c7..7cb5432d009 100644 --- a/litellm/types/llms/vertex_ai.py +++ b/litellm/types/llms/vertex_ai.py @@ -376,7 +376,7 @@ class RequestBody(TypedDict, total=False): class CachedContentRequestBody(TypedDict, total=False): contents: Required[list[ContentType]] system_instruction: SystemInstructions - tools: Tools + tools: Tools | None toolConfig: ToolConfig model: Required[str] # Format: models/{model} ttl: str # ending in 's' - Example: "3.5s". diff --git a/tests/integration/cost_calculation/cases/gemini.py b/tests/integration/cost_calculation/cases/gemini.py index 86a0d3b9067..9fe4b1c50d5 100644 --- a/tests/integration/cost_calculation/cases/gemini.py +++ b/tests/integration/cost_calculation/cases/gemini.py @@ -3348,10 +3348,10 @@ GEMINI_GEMINI_3_1_PRO_CACHE_CREATION: Final = CostTrackingTestCase( response=RoutedResponse( content_type="application/x-routed", routes={ - "GET /models/gemini-3.1-pro:cachedContents": JsonResponse( + "GET /cachedContents": JsonResponse( content_type="application/json", body={"cachedContents": []} ), - "POST /models/gemini-3.1-pro:cachedContents": JsonResponse( + "POST /cachedContents": JsonResponse( content_type="application/json", body={ "name": "cachedContents/$REQUEST_ID", diff --git a/tests/integration/messages_endpoint/providers/gemini/test_gemini_messages_cache_control_wire.py b/tests/integration/messages_endpoint/providers/gemini/test_gemini_messages_cache_control_wire.py index 71073b3dd5e..cd53b54e601 100644 --- a/tests/integration/messages_endpoint/providers/gemini/test_gemini_messages_cache_control_wire.py +++ b/tests/integration/messages_endpoint/providers/gemini/test_gemini_messages_cache_control_wire.py @@ -39,11 +39,11 @@ def test_gemini_messages_cache_control_creates_cached_content_and_generates_from def respond(request: Request) -> Reply: assert request.headers["x-goog-api-key"] == _API_KEY, request.headers if request.method == "GET": - assert request.target == f"/models/{_BACKEND}:cachedContents", request.target + assert request.target == "/cachedContents", request.target return Reply(body=b"{}") assert request.method == "POST", request.method body: Final = _JSON_OBJECT.validate_json(request.body) - if request.target == f"/models/{_BACKEND}:cachedContents": + if request.target == "/cachedContents": assert isinstance(body["displayName"], str) and body["displayName"], body assert body == { "contents": [{"role": "user", "parts": [{"text": "."}]}], @@ -80,7 +80,7 @@ def test_gemini_messages_cache_control_creates_cached_content_and_generates_from payload: Final = _JSON_OBJECT.validate_json(response.content) assert payload["content"] == [{"type": "text", "text": "The policy applies."}], response.text assert [(request.method, request.target) for request in wire.drain()] == [ - ("GET", f"/models/{_BACKEND}:cachedContents"), - ("POST", f"/models/{_BACKEND}:cachedContents"), + ("GET", "/cachedContents"), + ("POST", "/cachedContents"), ("POST", f"/models/{_BACKEND}:generateContent"), ] diff --git a/tests/integration/providers/test_cache_control_tool_call_marks_wire.py b/tests/integration/providers/test_cache_control_tool_call_marks_wire.py index 6389d90597e..a18224a9f36 100644 --- a/tests/integration/providers/test_cache_control_tool_call_marks_wire.py +++ b/tests/integration/providers/test_cache_control_tool_call_marks_wire.py @@ -406,8 +406,8 @@ def _gemini_peer(request: Request) -> Reply: _GEMINI: Final = "gemini/gemini-3.8-flash" _GEMINI_CACHE_WRITE: Final = [ - ("GET", "/models/gemini-3.8-flash:cachedContents"), - ("POST", "/models/gemini-3.8-flash:cachedContents"), + ("GET", "/cachedContents"), + ("POST", "/cachedContents"), ("POST", "/models/gemini-3.8-flash:generateContent"), ] diff --git a/tests/integration/providers/test_gemini_context_caching_custom_api_base_wire.py b/tests/integration/providers/test_gemini_context_caching_custom_api_base_wire.py new file mode 100644 index 00000000000..b76f0cd162b --- /dev/null +++ b/tests/integration/providers/test_gemini_context_caching_custom_api_base_wire.py @@ -0,0 +1,130 @@ +import json +from collections.abc import Callable +from typing import Final + +import pytest + +from integration._support.client import Gateway +from integration._support.vertex import service_account_json +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue, TypeAdapter + +_BACKEND: Final = "gemini-3.8-flash" +_PROJECT: Final = "scripted-project" +_LOCATION: Final = "us-central1" +_MODEL_LOCATION: Final = "europe-west4" +_VERTEX_MODEL_RESOURCE: Final = f"projects/{_PROJECT}/locations/{_MODEL_LOCATION}/publishers/google/models/{_BACKEND}" +_VERTEX_COLLECTION: Final = f"/v1/projects/{_PROJECT}/locations/{_MODEL_LOCATION}/cachedContents" +_GEMINI_COLLECTION: Final = "/v1beta/cachedContents" +_GEMINI_MODEL_PATH: Final = f"/v1beta/models/{_BACKEND}" +_GEMINI_KEY: Final = "scripted-gemini-key" +_CACHE_NAME: Final = "cachedContents/scripted-cache" +_CACHED_TOKENS: Final = 5120 +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_CACHED_TEXT: Final = " ".join( + f"Policy clause {index}: the gateway keeps this clause in the cache." for index in range(700) +) +_MESSAGES: Final = ( + { + "role": "user", + "content": [{"type": "text", "text": _CACHED_TEXT, "cache_control": {"type": "ephemeral"}}], + }, + {"role": "user", "content": "How many clauses are there?"}, +) + + +def _generate_content_reply() -> Reply: + return Reply( + body=json.dumps( + { + "candidates": [{"content": {"role": "model", "parts": [{"text": "700"}]}, "finishReason": "STOP"}], + "usageMetadata": { + "promptTokenCount": _CACHED_TOKENS + 9, + "cachedContentTokenCount": _CACHED_TOKENS, + "candidatesTokenCount": 2, + "totalTokenCount": _CACHED_TOKENS + 11, + }, + "modelVersion": _BACKEND, + } + ).encode() + ) + + +def _google_responder(collection: str, model_path: str) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + if request.target == collection and request.method == "GET": + return Reply(body=b'{"cachedContents": []}') + if request.target == collection and request.method == "POST": + return Reply(body=json.dumps({"name": _CACHE_NAME, "model": f"models/{_BACKEND}"}).encode()) + assert request.target == f"{model_path}:generateContent", f"{request.method} {request.target}" + return _generate_content_reply() + + return respond + + +def _chat(gateway: Gateway, model: str) -> dict[str, JsonValue]: + response: Final = gateway.client.post( + "/v1/chat/completions", + json={"model": model, "messages": list(_MESSAGES), "max_tokens": 32}, + headers={"Authorization": f"Bearer {gateway.key}"}, + timeout=30, + ) + assert response.status_code == 200, response.text + return _JSON_OBJECT.validate_json(response.content) + + +def _assert_cache_round_trip(requests: tuple[Request, ...], collection: str, model_path: str) -> None: + assert tuple((request.method, request.target) for request in requests) == ( + ("GET", collection), + ("POST", collection), + ("POST", f"{model_path}:generateContent"), + ), [(request.method, request.target) for request in requests] + created: Final = _JSON_OBJECT.validate_json(requests[1].body) + assert "contents" in created, created + generated: Final = _JSON_OBJECT.validate_json(requests[2].body) + assert generated["cachedContent"] == _CACHE_NAME, generated + + +def _cached_tokens_of(response: dict[str, JsonValue]) -> JsonValue: + usage: Final = response["usage"] + assert isinstance(usage, dict) + details: Final = usage["prompt_tokens_details"] + assert isinstance(details, dict) + return details["cached_tokens"] + + +def test_gemini_custom_api_base_lists_and_creates_the_cached_contents_collection(gateway: Gateway) -> None: + with ( + wire_server(_google_responder(_GEMINI_COLLECTION, _GEMINI_MODEL_PATH)) as wire, + gateway.scenario() as scenario, + ): + model: Final = scenario.model(model=f"gemini/{_BACKEND}", api_base=f"{wire.url}/v1beta", api_key=_GEMINI_KEY) + response: Final = _chat(gateway, model) + requests: Final = wire.drain() + _assert_cache_round_trip(requests, _GEMINI_COLLECTION, _GEMINI_MODEL_PATH) + assert all(request.headers.get("x-goog-api-key") == _GEMINI_KEY for request in requests) + assert _cached_tokens_of(response) == _CACHED_TOKENS + + +@pytest.mark.parametrize("gateway_prefix", ["", "/vertex-proxy"], ids=["model-path", "prefixed-model-path"]) +def test_vertex_custom_api_base_creates_the_cache_where_its_model_lives(gateway: Gateway, gateway_prefix: str) -> None: + api_base_path: Final = f"{gateway_prefix}/v1/{_VERTEX_MODEL_RESOURCE}" + collection: Final = f"{gateway_prefix}{_VERTEX_COLLECTION}" + with ( + wire_server(_google_responder(collection, api_base_path)) as wire, + gateway.scenario() as scenario, + ): + model: Final = scenario.model( + model=f"vertex_ai/{_BACKEND}", + api_base=f"{wire.url}{api_base_path}", + api_key=None, + vertex_project=_PROJECT, + vertex_location=_LOCATION, + vertex_credentials=service_account_json(_PROJECT, gateway.upstream_url.rstrip("/")), + ) + response: Final = _chat(gateway, model) + requests: Final = wire.drain() + _assert_cache_round_trip(requests, collection, api_base_path) + assert _JSON_OBJECT.validate_json(requests[1].body)["model"] == _VERTEX_MODEL_RESOURCE + assert all(request.headers.get("authorization") == "Bearer scripted-token" for request in requests) + assert _cached_tokens_of(response) == _CACHED_TOKENS diff --git a/tests/unit/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py b/tests/unit/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py index 19e7e2dee52..78b9e930001 100644 --- a/tests/unit/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py +++ b/tests/unit/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py @@ -1,4 +1,4 @@ -from typing import Final +from typing import Final, Literal from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -9,6 +9,7 @@ import litellm from litellm.litellm_core_utils.litellm_logging import Logging 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.known_cached_contents import KnownCachedContents from litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching import ( ContextCachingEndpoints, ) @@ -84,18 +85,10 @@ class TestContextCachingEndpoints: """Teardown for each test method""" self._token_check_patcher.stop() - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj" - ) - def test_check_and_create_cache_with_cached_content( - self, mock_cache_obj, mock_separate, custom_llm_provider - ): + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages") + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj") + def test_check_and_create_cache_with_cached_content(self, mock_cache_obj, mock_separate, custom_llm_provider): """Test check_and_create_cache when cached_content is provided""" # Setup cached_content = "cached_content_123" @@ -130,15 +123,9 @@ class TestContextCachingEndpoints: mock_separate.assert_not_called() mock_cache_obj.get_cache_key.assert_not_called() - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" - ) - def test_check_and_create_cache_no_cached_messages( - self, mock_separate, custom_llm_provider - ): + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages") + def test_check_and_create_cache_no_cached_messages(self, mock_separate, custom_llm_provider): """Test check_and_create_cache when no cached messages are found""" # Setup mock_separate.return_value = ([], self.sample_messages) # No cached messages @@ -168,15 +155,9 @@ class TestContextCachingEndpoints: assert returned_params == optional_params assert returned_cache is None - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj" - ) + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages") + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj") @patch.object(ContextCachingEndpoints, "check_cache") def test_check_and_create_cache_existing_cache_found( self, mock_check_cache, mock_cache_obj, mock_separate, custom_llm_provider @@ -224,15 +205,181 @@ class TestContextCachingEndpoints: model="gemini-1.5-pro", ) - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj" - ) + @pytest.mark.parametrize("vertex_location", ["us-central1", None]) + def test_check_and_create_cache_creates_the_cache_where_the_api_base_model_lives(self, vertex_location: str | None): + context_caching = ContextCachingEndpoints(known_caches=KnownCachedContents()) + listed = MagicMock() + listed.json.return_value = {"cachedContents": []} + created = MagicMock() + created.json.return_value = {"name": "cachedContents/base-cache"} + self.mock_client.get.return_value = listed + self.mock_client.post.return_value = created + base_model_location = "https://gateway.example.com/v1/projects/base-project/locations/europe-west4" + + _, _, cache_name = context_caching.check_and_create_cache( + messages=[ + { + "role": "user", + "content": [{"type": "text", "text": "Long policy text", "cache_control": {"type": "ephemeral"}}], + }, + {"role": "user", "content": "Summarize it"}, + ], + optional_params={}, + api_key="unused", + api_base=f"{base_model_location}/publishers/google/models/gemini-2.5-pro", + model="gemini-2.5-flash", + client=self.mock_client, + timeout=30.0, + logging_obj=self.mock_logging, + custom_llm_provider="vertex_ai", + vertex_project="config-project", + vertex_location=vertex_location, + vertex_auth_header="vertex-token", + ) + + assert cache_name == "cachedContents/base-cache" + assert self.mock_client.get.call_args.kwargs["url"] == f"{base_model_location}/cachedContents" + assert self.mock_client.post.call_args.kwargs["url"] == f"{base_model_location}/cachedContents" + assert ( + self.mock_client.post.call_args.kwargs["json"]["model"] + == "projects/base-project/locations/europe-west4/publishers/google/models/gemini-2.5-pro" + ) + + def test_check_and_create_cache_keeps_one_cache_per_api_base_model(self): + context_caching = ContextCachingEndpoints(known_caches=KnownCachedContents()) + created_bodies: Final = [] + + def create(**kwargs): + created_bodies.append(kwargs["json"]) + response = MagicMock() + response.json.return_value = {"name": f"cachedContents/cache-{len(created_bodies)}"} + return response + + def list_created(**kwargs): + response = MagicMock() + response.json.return_value = { + "cachedContents": [ + {"name": f"cachedContents/cache-{index}", "displayName": body["displayName"]} + for index, body in enumerate(created_bodies, start=1) + ] + } + return response + + self.mock_client.get.side_effect = list_created + self.mock_client.post.side_effect = create + base_model_location = "https://gateway.example.com/v1/projects/base-project/locations/europe-west4" + + def cache_for(base_model: str) -> str | None: + _, _, cache_name = context_caching.check_and_create_cache( + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "Long policy text", "cache_control": {"type": "ephemeral"}} + ], + }, + {"role": "user", "content": "Summarize it"}, + ], + optional_params={}, + api_key="unused", + api_base=f"{base_model_location}/publishers/google/models/{base_model}", + model="gemini-2.5-flash", + client=self.mock_client, + timeout=30.0, + logging_obj=self.mock_logging, + custom_llm_provider="vertex_ai", + vertex_project="config-project", + vertex_location="us-central1", + vertex_auth_header="vertex-token", + ) + return cache_name + + pro_cache: Final = cache_for("gemini-2.5-pro") + lite_cache: Final = cache_for("gemini-2.5-flash-lite") + + assert pro_cache != lite_cache + assert [body["model"].rsplit("/", 1)[-1] for body in created_bodies] == [ + "gemini-2.5-pro", + "gemini-2.5-flash-lite", + ] + assert cache_for("gemini-2.5-pro") == pro_cache + + def _gemini_v1_gateway_call_args(self, list_response: MagicMock): + self.mock_client.get.return_value = list_response + self.mock_async_client.get = AsyncMock(return_value=list_response) + created: Final = MagicMock() + created.json.return_value = {"name": "cachedContents/gateway-cache"} + self.mock_client.post.return_value = created + self.mock_async_client.post = AsyncMock(return_value=created) + messages: Final = [ + { + "role": "user", + "content": [{"type": "text", "text": "Long policy text", "cache_control": {"type": "ephemeral"}}], + }, + {"role": "user", "content": "Summarize it"}, + ] + call_args: Final = dict( + messages=messages, + optional_params={"tools": self.sample_tools}, + api_key="unused", + api_base="https://gateway.example.com/gemini/v1", + model="gemini-2.5-flash", + timeout=30.0, + logging_obj=self.mock_logging, + custom_llm_provider="gemini", + vertex_project=None, + vertex_location=None, + vertex_auth_header=None, + ) + return call_args + + @pytest.mark.parametrize("is_async", [False, True]) + @pytest.mark.asyncio + async def test_check_and_create_cache_sends_uncached_when_the_gemini_api_base_serves_no_cached_contents( + self, is_async: bool + ): + not_served: Final = MagicMock() + not_served.raise_for_status.side_effect = httpx.HTTPStatusError( + "Not Found", request=MagicMock(), response=MagicMock(status_code=404, text="") + ) + + call_args: Final = self._gemini_v1_gateway_call_args(not_served) + + result: Final = ( + await self.context_caching.async_check_and_create_cache(client=self.mock_async_client, **call_args) + if is_async + else self.context_caching.check_and_create_cache(client=self.mock_client, **call_args) + ) + + assert result == (call_args["messages"], {"tools": self.sample_tools}, None) + self.mock_client.post.assert_not_called() + self.mock_async_client.post.assert_not_called() + + @pytest.mark.parametrize("is_async", [False, True]) + @pytest.mark.asyncio + async def test_check_and_create_cache_caches_through_a_gemini_v1_gateway_that_serves_cached_contents( + self, is_async: bool + ): + empty_list: Final = MagicMock() + empty_list.json.return_value = {} + + call_args: Final = self._gemini_v1_gateway_call_args(empty_list) + + _, params, cache_name = ( + await self.context_caching.async_check_and_create_cache(client=self.mock_async_client, **call_args) + if is_async + else self.context_caching.check_and_create_cache(client=self.mock_client, **call_args) + ) + client: Final = self.mock_async_client if is_async else self.mock_client + + assert params == {} + assert cache_name == "cachedContents/gateway-cache" + assert client.get.call_args.kwargs["url"] == "https://gateway.example.com/gemini/v1/cachedContents" + assert client.post.call_args.kwargs["url"].startswith("https://gateway.example.com/gemini/v1/cachedContents") + + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages") + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj") @patch( "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.transform_openai_messages_to_gemini_context_caching" ) @@ -299,15 +446,9 @@ class TestContextCachingEndpoints: assert "tools" in call_args.kwargs["json"] assert call_args.kwargs["json"]["tools"] == self.sample_tools - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj" - ) + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages") + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj") @patch.object(ContextCachingEndpoints, "check_cache") @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching") def test_check_and_create_cache_http_error( @@ -332,9 +473,7 @@ class TestContextCachingEndpoints: mock_response = MagicMock() mock_response.status_code = 400 mock_response.text = "Bad Request" - http_error = httpx.HTTPStatusError( - "Error", request=MagicMock(), response=mock_response - ) + http_error = httpx.HTTPStatusError("Error", request=MagicMock(), response=mock_response) self.mock_client.post.side_effect = http_error optional_params = self.sample_optional_params.copy() @@ -362,15 +501,9 @@ class TestContextCachingEndpoints: assert "Bad Request" in str(exc_info.value.message) @pytest.mark.asyncio - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj" - ) + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages") + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj") async def test_async_check_and_create_cache_with_cached_content( self, mock_cache_obj, mock_separate, custom_llm_provider ): @@ -405,15 +538,9 @@ class TestContextCachingEndpoints: assert returned_cache == cached_content @pytest.mark.asyncio - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" - ) - async def test_async_check_and_create_cache_no_cached_messages( - self, mock_separate, custom_llm_provider - ): + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages") + async def test_async_check_and_create_cache_no_cached_messages(self, mock_separate, custom_llm_provider): """Test async_check_and_create_cache when no cached messages are found""" # Setup mock_separate.return_value = ([], self.sample_messages) @@ -444,15 +571,9 @@ class TestContextCachingEndpoints: assert returned_cache is None @pytest.mark.asyncio - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj" - ) + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages") + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj") @patch.object(ContextCachingEndpoints, "async_check_cache") async def test_async_check_and_create_cache_existing_cache_found( self, mock_async_check_cache, mock_cache_obj, mock_separate, custom_llm_provider @@ -501,23 +622,15 @@ class TestContextCachingEndpoints: ) @pytest.mark.asyncio - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj" - ) + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages") + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj") @patch( "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.transform_openai_messages_to_gemini_context_caching" ) @patch.object(ContextCachingEndpoints, "async_check_cache") @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching") - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.get_async_httpx_client" - ) + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.get_async_httpx_client") async def test_async_check_and_create_cache_create_new_cache( self, mock_get_client, @@ -581,20 +694,12 @@ class TestContextCachingEndpoints: assert call_args.kwargs["json"]["tools"] == self.sample_tools @pytest.mark.asyncio - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj" - ) + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages") + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj") @patch.object(ContextCachingEndpoints, "async_check_cache") @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching") - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.get_async_httpx_client" - ) + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.get_async_httpx_client") async def test_async_check_and_create_cache_timeout_error( self, mock_get_client, @@ -615,9 +720,7 @@ class TestContextCachingEndpoints: mock_get_token_url.return_value = ("token", "https://test-url.com") # Mock timeout error - self.mock_async_client.post = AsyncMock( - side_effect=httpx.TimeoutException("Timeout") - ) + self.mock_async_client.post = AsyncMock(side_effect=httpx.TimeoutException("Timeout")) optional_params = self.sample_optional_params.copy() test_project = "test_project" @@ -643,20 +746,14 @@ class TestContextCachingEndpoints: assert exc_info.value.status_code == 408 assert "Timeout error occurred" in str(exc_info.value.message) - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - def test_check_and_create_cache_tools_popped_from_optional_params( - self, custom_llm_provider - ): + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + def test_check_and_create_cache_tools_popped_from_optional_params(self, custom_llm_provider): """Test that tools are properly popped from optional_params when there are cached messages""" with patch( "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" ) as mock_separate: # Mock to return cached messages so tools get popped - cached_messages = [ - self.sample_messages[0] - ] # System message with cache_control + cached_messages = [self.sample_messages[0]] # System message with cache_control non_cached_messages = [self.sample_messages[1]] # User message mock_separate.return_value = (cached_messages, non_cached_messages) @@ -666,9 +763,7 @@ class TestContextCachingEndpoints: test_location = "us-central1" # Mock the check_cache to return existing cache so we don't make HTTP calls - with patch.object( - self.context_caching, "check_cache", return_value="existing_cache" - ): + with patch.object(self.context_caching, "check_cache", return_value="existing_cache"): # Execute result = self.context_caching.check_and_create_cache( messages=self.sample_messages, @@ -691,12 +786,8 @@ class TestContextCachingEndpoints: # But original tools should still be available for comparison assert original_tools == self.sample_tools - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - def test_check_and_create_cache_tools_not_popped_when_no_cached_messages( - self, custom_llm_provider - ): + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + def test_check_and_create_cache_tools_not_popped_when_no_cached_messages(self, custom_llm_provider): """Test that tools are NOT popped from optional_params when there are no cached messages""" with patch( "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" @@ -732,12 +823,8 @@ class TestContextCachingEndpoints: assert optional_params["tools"] == original_tools @pytest.mark.asyncio - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - async def test_async_check_and_create_cache_tools_not_popped_when_no_cached_messages( - self, custom_llm_provider - ): + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + async def test_async_check_and_create_cache_tools_not_popped_when_no_cached_messages(self, custom_llm_provider): """Test that tools are NOT popped from optional_params in async version when there are no cached messages""" with patch( "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" @@ -773,20 +860,14 @@ class TestContextCachingEndpoints: assert optional_params["tools"] == original_tools @pytest.mark.asyncio - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - async def test_async_check_and_create_cache_tools_popped_from_optional_params( - self, custom_llm_provider - ): + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + async def test_async_check_and_create_cache_tools_popped_from_optional_params(self, custom_llm_provider): """Test that tools are properly popped from optional_params in async version when there are cached messages""" with patch( "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" ) as mock_separate: # Mock to return cached messages so tools get popped - cached_messages = [ - self.sample_messages[0] - ] # System message with cache_control + cached_messages = [self.sample_messages[0]] # System message with cache_control non_cached_messages = [self.sample_messages[1]] # User message mock_separate.return_value = (cached_messages, non_cached_messages) @@ -796,9 +877,7 @@ class TestContextCachingEndpoints: test_location = "us-central1" # Mock the async_check_cache to return existing cache so we don't make HTTP calls - with patch.object( - self.context_caching, "async_check_cache", return_value="existing_cache" - ): + with patch.object(self.context_caching, "async_check_cache", return_value="existing_cache"): # Execute result = await self.context_caching.async_check_and_create_cache( messages=self.sample_messages, @@ -821,12 +900,8 @@ class TestContextCachingEndpoints: # But original tools should still be available for comparison assert original_tools == self.sample_tools - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - def test_check_and_create_cache_tool_choice_popped_from_optional_params( - self, custom_llm_provider - ): + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + def test_check_and_create_cache_tool_choice_popped_from_optional_params(self, custom_llm_provider): """tool_choice is popped from optional_params when cached messages exist.""" with patch( "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" @@ -838,9 +913,7 @@ class TestContextCachingEndpoints: optional_params = self.sample_optional_params.copy() optional_params["tool_choice"] = {"functionCallingConfig": {"mode": "ANY"}} - with patch.object( - self.context_caching, "check_cache", return_value="existing_cache" - ): + with patch.object(self.context_caching, "check_cache", return_value="existing_cache"): self.context_caching.check_and_create_cache( messages=self.sample_messages, optional_params=optional_params, @@ -858,12 +931,8 @@ class TestContextCachingEndpoints: assert "tool_choice" not in optional_params - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - def test_check_and_create_cache_tool_choice_not_popped_when_no_cached_messages( - self, custom_llm_provider - ): + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + def test_check_and_create_cache_tool_choice_not_popped_when_no_cached_messages(self, custom_llm_provider): """tool_choice is NOT popped when there are no cached messages (early return).""" with patch( "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" @@ -892,12 +961,8 @@ class TestContextCachingEndpoints: assert optional_params.get("tool_choice") == tool_choice @pytest.mark.asyncio - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - async def test_async_check_and_create_cache_tool_choice_popped_from_optional_params( - self, custom_llm_provider - ): + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + async def test_async_check_and_create_cache_tool_choice_popped_from_optional_params(self, custom_llm_provider): """Async equivalent of test_check_and_create_cache_tool_choice_popped_from_optional_params.""" with patch( "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" @@ -909,9 +974,7 @@ class TestContextCachingEndpoints: optional_params = self.sample_optional_params.copy() optional_params["tool_choice"] = {"functionCallingConfig": {"mode": "ANY"}} - with patch.object( - self.context_caching, "async_check_cache", return_value="existing_cache" - ): + with patch.object(self.context_caching, "async_check_cache", return_value="existing_cache"): await self.context_caching.async_check_and_create_cache( messages=self.sample_messages, optional_params=optional_params, @@ -930,9 +993,7 @@ class TestContextCachingEndpoints: assert "tool_choice" not in optional_params @pytest.mark.asyncio - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) async def test_async_check_and_create_cache_tool_choice_not_popped_when_no_cached_messages( self, custom_llm_provider ): @@ -963,15 +1024,9 @@ class TestContextCachingEndpoints: assert optional_params.get("tool_choice") == tool_choice - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj" - ) + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages") + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj") @patch( "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.transform_openai_messages_to_gemini_context_caching" ) @@ -1033,15 +1088,9 @@ class TestContextCachingEndpoints: ) @pytest.mark.asyncio - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj" - ) + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages") + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj") @patch( "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.transform_openai_messages_to_gemini_context_caching" ) @@ -1101,15 +1150,9 @@ class TestContextCachingEndpoints: model="gemini-1.5-pro", ) - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj" - ) + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages") + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj") @patch( "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.transform_openai_messages_to_gemini_context_caching" ) @@ -1161,15 +1204,9 @@ class TestContextCachingEndpoints: assert "tools" in call_args.kwargs["json"] assert "toolConfig" not in call_args.kwargs["json"] - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj" - ) + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages") + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj") @patch( "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.transform_openai_messages_to_gemini_context_caching" ) @@ -1227,15 +1264,9 @@ class TestContextCachingEndpoints: call_args = self.mock_client.post.call_args assert call_args.kwargs["json"]["toolConfig"] == function_pin - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj" - ) + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages") + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj") @patch( "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.transform_openai_messages_to_gemini_context_caching" ) @@ -1278,9 +1309,7 @@ class TestContextCachingEndpoints: } self.mock_client.post.return_value = mock_response - tool_choice = ToolConfig( - functionCallingConfig=FunctionCallingConfig(mode="ANY") - ) + tool_choice = ToolConfig(functionCallingConfig=FunctionCallingConfig(mode="ANY")) optional_params = self.sample_optional_params.copy() optional_params["tool_choice"] = tool_choice @@ -1301,9 +1330,7 @@ class TestContextCachingEndpoints: call_args = self.mock_client.post.call_args assert call_args.kwargs["json"]["toolConfig"] == tool_choice - assert call_args.kwargs["json"]["toolConfig"] == { - "functionCallingConfig": {"mode": "ANY"} - } + assert call_args.kwargs["json"]["toolConfig"] == {"functionCallingConfig": {"mode": "ANY"}} mock_cache_obj.get_cache_key.assert_called_once_with( messages=cached_messages, tools=self.sample_tools, @@ -1311,12 +1338,8 @@ class TestContextCachingEndpoints: model="gemini-1.5-pro", ) - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" - ) + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages") @patch.object(ContextCachingEndpoints, "check_cache") def test_check_and_create_cache_distinct_tool_choices_use_distinct_keys( self, @@ -1361,15 +1384,9 @@ class TestContextCachingEndpoints: second_cache_key = check_cache_calls[1].kwargs["cache_key"] assert first_cache_key != second_cache_key - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" - ) - def test_check_and_create_cache_skips_when_below_min_tokens( - self, mock_separate, custom_llm_provider - ): + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages") + def test_check_and_create_cache_skips_when_below_min_tokens(self, mock_separate, custom_llm_provider): """Test that context caching is skipped when cached content is below 1024 tokens. Gemini requires a minimum of 1024 tokens for context caching. If the cached @@ -1416,9 +1433,7 @@ class TestContextCachingEndpoints: # Restart the patcher so teardown_method can stop it cleanly self._token_check_patcher.start() - def test_check_and_create_cache_skips_between_default_and_gemini_2_5_pro_minimum( - self, local_model_cost_map - ): + def test_check_and_create_cache_skips_between_default_and_gemini_2_5_pro_minimum(self, local_model_cost_map): model = "gemini-2.5-pro" self._token_check_patcher.stop() @@ -1454,15 +1469,48 @@ class TestContextCachingEndpoints: self._token_check_patcher.start() @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" + ("api_base_model", "deployment_model"), + [("gemini-2.5-pro", "gemini-2.5-flash"), ("gemini-not-in-the-cost-map", "gemini-2.5-pro")], ) + @pytest.mark.parametrize("is_async", [False, True]) @pytest.mark.asyncio - async def test_async_check_and_create_cache_skips_when_below_min_tokens( - self, mock_separate, custom_llm_provider + async def test_check_and_create_cache_applies_the_stricter_model_minimum( + self, local_model_cost_map, is_async: bool, api_base_model: str, deployment_model: str ): + self._token_check_patcher.stop() + messages: Final = [ + {"role": "system", "content": " ".join(["word"] * 1500), "cache_control": {"type": "ephemeral"}}, + {"role": "user", "content": "Hello"}, + ] + call_args: Final = dict( + messages=messages, + optional_params={}, + api_key="unused", + api_base=f"https://gateway.example.com/v1/projects/base-project/locations/global/publishers/google/models/{api_base_model}", + model=deployment_model, + timeout=30.0, + logging_obj=self.mock_logging, + custom_llm_provider="vertex_ai", + vertex_project="config-project", + vertex_location="us-central1", + vertex_auth_header="vertex-token", + ) + + result: Final = ( + await self.context_caching.async_check_and_create_cache(client=self.mock_async_client, **call_args) + if is_async + else self.context_caching.check_and_create_cache(client=self.mock_client, **call_args) + ) + + self._token_check_patcher.start() + assert result == (messages, {}, None) + assert not self.mock_client.mock_calls + assert not self.mock_async_client.mock_calls + + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages") + @pytest.mark.asyncio + async def test_async_check_and_create_cache_skips_when_below_min_tokens(self, mock_separate, custom_llm_provider): """Test that async context caching is skipped when cached content is below 1024 tokens.""" # Stop the default mock so the real token count check runs self._token_check_patcher.stop() @@ -1505,13 +1553,9 @@ class TestContextCachingEndpoints: self._token_check_patcher.start() @pytest.mark.parametrize("is_async", [False, True]) - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai"] - ) + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai"]) @pytest.mark.asyncio - async def test_check_and_create_cache_considers_tools_for_min_tokens( - self, custom_llm_provider, is_async - ): + async def test_check_and_create_cache_considers_tools_for_min_tokens(self, custom_llm_provider, is_async): """Test that context caching accounts for tools when validating minimum token count. Fixes #42804: When messages alone are below the threshold, but tools push the total @@ -1536,11 +1580,15 @@ class TestContextCachingEndpoints: "type": "function", "function": { "name": f"synthetic_tool_{i}", - "description": "A very descriptive explanation of a synthetic tool designed to add tokens to the prompt cache prefix " * 8, + "description": "A very descriptive explanation of a synthetic tool designed to add tokens to the prompt cache prefix " + * 8, "parameters": { "type": "object", "properties": { - f"arg_{j}": {"type": "string", "description": "Argument description for caching verification " * 4} + f"arg_{j}": { + "type": "string", + "description": "Argument description for caching verification " * 4, + } for j in range(10) }, "required": [f"arg_{j}" for j in range(5)], @@ -1569,19 +1617,23 @@ class TestContextCachingEndpoints: self.mock_client.post.return_value = mock_response self.mock_async_client.post = AsyncMock(return_value=mock_response) - with patch.object( - self.context_caching, - "_get_token_and_url_context_caching", - return_value=("fake_token", "https://fake.url/cachedContents"), - ), patch.object( - self.context_caching, - "check_cache", - return_value=None, - ), patch.object( - self.context_caching, - "async_check_cache", - new_callable=AsyncMock, - return_value=None, + with ( + patch.object( + self.context_caching, + "_get_token_and_url_context_caching", + return_value=("fake_token", "https://fake.url/cachedContents"), + ), + patch.object( + self.context_caching, + "check_cache", + return_value=None, + ), + patch.object( + self.context_caching, + "async_check_cache", + new_callable=AsyncMock, + return_value=None, + ), ): if is_async: result = await self.context_caching.async_check_and_create_cache( @@ -1625,13 +1677,10 @@ class TestContextCachingEndpoints: post_mock.assert_called_once() call_kwargs = post_mock.call_args.kwargs assert call_kwargs["json"]["tools"] == large_tools - assert call_kwargs["json"]["contents"] == [ - {"role": "user", "parts": [{"text": "Short system instruction."}]} - ] + assert call_kwargs["json"]["contents"] == [{"role": "user", "parts": [{"text": "Short system instruction."}]}] self._token_check_patcher.start() - def _model_turn_final_messages(self, final_cached_role): tool_call = { "id": "call_abc123", @@ -1678,9 +1727,7 @@ class TestContextCachingEndpoints: ] @pytest.mark.parametrize("final_cached_role", ["assistant", "tool", "system"]) - def test_check_and_create_cache_skips_when_cached_block_ends_on_model_turn( - self, final_cached_role - ): + def test_check_and_create_cache_skips_when_cached_block_ends_on_model_turn(self, final_cached_role): """The cachedContents API rejects contents ending on an assistant or tool turn with HTTP 400 "Requests ending with a model turn are not supported", so the request must proceed uncached instead of failing. @@ -1713,9 +1760,7 @@ class TestContextCachingEndpoints: @pytest.mark.parametrize("final_cached_role", ["assistant", "tool", "system"]) @pytest.mark.asyncio - async def test_async_check_and_create_cache_skips_when_cached_block_ends_on_model_turn( - self, final_cached_role - ): + async def test_async_check_and_create_cache_skips_when_cached_block_ends_on_model_turn(self, final_cached_role): """Async variant: an unsupported terminal turn skips caching instead of failing.""" all_messages = self._model_turn_final_messages(final_cached_role) optional_params = self.sample_optional_params.copy() @@ -1773,14 +1818,8 @@ def test_cached_messages_end_on_supported_turn(): ) is True ) - assert ( - cached_messages_end_on_supported_turn([{"role": "tool", "tool_call_id": "x", "content": "y"}]) - is False - ) - assert ( - cached_messages_end_on_supported_turn([{"role": "function", "name": "f", "content": "y"}]) - is False - ) + assert cached_messages_end_on_supported_turn([{"role": "tool", "tool_call_id": "x", "content": "y"}]) is False + assert cached_messages_end_on_supported_turn([{"role": "function", "name": "f", "content": "y"}]) is False assert cached_messages_end_on_supported_turn([]) is False @@ -1795,13 +1834,9 @@ class TestCheckCachePagination: self.mock_client = MagicMock(spec=HTTPHandler) self.mock_async_client = MagicMock(spec=AsyncHTTPHandler) - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching") - def test_check_cache_pagination_finds_cache_on_second_page( - self, mock_get_token_url, custom_llm_provider - ): + def test_check_cache_pagination_finds_cache_on_second_page(self, mock_get_token_url, custom_llm_provider): """Test that check_cache correctly handles pagination and finds cache on second page""" # Setup mock_get_token_url.return_value = ("token", "https://test-url.com") @@ -1850,13 +1885,9 @@ class TestCheckCachePagination: second_call_url = self.mock_client.get.call_args_list[1].kwargs["url"] assert "pageToken=token_page_2" in second_call_url - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching") - def test_check_cache_pagination_stops_when_no_next_token( - self, mock_get_token_url, custom_llm_provider - ): + def test_check_cache_pagination_stops_when_no_next_token(self, mock_get_token_url, custom_llm_provider): """Test that check_cache stops pagination when no nextPageToken is present""" # Setup mock_get_token_url.return_value = ("token", "https://test-url.com") @@ -1891,13 +1922,9 @@ class TestCheckCachePagination: assert result is None assert self.mock_client.get.call_count == 1 - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching") - def test_check_cache_pagination_multiple_pages( - self, mock_get_token_url, custom_llm_provider - ): + def test_check_cache_pagination_multiple_pages(self, mock_get_token_url, custom_llm_provider): """Test that check_cache correctly iterates through multiple pages""" # Setup mock_get_token_url.return_value = ("token", "https://test-url.com") @@ -1942,9 +1969,7 @@ class TestCheckCachePagination: assert self.mock_client.get.call_count == 3 @pytest.mark.asyncio - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching") async def test_async_check_cache_pagination_finds_cache_on_second_page( self, mock_get_token_url, custom_llm_provider @@ -1974,9 +1999,7 @@ class TestCheckCachePagination: } # Setup mock async client to return different responses - self.mock_async_client.get = AsyncMock( - side_effect=[first_page_response, second_page_response] - ) + self.mock_async_client.get = AsyncMock(side_effect=[first_page_response, second_page_response]) # Execute result = await self.context_caching.async_check_cache( @@ -2000,13 +2023,9 @@ class TestCheckCachePagination: assert "pageToken=token_page_2" in second_call_url @pytest.mark.asyncio - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching") - async def test_async_check_cache_pagination_stops_when_no_next_token( - self, mock_get_token_url, custom_llm_provider - ): + async def test_async_check_cache_pagination_stops_when_no_next_token(self, mock_get_token_url, custom_llm_provider): """Test that async_check_cache stops pagination when no nextPageToken is present""" # Setup mock_get_token_url.return_value = ("token", "https://test-url.com") @@ -2041,29 +2060,21 @@ class TestCheckCachePagination: assert result is None assert self.mock_async_client.get.call_count == 1 - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching") - def test_check_cache_pagination_max_pages_limit( - self, mock_get_token_url, custom_llm_provider - ): + def test_check_cache_pagination_max_pages_limit(self, mock_get_token_url, custom_llm_provider): mock_get_token_url.return_value = ("token", "https://test-url.com") cache_key_to_find = "nonexistent_cache_key" def create_page_response(page_num): response = MagicMock() response.json.return_value = { - "cachedContents": [ - {"name": f"cache_{page_num}", "displayName": f"key_{page_num}"} - ], + "cachedContents": [{"name": f"cache_{page_num}", "displayName": f"key_{page_num}"}], "nextPageToken": f"token_page_{page_num + 1}", } return response - self.mock_client.get.side_effect = [ - create_page_response(i) for i in range(100) - ] + self.mock_client.get.side_effect = [create_page_response(i) for i in range(100)] result = self.context_caching.check_cache( cache_key=cache_key_to_find, @@ -2082,29 +2093,21 @@ class TestCheckCachePagination: assert self.mock_client.get.call_count == 100 @pytest.mark.asyncio - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching") - async def test_async_check_cache_pagination_max_pages_limit( - self, mock_get_token_url, custom_llm_provider - ): + async def test_async_check_cache_pagination_max_pages_limit(self, mock_get_token_url, custom_llm_provider): mock_get_token_url.return_value = ("token", "https://test-url.com") cache_key_to_find = "nonexistent_cache_key" def create_page_response(page_num): response = MagicMock() response.json.return_value = { - "cachedContents": [ - {"name": f"cache_{page_num}", "displayName": f"key_{page_num}"} - ], + "cachedContents": [{"name": f"cache_{page_num}", "displayName": f"key_{page_num}"}], "nextPageToken": f"token_page_{page_num + 1}", } return response - self.mock_async_client.get = AsyncMock( - side_effect=[create_page_response(i) for i in range(100)] - ) + self.mock_async_client.get = AsyncMock(side_effect=[create_page_response(i) for i in range(100)]) result = await self.context_caching.async_check_cache( cache_key=cache_key_to_find, @@ -2123,8 +2126,6 @@ class TestCheckCachePagination: assert self.mock_async_client.get.call_count == 100 - - class TestVertexAIGlobalLocation: """Test global location handling in context caching.""" @@ -2150,9 +2151,7 @@ class TestVertexAIGlobalLocation: # Assert correct URL format for global expected_url = "https://aiplatform.googleapis.com/v1/projects/test-project/locations/global/cachedContents" assert url == expected_url, f"Expected {expected_url}, got {url}" - assert ( - "global-aiplatform" not in url - ), "URL should not contain 'global-aiplatform' prefix" + assert "global-aiplatform" not in url, "URL should not contain 'global-aiplatform' prefix" def test_regional_location_url_construction_v1(self): """Test that regional location uses correct URL (with location prefix) for v1 API.""" @@ -2195,32 +2194,117 @@ class TestVertexAIGlobalLocation: ) # Assert correct URL format for global with beta API - expected_url = "https://aiplatform.googleapis.com/v1beta1/projects/test-project/locations/global/cachedContents" + expected_url = ( + "https://aiplatform.googleapis.com/v1beta1/projects/test-project/locations/global/cachedContents" + ) assert url == expected_url, f"Expected {expected_url}, got {url}" - assert ( - "global-aiplatform" not in url - ), "URL should not contain 'global-aiplatform' prefix" + assert "global-aiplatform" not in url, "URL should not contain 'global-aiplatform' prefix" - def test_gemini_context_caching_with_custom_api_base_passes_model(self): - """Gemini context caching with custom api_base must pass model to _check_custom_proxy. - - Regression test for https://github.com/BerriAI/litellm/issues/23846 - Previously model was hardcoded to None, causing ValueError when api_base was set. - """ + @pytest.mark.parametrize( + ("api_base", "expected_url"), + [ + ("https://my-proxy.example.com", "https://my-proxy.example.com/cachedContents"), + ("https://my-proxy.example.com/v1beta/", "https://my-proxy.example.com/v1beta/cachedContents"), + ( + "https://gateway.example.com/publishers/relay/v1beta", + "https://gateway.example.com/publishers/relay/v1beta/cachedContents", + ), + ], + ) + def test_gemini_context_caching_with_custom_api_base_uses_collection_endpoint( + self, api_base: str, expected_url: str + ): caching = ContextCachingEndpoints() auth_header, url = caching._get_token_and_url_context_caching( gemini_api_key="test-key", custom_llm_provider="gemini", - api_base="https://my-proxy.example.com", + api_base=api_base, vertex_project=None, vertex_location=None, vertex_auth_header=None, model="gemini-1.5-pro", ) - assert "models/gemini-1.5-pro" in url - assert url.startswith("https://my-proxy.example.com/") + assert auth_header == {"x-goog-api-key": "test-key"} + assert url == expected_url + + @pytest.mark.parametrize( + ("custom_llm_provider", "api_base", "expected_url"), + [ + ( + "vertex_ai", + "https://gateway.example.com", + "https://gateway.example.com/v1/projects/my-project/locations/us-central1/cachedContents", + ), + ( + "vertex_ai", + "https://gateway.example.com/v1", + "https://gateway.example.com/v1/projects/my-project/locations/us-central1/cachedContents", + ), + ( + "vertex_ai", + "https://gateway.example.com/v1/projects/my-project/locations/us-central1" + "/publishers/google/models/gemini-2.5-pro", + "https://gateway.example.com/v1/projects/my-project/locations/us-central1/cachedContents", + ), + ( + "vertex_ai", + "https://gateway.example.com/publishers/relay/v1/projects/my-project/locations/us-central1" + "/publishers/google/models/gemini-2.5-pro/", + "https://gateway.example.com/publishers/relay/v1/projects/my-project/locations/us-central1" + "/cachedContents", + ), + ( + "vertex_ai_beta", + "https://gateway.example.com/v1beta1/projects/my-project/locations/us-central1" + "/publishers/google/models/gemini-2.5-pro", + "https://gateway.example.com/v1beta1/projects/my-project/locations/us-central1/cachedContents", + ), + ], + ) + def test_vertex_context_caching_with_custom_api_base_uses_collection_endpoint( + self, custom_llm_provider: Literal["vertex_ai", "vertex_ai_beta"], api_base: str, expected_url: str + ): + caching = ContextCachingEndpoints() + + auth_header, url = caching._get_token_and_url_context_caching( + gemini_api_key=None, + custom_llm_provider=custom_llm_provider, + api_base=api_base, + vertex_project="my-project", + vertex_location="us-central1", + vertex_auth_header="vertex-token", + model="gemini-2.5-pro", + ) + + assert auth_header == "vertex-token" + assert url == expected_url + + @pytest.mark.parametrize( + "api_base", + [ + "https://gateway.example.com/v1/projects/tuned-project/locations/europe-west4/endpoints/123", + "https://gateway.example.com/v1/projects/my-project/locations/us-central1", + "https://gateway.example.com/vertex-proxy", + ], + ) + def test_vertex_context_caching_without_a_publisher_model_in_api_base_keeps_the_api_base_resource( + self, api_base: str + ): + caching = ContextCachingEndpoints() + + _, url = caching._get_token_and_url_context_caching( + gemini_api_key=None, + custom_llm_provider="vertex_ai", + api_base=api_base, + vertex_project="my-project", + vertex_location="us-central1", + vertex_auth_header="vertex-token", + model="gemini-2.5-pro", + ) + + assert url == f"{api_base}:cachedContents" def test_gemini_context_caching_without_api_base_ignores_model(self): """Without custom api_base, model param is not needed (default URL is used).""" @@ -2305,9 +2389,7 @@ def _cacheable_messages() -> list: return [ { "role": "system", - "content": [ - {"type": "text", "text": "cacheable prefix " * 1200, "cache_control": {"type": "ephemeral"}} - ], + "content": [{"type": "text", "text": "cacheable prefix " * 1200, "cache_control": {"type": "ephemeral"}}], }, {"role": "user", "content": "question"}, ] @@ -2404,7 +2486,9 @@ def test_check_and_create_cache_clamps_negative_creation_tokens() -> None: with ( patch.object(ContextCachingEndpoints, "check_cache", return_value=None), - patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching", return_value=("token", "https://u")), + patch.object( + ContextCachingEndpoints, "_get_token_and_url_context_caching", return_value=("token", "https://u") + ), ): _, _, cache_name = ContextCachingEndpoints().check_and_create_cache( messages=_cacheable_messages(),