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(),