From df8e5397afbc8bc62df054e263dd21d995845d10 Mon Sep 17 00:00:00 2001 From: Dominastorm Date: Wed, 14 Oct 2026 19:47:00 +0530 Subject: [PATCH 1/2] fix(vertex): share context cache IDs through Redis --- .../context_caching/redis_cache_ids.py | 586 ++++++++++++++++++ .../vertex_ai_context_caching.py | 48 ++ .../llms/vertex_ai/gemini/transformation.py | 8 + .../vertex_and_google_ai_studio_gemini.py | 52 +- .../context_caching/test_redis_cache_ids.py | 184 ++++++ 5 files changed, 870 insertions(+), 8 deletions(-) create mode 100644 litellm/llms/vertex_ai/context_caching/redis_cache_ids.py create mode 100644 tests/unit/llms/vertex_ai/context_caching/test_redis_cache_ids.py diff --git a/litellm/llms/vertex_ai/context_caching/redis_cache_ids.py b/litellm/llms/vertex_ai/context_caching/redis_cache_ids.py new file mode 100644 index 00000000000..62272495ba4 --- /dev/null +++ b/litellm/llms/vertex_ai/context_caching/redis_cache_ids.py @@ -0,0 +1,586 @@ +"""Opt-in Vertex cache-name registry, independent of LiteLLM response caching. + +Redis stores hashes, resource names and deadlines, never prompts or credentials. +No Google LIST/GET is used. All coordination operations have bounded timeouts. +""" + +from __future__ import annotations + +import asyncio +import copy +import hashlib +import json +import math +import os +import re +import time +import uuid +from collections.abc import Awaitable, Callable, Mapping +from dataclasses import dataclass +from typing import TYPE_CHECKING, Final, Literal, Protocol + +from pydantic import JsonValue, TypeAdapter +from typing_extensions import ReadOnly, TypedDict + +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler +from litellm.types.llms.openai import AllMessageValues +from litellm.types.llms.vertex_ai import CachedContentRequestBody, RequestBody + +if TYPE_CHECKING: + from redis import Redis + + from .vertex_ai_context_caching import ContextCachingEndpoints +from datetime import datetime +from functools import lru_cache + +import httpx + + +def redis_error_type() -> type[Exception]: + try: + from redis.exceptions import RedisError + + return RedisError + except ImportError: + return ImportError + + +RedisError: Final = redis_error_type() +JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) + + +Params = dict[str, JsonValue] +Provider = Literal["vertex_ai", "vertex_ai_beta", "gemini"] +CacheResult = tuple[list[AllMessageValues], Params, str | None] + + +class CacheArguments(TypedDict): + messages: ReadOnly[list[AllMessageValues]] + optional_params: ReadOnly[Params] + api_key: ReadOnly[str] + api_base: ReadOnly[str | None] + model: ReadOnly[str] + client: ReadOnly[HTTPHandler | AsyncHTTPHandler | None] + timeout: ReadOnly[float | httpx.Timeout | None] + logging_obj: ReadOnly[Logging] + custom_llm_provider: ReadOnly[Provider] + vertex_project: ReadOnly[str | None] + vertex_location: ReadOnly[str | None] + vertex_auth_header: ReadOnly[str | None] + extra_headers: ReadOnly[dict[str, str] | None] + cache_id_auth_scope: ReadOnly[str | None] + + +class RequestTransform(Protocol): + def __call__( + self, + *, + messages: list[AllMessageValues], + model: str, + optional_params: Params, + custom_llm_provider: Provider, + litellm_params: Params, + cached_content: str | None, + ) -> RequestBody: ... + + +@dataclass(frozen=True, slots=True) +class PreparedCache: + registry: Registry + url: str + headers: dict[str, str] + body: CachedContentRequestBody + messages: list[AllMessageValues] + params: Params + + +ENV = "LITELLM_VERTEX_CACHE_ID_REDIS_URL" +CACHE_NAME = re.compile(r"projects/[^/]+/locations/[^/]+/cachedContents/[^/]+") +# A bounded wait avoids turning cache contention into generation latency. +WAIT_SECONDS = 2.0 +LEASE_MS = 60000 +SAFETY_SECONDS = 60.0 +RELEASE = "if redis.call('get',KEYS[1]) == ARGV[1] then return redis.call('del',KEYS[1]) end return 0" +PUBLISH = """if redis.call('get',KEYS[1]) == ARGV[1] then +redis.call('psetex',KEYS[2],ARGV[3],ARGV[2]); redis.call('del',KEYS[1]); return 1 end return 0""" +INVALIDATE = """local v=redis.call('get',KEYS[1]); if v then +local ok,r=pcall(cjson.decode,v); if ok and r.name == ARGV[1] then return redis.call('del',KEYS[1]) end end return 0""" + + +def valid_name(name: object) -> bool: + return isinstance(name, str) and CACHE_NAME.fullmatch(name) is not None + + +def enabled(provider: str) -> bool: + return provider == "vertex_ai" and bool(os.environ.get(ENV)) + + +@lru_cache(maxsize=4) +def connection(url: str) -> Redis: + from redis import Redis + + return Redis.from_url( + url, + socket_connect_timeout=0.2, + socket_timeout=0.2, + retry_on_timeout=False, + decode_responses=True, + ) + + +def digest(value: object) -> str: + return hashlib.sha256( + json.dumps( + value, + sort_keys=True, + separators=(",", ":"), + ensure_ascii=False, + allow_nan=False, + ).encode() + ).hexdigest() + + +def credential_scope(credentials: object) -> str | None: + """Fingerprint explicit credentials, independent of short-lived access tokens. + + Unknown/ADC credentials use the token-scoped fallback in prepare(). Raw + credentials never enter Redis. File contents are read only for a configured + credential path already accepted by the Vertex authentication code. + """ + if not enabled("vertex_ai"): + return None + try: + if isinstance(credentials, str): + if os.path.isfile(credentials): + with open(credentials) as stream: + credentials = json.load(stream) + else: + credentials = json.loads(credentials) + if isinstance(credentials, dict) and credentials: + return digest(credentials) + except (OSError, ValueError, TypeError): + pass + return None + + +def transient_creation_error(error: Exception) -> bool: + import litellm + + if isinstance(error, (httpx.TransportError, litellm.Timeout)): + return True + return isinstance(error, httpx.HTTPStatusError) and error.response.status_code in ( + 429, + 500, + 502, + 503, + 504, + ) + + +class Registry: + def __init__( + self, + client: Redis, + scope: object, + body: object, + *, + safety: float = SAFETY_SECONDS, + wait: float = WAIT_SECONDS, + lease_ms: int = LEASE_MS, + clock: Callable[[], float] = time.time, + monotonic: Callable[[], float] = time.monotonic, + sleep: Callable[[float], None] = time.sleep, + async_sleep: Callable[[float], Awaitable[None]] = asyncio.sleep, + ) -> None: + self.client = client + self.key = "litellm:vertex-cache-id:v1:{" + digest([scope, body]) + "}" + self.lock = self.key + ":lock" + self.cooldown = "litellm:vertex-cache-id:cooldown:" + digest(scope) + self.safety, self.wait, self.lease_ms = safety, wait, lease_ms + self.clock, self.monotonic = clock, monotonic + self.sleep, self.async_sleep = sleep, async_sleep + + def read(self) -> str | None: + raw = self.client.get(self.key) + if not raw: + return None + try: + record = JSON_OBJECT.validate_json(raw) + if not isinstance(record.get("name"), str) or not valid_name(record["name"]): + return None + if not math.isfinite(float(record["expires"])) or float(record["expires"]) - self.safety <= self.clock(): + return None + return record["name"] + except (ValueError, TypeError, KeyError): + return None + + def publish(self, token: str, result: Mapping[str, object]) -> str | None: + try: + name = result["name"] + if not isinstance(result["expireTime"], str): + return None + expires = datetime.fromisoformat(result["expireTime"].replace("Z", "+00:00")) + if expires.tzinfo is None or not isinstance(name, str) or not valid_name(name): + return None + deadline = expires.timestamp() + ttl = int((deadline - self.safety - self.clock()) * 1000) + if ttl <= 0: + return None + record = json.dumps({"name": name, "expires": deadline}) + if self.client.eval(PUBLISH, 2, self.lock, self.key, token, record, ttl): + return name + except (ValueError, TypeError, KeyError): + pass + return None + + def release(self, token: str) -> None: + try: + self.client.eval(RELEASE, 1, self.lock, token) + except RedisError: + pass + + def invalidate(self, name: str) -> None: + try: + self.client.eval(INVALIDATE, 1, self.key, str(name)) + except RedisError: + pass + + def acquire(self, token: str) -> bool: + return bool(self.client.set(self.lock, token, nx=True, px=self.lease_ms)) + + def cool_down(self) -> None: + # Shared per auth/project/endpoint scope. Do not retry cache creation on 429. + try: + self.client.set(self.cooldown, "1", ex=30) + except RedisError: + pass + + def lookup(self) -> tuple[bool, str | None]: + """Return (finished, name); a cooldown finishes without a cache ID.""" + name = self.read() + return bool(name or self.client.exists(self.cooldown)), name + + def resolve(self, create: Callable[[], Mapping[str, object]]) -> str | None: + token = uuid.uuid4().hex + end = self.monotonic() + self.wait + try: + while True: + finished, name = self.lookup() + if finished: + return name + if self.acquire(token): + try: + # Prior owner may have published before acquisition. + finished, name = self.lookup() + if finished: + return name + return self.publish(token, create()) + except Exception as error: + if not transient_creation_error(error): + raise + self.cool_down() + return None + finally: + self.release(token) + if self.monotonic() >= end: + return None + self.sleep(0.02) + except RedisError: + return None + + async def async_resolve(self, create: Callable[[], Awaitable[Mapping[str, object]]]) -> str | None: + token = uuid.uuid4().hex + end = self.monotonic() + self.wait + try: + while True: + finished, name = await asyncio.to_thread(self.lookup) + if finished: + return name + if await asyncio.to_thread(self.acquire, token): + try: + finished, name = await asyncio.to_thread(self.lookup) + if finished: + return name + result = await create() + return await asyncio.to_thread(self.publish, token, result) + except Exception as error: + if not transient_creation_error(error): + raise + await asyncio.to_thread(self.cool_down) + return None + finally: + await asyncio.shield(asyncio.to_thread(self.release, token)) + if self.monotonic() >= end: + return None + await self.async_sleep(0.02) + except RedisError: + return None + + +class CacheReference(str): + """String on the wire; local metadata supports one stale-cache recovery.""" + + registry: Registry + messages: list[AllMessageValues] + params: Params + rebuild: Callable[[], RequestBody] + + def __new__(cls, name: str, registry: Registry, messages: list[AllMessageValues], params: Params) -> CacheReference: + obj = super().__new__(cls, name) + obj.registry = registry + obj.messages = copy.deepcopy(messages) + obj.params = copy.deepcopy(params) + return obj + + def __deepcopy__(self, memo: dict[int, object]) -> CacheReference: + return CacheReference(str(self), self.registry, self.messages, self.params) + + def restore(self, body: RequestBody) -> None: + replacement = self.rebuild() + for key in tuple(body): + del body[key] + body.update(replacement) + + +def prepare(endpoint: ContextCachingEndpoints, kw: CacheArguments) -> PreparedCache | None: + from .transformation import ( + cached_messages_end_on_supported_turn, + separate_cached_messages, + transform_openai_messages_to_gemini_context_caching, + ) + from .vertex_ai_context_caching import is_prompt_caching_valid_prompt + + original_messages = kw["messages"] + original_params = kw["optional_params"] + cached, uncached = separate_cached_messages(copy.deepcopy(original_messages)) + if ( + not cached + or not cached_messages_end_on_supported_turn(cached) + or not is_prompt_caching_valid_prompt( + model=kw["model"], + messages=cached, + tools=original_params.get("tools"), + custom_llm_provider=kw["custom_llm_provider"], + ) + ): + return None + params = copy.deepcopy(original_params) + tools, tool_choice = params.pop("tools", None), params.pop("tool_choice", None) + token, url = endpoint._get_token_and_url_context_caching( + gemini_api_key=kw["api_key"], + custom_llm_provider=kw["custom_llm_provider"], + api_base=kw["api_base"], + vertex_project=kw["vertex_project"], + vertex_location=kw["vertex_location"], + vertex_auth_header=kw["vertex_auth_header"], + model=kw["model"], + ) + headers = {"Content-Type": "application/json"} + if token is not None: + headers["Authorization"] = token if token.startswith("Bearer ") else "Bearer " + token + headers.update(kw.get("extra_headers") or {}) + scoped_headers = {key.lower(): value for key, value in headers.items()} + extra = {key.lower(): value for key, value in (kw.get("extra_headers") or {}).items()} + stable_auth = kw.get("cache_id_auth_scope") + if stable_auth and "authorization" not in extra: + scoped_headers.pop("authorization", None) + else: + stable_auth = None + scope = [ + kw["custom_llm_provider"], + kw["vertex_project"], + kw["vertex_location"], + url, + stable_auth, + digest(scoped_headers), + ] + body = transform_openai_messages_to_gemini_context_caching( + model=kw["model"], + messages=cached, + cache_key="", + custom_llm_provider=kw["custom_llm_provider"], + vertex_project=kw["vertex_project"], + vertex_location=kw["vertex_location"], + ) + if "ttl" in body and float(body["ttl"][:-1]) <= SAFETY_SECONDS: + return None + body["tools"] = tools + if tool_choice is not None: + body["toolConfig"] = tool_choice + registry = Registry(connection(os.environ[ENV]), scope, body) + body["displayName"] = registry.key.split("{")[1].rstrip("}") + return PreparedCache(registry, url, headers, body, uncached, params) + + +def check_and_create(endpoint: ContextCachingEndpoints, kw: CacheArguments) -> CacheResult: + import litellm + from litellm.llms.custom_httpx.http_handler import HTTPHandler + + original = (kw["messages"], kw["optional_params"], None) + try: + prepared = prepare(endpoint, kw) + if prepared is None: + return original + registry, url, headers, body, messages, params = ( + prepared.registry, + prepared.url, + prepared.headers, + prepared.body, + prepared.messages, + prepared.params, + ) + + def create() -> Mapping[str, object]: + owned = not isinstance(kw["client"], HTTPHandler) + supplied = kw["client"] + client = supplied if isinstance(supplied, HTTPHandler) else HTTPHandler(timeout=20) + try: + response = client.post(url=url, headers=headers, data=json.dumps(body), timeout=20) + response.raise_for_status() + return record_creation(kw["logging_obj"], response) + finally: + if owned: + client.close() + + name = registry.resolve(create) + if name: + return ( + messages, + params, + CacheReference(name, registry, kw["messages"], kw["optional_params"]), + ) + except (RedisError, httpx.TransportError, litellm.Timeout): + pass + return original + + +async def async_check_and_create(endpoint: ContextCachingEndpoints, kw: CacheArguments) -> CacheResult: + import litellm + from litellm.llms.custom_httpx.http_handler import ( + AsyncHTTPHandler, + get_async_httpx_client, + ) + + original = (kw["messages"], kw["optional_params"], None) + try: + prepared = prepare(endpoint, kw) + if prepared is None: + return original + registry, url, headers, body, messages, params = ( + prepared.registry, + prepared.url, + prepared.headers, + prepared.body, + prepared.messages, + prepared.params, + ) + + async def create() -> Mapping[str, object]: + # This helper returns a shared client; do not close it here. + client = ( + kw["client"] + if isinstance(kw["client"], AsyncHTTPHandler) + else get_async_httpx_client(params={"timeout": 20}, llm_provider=litellm.LlmProviders.VERTEX_AI) + ) + response = await client.post(url=url, headers=headers, data=json.dumps(body), timeout=20) + response.raise_for_status() + return record_creation(kw["logging_obj"], response) + + name = await registry.async_resolve(create) + if name: + return ( + messages, + params, + CacheReference(name, registry, kw["messages"], kw["optional_params"]), + ) + except (RedisError, httpx.TransportError, litellm.Timeout): + pass + return original + + +def attach_recovery( + reference: str | None, transform: RequestTransform, model: str, provider: Provider, litellm_params: Params +) -> None: + if isinstance(reference, CacheReference): + saved_params = copy.deepcopy(litellm_params) + messages, params = reference.messages, reference.params + reference.rebuild = lambda: transform( + messages=copy.deepcopy(messages), + optional_params=copy.deepcopy(params), + model=model, + custom_llm_provider=provider, + litellm_params=copy.deepcopy(saved_params), + cached_content=None, + ) + + +def stale_reference(body: RequestBody | None, response: httpx.Response) -> CacheReference | None: + reference = body.get("cachedContent") if isinstance(body, dict) else None + if ( + not isinstance(reference, CacheReference) + or not hasattr(reference, "rebuild") + or response.status_code not in (400, 404) + ): + return None + try: + message = response.json().get("error", {}).get("message", "").lower() + except (ValueError, AttributeError): + return None + cache_error = any(word in message for word in ("cached content", "cachedcontent", "cached_content")) + missing = any(word in message for word in ("expired", "not found", "does not exist")) + return reference if cache_error and missing else None + + +def post_with_recovery( + post: Callable[[RequestBody | None], httpx.Response], body: RequestBody | None +) -> httpx.Response: + """Retry exactly once, uncached, only for a recognized stale managed cache.""" + if not isinstance((body or {}).get("cachedContent"), CacheReference): + return post(body) + try: + response = post(body) + response.raise_for_status() + return response + except httpx.HTTPStatusError as err: + err.response.read() + reference = stale_reference(body, err.response) + if reference is None: + raise + err.response.close() + reference.registry.invalidate(reference) + reference.restore(body) + response = post(body) + response.raise_for_status() + return response + + +async def async_post_with_recovery( + post: Callable[[RequestBody | None], Awaitable[httpx.Response]], body: RequestBody | None +) -> httpx.Response: + if not isinstance((body or {}).get("cachedContent"), CacheReference): + return await post(body) + try: + response = await post(body) + response.raise_for_status() + return response + except httpx.HTTPStatusError as err: + await err.response.aread() + reference = stale_reference(body, err.response) + if reference is None: + raise + await err.response.aclose() + await asyncio.to_thread(reference.registry.invalidate, reference) + await asyncio.to_thread(reference.restore, body) + response = await post(body) + response.raise_for_status() + return response + + +def record_creation(logging_obj: Logging, response: httpx.Response) -> Mapping[str, object]: + from .vertex_ai_context_caching import _CreatedCachedContentRoot, _record_cache_creation_tokens + + raw = JSON_OBJECT.validate_json(response.content) + created = _CreatedCachedContentRoot.model_validate(raw).root + _record_cache_creation_tokens(logging_obj, created) + return raw diff --git a/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py b/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py index 06b8f3a3fbb..3ae89dd579e 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 @@ -315,6 +315,7 @@ class ContextCachingEndpoints(VertexBase): vertex_auth_header: str | None, extra_headers: dict | None = None, cached_content: str | None = None, + cache_id_auth_scope: str | None = None, ) -> tuple[list[AllMessageValues], dict, str | None]: """ Receives @@ -329,6 +330,29 @@ class ContextCachingEndpoints(VertexBase): if cached_content is not None: return messages, optional_params, cached_content + from . import redis_cache_ids + + if redis_cache_ids.enabled(custom_llm_provider): + return redis_cache_ids.check_and_create( + self, + redis_cache_ids.CacheArguments( + messages=messages, + optional_params=optional_params, + api_key=api_key, + api_base=api_base, + model=model, + client=client, + timeout=timeout, + logging_obj=logging_obj, + custom_llm_provider=custom_llm_provider, + vertex_project=vertex_project, + vertex_location=vertex_location, + vertex_auth_header=vertex_auth_header, + extra_headers=extra_headers, + cache_id_auth_scope=cache_id_auth_scope, + ), + ) + cached_messages, non_cached_messages = separate_cached_messages(messages=messages) if len(cached_messages) == 0: @@ -473,6 +497,7 @@ class ContextCachingEndpoints(VertexBase): vertex_auth_header: str | None, extra_headers: dict | None = None, cached_content: str | None = None, + cache_id_auth_scope: str | None = None, ) -> tuple[list[AllMessageValues], dict, str | None]: """ Receives @@ -487,6 +512,29 @@ class ContextCachingEndpoints(VertexBase): if cached_content is not None: return messages, optional_params, cached_content + from . import redis_cache_ids + + if redis_cache_ids.enabled(custom_llm_provider): + return await redis_cache_ids.async_check_and_create( + self, + redis_cache_ids.CacheArguments( + messages=messages, + optional_params=optional_params, + api_key=api_key, + api_base=api_base, + model=model, + client=client, + timeout=timeout, + logging_obj=logging_obj, + custom_llm_provider=custom_llm_provider, + vertex_project=vertex_project, + vertex_location=vertex_location, + vertex_auth_header=vertex_auth_header, + extra_headers=extra_headers, + cache_id_auth_scope=cache_id_auth_scope, + ), + ) + cached_messages, non_cached_messages = separate_cached_messages(messages=messages) if len(cached_messages) == 0: diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py index 15e5e440b91..6e65abf1683 100644 --- a/litellm/llms/vertex_ai/gemini/transformation.py +++ b/litellm/llms/vertex_ai/gemini/transformation.py @@ -1138,6 +1138,10 @@ def transform_request_body( """ Common transformation logic across sync + async Gemini /generateContent calls. """ + from ..context_caching.redis_cache_ids import attach_recovery + + attach_recovery(cached_content, transform_request_body, model, custom_llm_provider, litellm_params) + # Separate system prompt from rest of message supports_system_message: Final = get_supports_system_message(model=model, custom_llm_provider=custom_llm_provider) system_instructions, messages = transform_system_message( @@ -1258,6 +1262,7 @@ def sync_transform_request_body( vertex_project: str | None, vertex_location: str | None, vertex_auth_header: str | None, + cache_id_auth_scope: str | None = None, ) -> RequestBody: from ..context_caching.vertex_ai_context_caching import ContextCachingEndpoints @@ -1282,6 +1287,7 @@ def sync_transform_request_body( vertex_project=vertex_project, vertex_location=vertex_location, vertex_auth_header=vertex_auth_header, + cache_id_auth_scope=cache_id_auth_scope, ) return transform_request_body( @@ -1326,6 +1332,7 @@ async def async_transform_request_body( vertex_project: str | None, vertex_location: str | None, vertex_auth_header: str | None, + cache_id_auth_scope: str | None = None, ) -> RequestBody: from ..context_caching.vertex_ai_context_caching import ContextCachingEndpoints @@ -1350,6 +1357,7 @@ async def async_transform_request_body( vertex_project=vertex_project, vertex_location=vertex_location, vertex_auth_header=vertex_auth_header, + cache_id_auth_scope=cache_id_auth_scope, ) inlined_messages: Final = await async_inline_remote_media( diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index 517ac9e98b1..c7f84303e70 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -65,6 +65,7 @@ from litellm.types.llms.vertex_ai import ( GenerateContentResponseBody, HttpxPartType, LogprobsResult, + RequestBody, ToolConfig, Tools, UsageMetadata, @@ -94,6 +95,11 @@ from ..common_utils import ( build_vertex_schema, supports_response_json_schema, ) +from ..context_caching.redis_cache_ids import ( + async_post_with_recovery, + credential_scope, + post_with_recovery, +) from ..vertex_llm_base import VertexBase from .grounding_requests import calculate_grounding_requests from .transformation import ( @@ -2748,6 +2754,7 @@ async def make_call( model: str, messages: list, logging_obj, + cache_body: RequestBody | None = None, ): if gemini_client is not None: client = gemini_client @@ -2757,7 +2764,16 @@ async def make_call( ) try: - response: Final = await client.post(api_base, headers=headers, data=data, stream=True, logging_obj=logging_obj) + response: Final = await async_post_with_recovery( + lambda body: client.post( + api_base, + headers=headers, + data=json.dumps(body) if body is not None else data, + stream=True, + logging_obj=logging_obj, + ), + cache_body, + ) response.raise_for_status() except httpx.HTTPStatusError as e: exception_string: Final = str(await e.response.aread()) @@ -2800,13 +2816,28 @@ def make_sync_call( model: str, messages: list, logging_obj, + cache_body: RequestBody | None = None, ): if gemini_client is not None: client = gemini_client if client is None: client = HTTPHandler() # Create a new client if none provided - response: Final = client.post(api_base, headers=headers, data=data, stream=True, logging_obj=logging_obj) + try: + response: Final = post_with_recovery( + lambda body: client.post( + api_base, + headers=headers, + data=json.dumps(body) if body is not None else data, + stream=True, + logging_obj=logging_obj, + ), + cache_body, + ) + except httpx.HTTPStatusError as error: + raise VertexAIError( + status_code=error.response.status_code, message=str(error.response.read()), headers=error.response.headers + ) from error if response.status_code != 200 and response.status_code != 201: raise VertexAIError( @@ -2899,6 +2930,7 @@ class VertexLLM(VertexBase): vertex_project=vertex_project, vertex_location=vertex_location, vertex_auth_header=auth_header, + cache_id_auth_scope=credential_scope(vertex_credentials), ) ## LOGGING @@ -2921,6 +2953,7 @@ class VertexLLM(VertexBase): api_base=api_base, headers=headers, data=request_body_str, + cache_body=request_body, model=model, messages=messages, logging_obj=logging_obj, @@ -2992,6 +3025,7 @@ class VertexLLM(VertexBase): vertex_project=vertex_project, vertex_location=vertex_location, vertex_auth_header=auth_header, + cache_id_auth_scope=credential_scope(vertex_credentials), ) _async_client_params: Final = {} @@ -3013,11 +3047,9 @@ class VertexLLM(VertexBase): ) try: - response: Final = await client.post( - api_base, - headers=headers, - json=cast(dict, request_body), - logging_obj=logging_obj, + response: Final = await async_post_with_recovery( + lambda body: client.post(api_base, headers=headers, json=body, logging_obj=logging_obj), + request_body, ) response.raise_for_status() except httpx.HTTPStatusError as err: @@ -3173,6 +3205,7 @@ class VertexLLM(VertexBase): vertex_project=vertex_project, vertex_location=vertex_location, vertex_auth_header=auth_header, + cache_id_auth_scope=credential_scope(vertex_credentials), ) ## LOGGING @@ -3196,6 +3229,7 @@ class VertexLLM(VertexBase): gemini_client=(client if client is not None and isinstance(client, HTTPHandler) else None), api_base=url, data=request_data_str, + cache_body=data, model=model, messages=messages, logging_obj=logging_obj, @@ -3220,7 +3254,9 @@ class VertexLLM(VertexBase): client = client try: - response: Final = client.post(url=url, headers=headers, json=data, logging_obj=logging_obj) + response: Final = post_with_recovery( + lambda body: client.post(url=url, headers=headers, json=body, logging_obj=logging_obj), data + ) response.raise_for_status() except httpx.HTTPStatusError as err: error_code: Final = err.response.status_code diff --git a/tests/unit/llms/vertex_ai/context_caching/test_redis_cache_ids.py b/tests/unit/llms/vertex_ai/context_caching/test_redis_cache_ids.py new file mode 100644 index 00000000000..b46e9f6910d --- /dev/null +++ b/tests/unit/llms/vertex_ai/context_caching/test_redis_cache_ids.py @@ -0,0 +1,184 @@ +from datetime import datetime, timezone +from unittest.mock import Mock + +import httpx +import pytest +from redis import RedisError + +from litellm.llms.vertex_ai.context_caching.redis_cache_ids import Registry + +NAME = "projects/project/locations/global/cachedContents/cache" +NOW = 1_900_000_000.0 + + +def registry(store, **kwargs): + return Registry(store, ("project", "global", "auth"), {"model": "model"}, clock=lambda: NOW, **kwargs) + + +def created(seconds=3600): + return {"name": NAME, "expireTime": datetime.fromtimestamp(NOW + seconds, timezone.utc).isoformat()} + + +def test_only_owner_publishes_until_provider_expiry_minus_margin(): + store = Mock() + store.eval.return_value = 1 + resolved = registry(store) + assert resolved.publish("owner", created()) == NAME + assert store.eval.call_args.args[-1] == 3_540_000 + assert store.eval.call_args.args[2:5] == (resolved.lock, resolved.key, "owner") + store.eval.return_value = 0 + assert resolved.publish("expired-owner", created()) is None + + +@pytest.mark.parametrize("seconds", [-1, 0, 30, 60]) +def test_expired_or_near_expiry_names_are_not_published(seconds): + store = Mock() + assert registry(store).publish("owner", created(seconds)) is None + store.eval.assert_not_called() + + +@pytest.mark.parametrize( + "result", + [ + {}, + {"name": NAME}, + {"name": NAME, "expireTime": "not-a-date"}, + {"name": "invalid", "expireTime": "2030-01-01T00:00:00Z"}, + {"name": NAME, "expireTime": "2030-01-01T00:00:00"}, + ], +) +def test_unusable_provider_metadata_never_publishes(result): + store = Mock() + assert registry(store).publish("owner", result) is None + store.eval.assert_not_called() + + +@pytest.mark.parametrize( + "raw", + [ + "{}", + "null", + "[]", + '{"name":"bad","expires":1900003600}', + '{"name":"' + NAME + '","expires":1900000030}', + '{"name":"' + NAME + '","expires":"NaN"}', + ], +) +def test_untrusted_or_expiring_redis_records_are_misses(raw): + store = Mock() + store.get.return_value = raw + assert registry(store).read() is None + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.asyncio +async def test_redis_failure_bypasses_cache_creation(asynchronous): + store = Mock() + store.get.side_effect = RedisError("unavailable") + create = Mock() + resolved = registry(store) + result = await resolved.async_resolve(create) if asynchronous else resolved.resolve(create) + assert result is None + create.assert_not_called() + + +def test_transient_creation_failure_sets_shared_cooldown_and_releases_lease(): + store = Mock() + store.get.return_value = None + store.exists.return_value = 0 + store.set.return_value = True + response = httpx.Response(429, request=httpx.Request("POST", "https://provider.invalid")) + create = Mock(side_effect=httpx.HTTPStatusError("quota", request=response.request, response=response)) + resolved = registry(store) + assert resolved.resolve(create) is None + create.assert_called_once() + store.set.assert_any_call(resolved.cooldown, "1", ex=30) + assert store.eval.call_args.args[2] == resolved.lock + + +def test_keys_separate_tenant_and_payload_without_storing_either(): + store = Mock() + first = Registry(store, ("tenant-a",), {"text": "private prefix"}) + tenant = Registry(store, ("tenant-b",), {"text": "private prefix"}) + prefix = Registry(store, ("tenant-a",), {"text": "different prefix"}) + assert len({first.key, tenant.key, prefix.key}) == 3 + assert "private prefix" not in first.key + assert "tenant-a" not in first.key + assert first.key.split("{")[1].split("}")[0] == first.lock.split("{")[1].split("}")[0] + + +def test_valid_hit_does_not_create_or_acquire_a_lease(): + import json + + store = Mock() + store.get.return_value = json.dumps({"name": NAME, "expires": NOW + 3600}) + create = Mock() + assert registry(store).resolve(create) == NAME + create.assert_not_called() + store.set.assert_not_called() + + +def test_contention_exits_without_duplicate_creation_when_wait_budget_expires(): + store = Mock() + store.get.return_value = None + store.exists.return_value = 0 + store.set.return_value = False + ticks = iter((0.0, 3.0)) + create = Mock() + resolved = registry(store, monotonic=lambda: next(ticks), sleep=Mock()) + assert resolved.resolve(create) is None + create.assert_not_called() + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.asyncio +async def test_expired_managed_cache_retries_once_with_original_content(asynchronous): + from litellm.llms.vertex_ai.context_caching.redis_cache_ids import ( + CacheReference, + async_post_with_recovery, + post_with_recovery, + ) + + store = Mock() + reference = CacheReference(NAME, registry(store), [], {}) + full_body = {"contents": [{"role": "user", "parts": [{"text": "complete original request"}]}]} + reference.rebuild = lambda: full_body + body = {"cachedContent": reference, "contents": []} + received = [] + + def post(request_body): + received.append(dict(request_body)) + request = httpx.Request("POST", "https://provider.invalid/generateContent") + if len(received) == 1: + return httpx.Response(404, request=request, json={"error": {"message": "CachedContent not found"}}) + return httpx.Response(200, request=request, json={}) + + async def async_post(request_body): + return post(request_body) + + response = await async_post_with_recovery(async_post, body) if asynchronous else post_with_recovery(post, body) + assert response.status_code == 200 + assert len(received) == 2 + assert received[0]["cachedContent"] == NAME + assert received[1] == full_body + store.eval.assert_called_once() + assert store.eval.call_args.args[-1] == NAME + + +@pytest.mark.parametrize( + "status,message", + [(403, "permission denied"), (429, "quota exceeded"), (404, "model not found"), (400, "invalid request")], +) +def test_unrelated_provider_errors_are_not_retried(status, message): + from litellm.llms.vertex_ai.context_caching.redis_cache_ids import CacheReference, post_with_recovery + + store = Mock() + reference = CacheReference(NAME, registry(store), [], {}) + reference.rebuild = lambda: {"contents": []} + body = {"cachedContent": reference, "contents": []} + request = httpx.Request("POST", "https://provider.invalid/generateContent") + post = Mock(return_value=httpx.Response(status, request=request, json={"error": {"message": message}})) + with pytest.raises(httpx.HTTPStatusError): + post_with_recovery(post, body) + post.assert_called_once() + store.eval.assert_not_called() From 26088e98b5d8d9d192c83f700fd9e5c5cc4ea6fa Mon Sep 17 00:00:00 2001 From: Dominastorm Date: Wed, 14 Oct 2026 20:18:42 +0530 Subject: [PATCH 2/2] test(vertex): cover Redis cache ID lifecycle and recovery --- .../context_caching/test_redis_cache_ids.py | 471 ++++++++++++++++++ ...test_vertex_and_google_ai_studio_gemini.py | 40 ++ 2 files changed, 511 insertions(+) diff --git a/tests/unit/llms/vertex_ai/context_caching/test_redis_cache_ids.py b/tests/unit/llms/vertex_ai/context_caching/test_redis_cache_ids.py index b46e9f6910d..d4c9c8efbd5 100644 --- a/tests/unit/llms/vertex_ai/context_caching/test_redis_cache_ids.py +++ b/tests/unit/llms/vertex_ai/context_caching/test_redis_cache_ids.py @@ -43,6 +43,7 @@ def test_expired_or_near_expiry_names_are_not_published(seconds): {}, {"name": NAME}, {"name": NAME, "expireTime": "not-a-date"}, + {"name": NAME, "expireTime": None}, {"name": "invalid", "expireTime": "2030-01-01T00:00:00Z"}, {"name": NAME, "expireTime": "2030-01-01T00:00:00"}, ], @@ -182,3 +183,473 @@ def test_unrelated_provider_errors_are_not_retried(status, message): post_with_recovery(post, body) post.assert_called_once() store.eval.assert_not_called() + + +@pytest.fixture +def redis_store(monkeypatch): + from unittest.mock import patch + + from litellm.llms.vertex_ai.context_caching.redis_cache_ids import connection + + store = Mock() + store.get.return_value = None + store.exists.return_value = 0 + store.set.return_value = True + store.eval.return_value = 1 + monkeypatch.setenv("LITELLM_VERTEX_CACHE_ID_REDIS_URL", "redis://cache.invalid/0") + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + connection.cache_clear() + with patch("redis.Redis.from_url", return_value=store): + yield store + connection.cache_clear() + + +def cache_arguments(): + from litellm.litellm_core_utils.litellm_logging import Logging + + return dict( + messages=[ + { + "role": "system", + "content": "Synthetic reusable instructions. " * 10000, + "cache_control": {"type": "ephemeral"}, + }, + {"role": "user", "content": "Uncached question"}, + ], + optional_params={ + "tools": [{"functionDeclarations": [{"name": "lookup", "parameters": {"type": "OBJECT"}}]}], + "tool_choice": {"functionCallingConfig": {"mode": "AUTO"}}, + "temperature": 0.3, + }, + api_key="synthetic-api-key", + api_base=None, + model="gemini-2.5-flash", + timeout=10, + logging_obj=Mock(spec=Logging), + custom_llm_provider="vertex_ai", + vertex_project="project", + vertex_location="global", + vertex_auth_header="synthetic-token", + cache_id_auth_scope="synthetic-credential-fingerprint", + ) + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("supplied_client", [False, True]) +@pytest.mark.asyncio +async def test_public_entry_creates_then_reuses_id_preserving_tools_and_accounting( + redis_store, respx_mock, asynchronous, supplied_client +): + import copy + import json + + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler + from litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching import ContextCachingEndpoints + + arguments = cache_arguments() + before = copy.deepcopy((arguments["messages"], arguments["optional_params"])) + response = {**created(), "expireTime": "9999-01-01T00:00:00Z", "usageMetadata": {"totalTokenCount": 1234}} + route = respx_mock.post(url__regex=r"https://.*?/cachedContents$").respond(200, json=response) + client = (AsyncHTTPHandler() if asynchronous else HTTPHandler()) if supplied_client else None + endpoint = ContextCachingEndpoints() + try: + first = ( + await endpoint.async_check_and_create_cache(client=client, **arguments) + if asynchronous + else endpoint.check_and_create_cache(client=client, **arguments) + ) + assert first[2] == NAME + assert first[0] == before[0][1:] + assert first[1] == {"temperature": 0.3} + assert arguments["logging_obj"].vertex_ai_cache_creation_input_tokens == 1234 + outbound = json.loads(route.calls[0].request.content) + assert outbound["tools"] == before[1]["tools"] + assert outbound["toolConfig"] == before[1]["tool_choice"] + assert route.calls[0].request.headers["authorization"] == "Bearer synthetic-token" + redis_store.get.return_value = json.dumps({"name": NAME, "expires": 253370764800}) + second = ( + await endpoint.async_check_and_create_cache(client=client, **arguments) + if asynchronous + else endpoint.check_and_create_cache(client=client, **arguments) + ) + assert second == first + assert route.call_count == 1 + assert (arguments["messages"], arguments["optional_params"]) == before + if client is not None: + assert not client.client.is_closed + finally: + if isinstance(client, AsyncHTTPHandler): + await client.close() + elif client is not None: + client.close() + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("reason", ["unmarked", "too_short", "short_ttl", "assistant_turn", "redis_down", "cooldown"]) +@pytest.mark.asyncio +async def test_ineligible_or_unavailable_cache_keeps_full_request(redis_store, respx_mock, asynchronous, reason): + import copy + + from litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching import ContextCachingEndpoints + + arguments = cache_arguments() + if reason == "unmarked": + del arguments["messages"][0]["cache_control"] + elif reason == "too_short": + arguments["messages"][0]["content"] = "short" + elif reason == "short_ttl": + arguments["messages"][0]["content"] = [ + { + "type": "text", + "text": arguments["messages"][0]["content"], + "cache_control": {"type": "ephemeral", "ttl": "30s"}, + } + ] + elif reason == "assistant_turn": + arguments["messages"][0]["role"] = "assistant" + elif reason == "redis_down": + redis_store.get.side_effect = RedisError("unavailable") + else: + redis_store.exists.return_value = 1 + before = copy.deepcopy((arguments["messages"], arguments["optional_params"])) + endpoint = ContextCachingEndpoints() + result = ( + await endpoint.async_check_and_create_cache(client=None, **arguments) + if asynchronous + else endpoint.check_and_create_cache(client=None, **arguments) + ) + assert result == (*before, None) + assert result[0] is arguments["messages"] + assert result[1] is arguments["optional_params"] + assert len(respx_mock.calls) == 0 + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("status", [403, 429]) +@pytest.mark.asyncio +async def test_public_creation_error_preserves_auth_failure_or_fails_open_on_quota( + redis_store, respx_mock, asynchronous, status +): + from litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching import ContextCachingEndpoints + + arguments = cache_arguments() + route = respx_mock.post(url__regex=r"https://.*?/cachedContents$").respond(status, json={"error": "test"}) + endpoint = ContextCachingEndpoints() + + async def call(): + if asynchronous: + return await endpoint.async_check_and_create_cache(client=None, **arguments) + return endpoint.check_and_create_cache(client=None, **arguments) + + if status == 403: + with pytest.raises(httpx.HTTPStatusError) as error: + await call() + assert error.value.response.status_code == status + else: + assert await call() == (arguments["messages"], arguments["optional_params"], None) + assert any(call.kwargs.get("ex") == 30 for call in redis_store.set.call_args_list) + assert route.call_count == 1 + assert redis_store.eval.call_args.args[2].endswith(":lock") + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("published", [True, False]) +@pytest.mark.asyncio +async def test_waiter_rechecks_publication_or_cooldown_after_winning_lease(asynchronous, published): + import json + from unittest.mock import AsyncMock + + store = Mock() + store.get.side_effect = [None, json.dumps({"name": NAME, "expires": NOW + 3600}) if published else None] + store.exists.side_effect = [0, 1] + store.set.return_value = True + create = AsyncMock() if asynchronous else Mock() + resolved = registry(store) + result = await resolved.async_resolve(create) if asynchronous else resolved.resolve(create) + assert result == (NAME if published else None) + create.assert_not_called() + assert store.eval.call_args.args[2] == resolved.lock + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("published", [True, False]) +@pytest.mark.asyncio +async def test_contender_observes_winner_or_exits_at_deadline(asynchronous, published): + import json + from unittest.mock import AsyncMock + + store = Mock() + store.get.side_effect = [None, json.dumps({"name": NAME, "expires": NOW + 3600}) if published else None] + store.exists.return_value = 0 + store.set.return_value = False + ticks = iter((0.0, 0.1, 3.0)) + create = AsyncMock() if asynchronous else Mock() + resolved = registry(store, monotonic=lambda: next(ticks), sleep=Mock(), async_sleep=AsyncMock()) + result = await resolved.async_resolve(create) if asynchronous else resolved.resolve(create) + assert result == (NAME if published else None) + create.assert_not_called() + store.eval.assert_not_called() + + +@pytest.mark.asyncio +async def test_cancelled_async_creator_releases_lease(): + import asyncio + + store = Mock() + store.get.return_value = None + store.exists.return_value = 0 + store.set.return_value = True + started = asyncio.Event() + blocked = asyncio.Event() + + async def create(): + started.set() + await blocked.wait() + return created() + + resolved = registry(store) + task = asyncio.create_task(resolved.async_resolve(create)) + await started.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert store.eval.call_args.args[2] == resolved.lock + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("error", [httpx.ConnectError("connection failed"), ValueError("bad provider response")]) +@pytest.mark.asyncio +async def test_creation_transport_failure_fails_open_but_unexpected_failure_propagates(asynchronous, error): + from unittest.mock import AsyncMock + + store = Mock() + store.get.return_value = None + store.exists.return_value = 0 + store.set.return_value = True + create = AsyncMock(side_effect=error) if asynchronous else Mock(side_effect=error) + resolved = registry(store) + if isinstance(error, ValueError): + with pytest.raises(ValueError, match="bad provider response"): + if asynchronous: + await resolved.async_resolve(create) + else: + resolved.resolve(create) + else: + assert (await resolved.async_resolve(create) if asynchronous else resolved.resolve(create)) is None + store.set.assert_any_call(resolved.cooldown, "1", ex=30) + assert store.eval.call_args.args[2] == resolved.lock + + +@pytest.mark.parametrize("operation", ["release", "invalidate", "cool_down"]) +def test_redis_cleanup_failure_does_not_interrupt_generation(operation): + store = Mock() + store.eval.side_effect = RedisError("unavailable") + store.set.side_effect = RedisError("unavailable") + resolved = registry(store) + if operation == "release": + assert resolved.release("owner") is None + elif operation == "invalidate": + assert resolved.invalidate(NAME) is None + else: + assert resolved.cool_down() is None + + +def test_expiry_recreates_instead_of_reusing_stale_name(): + import json + + store = Mock() + store.get.return_value = json.dumps({"name": NAME, "expires": NOW + 30}) + store.exists.return_value = 0 + store.set.return_value = True + store.eval.return_value = 1 + replacement = NAME + "-replacement" + create = Mock(return_value={**created(), "name": replacement}) + assert registry(store).resolve(create) == replacement + create.assert_called_once() + + +@pytest.mark.parametrize("source", ["mapping", "json", "file"]) +def test_credential_fingerprint_stays_stable_without_exposing_secret(monkeypatch, source): + import json + from unittest.mock import mock_open, patch + + from litellm.llms.vertex_ai.context_caching.redis_cache_ids import credential_scope + + monkeypatch.setenv("LITELLM_VERTEX_CACHE_ID_REDIS_URL", "redis://cache.invalid/0") + credentials = {"client_email": "test@example.invalid", "private_key": "synthetic-secret"} + expected = credential_scope(credentials) + argument = credentials if source == "mapping" else json.dumps(credentials) if source == "json" else "key.json" + with ( + patch("os.path.isfile", return_value=source == "file"), + patch("builtins.open", mock_open(read_data=json.dumps(credentials))), + ): + assert credential_scope(argument) == expected + assert len(expected) == 64 + assert "synthetic-secret" not in expected + assert credential_scope({**credentials, "client_email": "other@example.invalid"}) != expected + + +@pytest.mark.parametrize("credentials", [None, {}, "not-json", "{}", "[]"]) +def test_unknown_credentials_fall_back_to_token_scope(monkeypatch, credentials): + from litellm.llms.vertex_ai.context_caching.redis_cache_ids import credential_scope + + monkeypatch.setenv("LITELLM_VERTEX_CACHE_ID_REDIS_URL", "redis://cache.invalid/0") + assert credential_scope(credentials) is None + monkeypatch.delenv("LITELLM_VERTEX_CACHE_ID_REDIS_URL") + assert credential_scope({"private_key": "synthetic-secret"}) is None + + +@pytest.mark.parametrize("override", [False, True]) +def test_rotating_token_reuses_stable_scope_unless_authorization_is_overridden(redis_store, override): + from litellm.llms.vertex_ai.context_caching.redis_cache_ids import prepare + from litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching import ContextCachingEndpoints + + endpoint = ContextCachingEndpoints() + first_args = cache_arguments() + second_args = {**first_args, "vertex_auth_header": "new-token"} + if override: + first_args["extra_headers"] = {"Authorization": "Bearer first"} + second_args["extra_headers"] = {"Authorization": "Bearer second"} + first = prepare(endpoint, first_args) + second = prepare(endpoint, second_args) + assert first is not None and second is not None + assert (first.registry.key == second.registry.key) is (not override) + assert first.headers["Authorization"] == ("Bearer first" if override else "Bearer synthetic-token") + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.asyncio +async def test_unusable_redis_connection_returns_original_request(monkeypatch, asynchronous): + from unittest.mock import patch + + from litellm.llms.vertex_ai.context_caching.redis_cache_ids import connection + from litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching import ContextCachingEndpoints + + monkeypatch.setenv("LITELLM_VERTEX_CACHE_ID_REDIS_URL", "redis://cache.invalid/0") + connection.cache_clear() + arguments = cache_arguments() + endpoint = ContextCachingEndpoints() + with patch("redis.Redis.from_url", side_effect=RedisError("unavailable")): + result = ( + await endpoint.async_check_and_create_cache(client=None, **arguments) + if asynchronous + else endpoint.check_and_create_cache(client=None, **arguments) + ) + assert result == (arguments["messages"], arguments["optional_params"], None) + connection.cache_clear() + + +def test_redis_is_optional_for_base_install(): + import builtins + from unittest.mock import patch + + from litellm.llms.vertex_ai.context_caching.redis_cache_ids import redis_error_type + + original_import = builtins.__import__ + + def without_redis(name, *args, **kwargs): + if name == "redis.exceptions": + raise ImportError("redis is not installed") + return original_import(name, *args, **kwargs) + + with patch("builtins.__import__", side_effect=without_redis): + assert redis_error_type() is ImportError + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("managed", [False, True]) +@pytest.mark.asyncio +async def test_successful_generation_never_invalidates_or_retries(asynchronous, managed): + from unittest.mock import AsyncMock + + from litellm.llms.vertex_ai.context_caching.redis_cache_ids import ( + CacheReference, + async_post_with_recovery, + post_with_recovery, + ) + + store = Mock() + body = {"cachedContent": CacheReference(NAME, registry(store), [], {}) if managed else NAME} + response = httpx.Response(200, request=httpx.Request("POST", "https://provider.invalid"), json={"answer": "ok"}) + post = AsyncMock(return_value=response) if asynchronous else Mock(return_value=response) + assert (await async_post_with_recovery(post, body) if asynchronous else post_with_recovery(post, body)) is response + post.assert_called_once_with(body) + store.eval.assert_not_called() + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("raw", [b"not json", b"[]", b'{"error":{"message":"model not found"}}']) +@pytest.mark.asyncio +async def test_unknown_error_body_propagates_without_retry(asynchronous, raw): + from unittest.mock import AsyncMock + + from litellm.llms.vertex_ai.context_caching.redis_cache_ids import ( + CacheReference, + async_post_with_recovery, + post_with_recovery, + ) + + store = Mock() + reference = CacheReference(NAME, registry(store), [], {}) + reference.rebuild = lambda: {"contents": []} + response = httpx.Response(404, request=httpx.Request("POST", "https://provider.invalid"), content=raw) + post = AsyncMock(return_value=response) if asynchronous else Mock(return_value=response) + with pytest.raises(httpx.HTTPStatusError) as error: + if asynchronous: + await async_post_with_recovery(post, {"cachedContent": reference}) + else: + post_with_recovery(post, {"cachedContent": reference}) + assert error.value.response is response + post.assert_called_once() + store.eval.assert_not_called() + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.asyncio +async def test_recovery_rebuilds_full_gemini_body_and_propagates_second_failure(asynchronous): + import copy + from unittest.mock import AsyncMock + + from litellm.llms.vertex_ai.context_caching.redis_cache_ids import ( + CacheReference, + async_post_with_recovery, + post_with_recovery, + ) + from litellm.llms.vertex_ai.gemini.transformation import transform_request_body + + arguments = cache_arguments() + params = {"temperature": 0.3} + original = copy.deepcopy((arguments["messages"], params)) + store = Mock() + reference = copy.deepcopy(CacheReference(NAME, registry(store), arguments["messages"], params)) + body = transform_request_body( + messages=arguments["messages"][1:], + model=arguments["model"], + optional_params=params, + custom_llm_provider="vertex_ai", + litellm_params={}, + cached_content=reference, + ) + expected = transform_request_body( + messages=copy.deepcopy(original[0]), + model=arguments["model"], + optional_params=copy.deepcopy(original[1]), + custom_llm_provider="vertex_ai", + litellm_params={}, + cached_content=None, + ) + request = httpx.Request("POST", "https://provider.invalid") + expired = httpx.Response(404, request=request, json={"error": {"message": "CachedContent expired"}}) + denied = httpx.Response(403, request=request, json={"error": {"message": "permission denied"}}) + post = AsyncMock(side_effect=[expired, denied]) if asynchronous else Mock(side_effect=[expired, denied]) + with pytest.raises(httpx.HTTPStatusError) as error: + if asynchronous: + await async_post_with_recovery(post, body) + else: + post_with_recovery(post, body) + assert error.value.response is denied + assert post.call_count == 2 + assert body == expected + assert "cachedContent" not in body + assert (arguments["messages"], params) == original + assert expired.is_closed diff --git a/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 3f49869b19f..a70fc4a906c 100644 --- a/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -9684,3 +9684,43 @@ def test_streaming_usage_reads_cache_creation_tokens_recorded_on_logging_obj(): assert chunk.usage.prompt_tokens == 1536 + 1280 assert chunk.usage.cache_creation_input_tokens == 1280 + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.asyncio +async def test_managed_cache_streaming_error_preserves_vertex_status(respx_mock, asynchronous): + from litellm.llms.vertex_ai.context_caching.redis_cache_ids import CacheReference, Registry + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import make_call, make_sync_call + + reference = CacheReference("projects/p/locations/global/cachedContents/c", Registry(Mock(), ("scope",), {}), [], {}) + reference.rebuild = lambda: {"contents": []} + body = {"cachedContent": reference, "contents": []} + route = respx_mock.post("https://provider.invalid/streamGenerateContent").respond( + 403, json={"error": {"message": "permission denied"}} + ) + client = AsyncHTTPHandler(transport=httpx.AsyncHTTPTransport()) if asynchronous else HTTPHandler() + try: + kwargs = dict( + client=client, + gemini_client=client, + api_base="https://provider.invalid/streamGenerateContent", + headers={}, + data=json.dumps(body), + model="gemini-2.5-flash", + messages=[], + logging_obj=Mock(), + cache_body=body, + ) + with pytest.raises(VertexAIError) as error: + if asynchronous: + await make_call(**kwargs) + else: + make_sync_call(**kwargs) + assert error.value.status_code == 403 + assert "permission denied" in str(error.value) + assert route.call_count == 1 + finally: + if asynchronous: + await client.close() + else: + client.close()