From f8ce8aab33bb058257924a97fd38bbaca4f029b3 Mon Sep 17 00:00:00 2001 From: Gabriel Mesquita Date: Tue, 18 Aug 2026 22:37:18 -0300 Subject: [PATCH] fix(models): enforce 32k token floor for Vertex AI context caching and add graceful fallback On Google Cloud Vertex AI, creating context cache requires a minimum of 32,768 tokens (unlike Google AI Studio which supports lower thresholds). When using models like gemini-3.5-flash on Vertex AI with prompts below 32k tokens, caches.create failed with INVALID_ARGUMENT and uncaught exceptions aborted model generation. - Add _VERTEX_AI_MIN_CACHE_TOKENS = 32768 in GeminiContextCacheManager - Check is_vertex in _minimum_cache_tokens before attempting cache creation - Add try/except error handling around handle_context_caching in Gemini.generate_content_async to gracefully fall back without breaking streams - Add unit tests verifying Vertex AI token floor and graceful cache handling --- .../models/gemini_context_cache_manager.py | 14 ++++-- src/google/adk/models/google_llm.py | 26 ++++++---- .../test_gemini_context_cache_manager.py | 47 +++++++++++++++++++ 3 files changed, 76 insertions(+), 11 deletions(-) diff --git a/src/google/adk/models/gemini_context_cache_manager.py b/src/google/adk/models/gemini_context_cache_manager.py index bbe0d0677ef..f7f51206b47 100644 --- a/src/google/adk/models/gemini_context_cache_manager.py +++ b/src/google/adk/models/gemini_context_cache_manager.py @@ -36,17 +36,22 @@ logger = logging.getLogger("google_adk." + __name__) -# Named Gemini model families have documented explicit-cache floors. For +# Named Gemini model families and backends have documented explicit-cache floors. For # opaque tuned-model and endpoint IDs, the server remains authoritative. _GEMINI_2_5_MIN_CACHE_TOKENS = 2048 _GEMINI_3_MIN_CACHE_TOKENS = 4096 +_VERTEX_AI_MIN_CACHE_TOKENS = 32768 if TYPE_CHECKING: from google.genai import Client -def _minimum_cache_tokens(model: Optional[str]) -> Optional[int]: +def _minimum_cache_tokens( + model: Optional[str], is_vertex: bool = False +) -> Optional[int]: """Return the explicit-cache token floor for a named Gemini model.""" + if is_vertex: + return _VERTEX_AI_MIN_CACHE_TOKENS model_name = (model or "").rsplit("/", maxsplit=1)[-1] if model_name.startswith("gemini-2.5-"): return _GEMINI_2_5_MIN_CACHE_TOKENS @@ -424,7 +429,10 @@ async def _create_new_cache_with_contents( cacheable_prefix_tokens = self._estimate_cacheable_prefix_tokens( llm_request, cache_contents_count ) - minimum_cache_tokens = _minimum_cache_tokens(llm_request.model) + minimum_cache_tokens = _minimum_cache_tokens( + llm_request.model, + is_vertex=bool(self.genai_client.vertexai), + ) if ( minimum_cache_tokens is not None and cacheable_prefix_tokens < minimum_cache_tokens diff --git a/src/google/adk/models/google_llm.py b/src/google/adk/models/google_llm.py index dd27d42c1b9..782e794a2d9 100644 --- a/src/google/adk/models/google_llm.py +++ b/src/google/adk/models/google_llm.py @@ -234,14 +234,24 @@ async def generate_content_async( from .gemini_context_cache_manager import GeminiContextCacheManager with tracer.start_as_current_span('handle_context_caching') as span: - cache_manager = GeminiContextCacheManager(self.api_client) - cache_metadata = await cache_manager.handle_context_caching(llm_request) - if cache_metadata: - if cache_metadata.cache_name: - span.set_attribute('cache_action', 'active_cache') - span.set_attribute('cache_name', cache_metadata.cache_name) - else: - span.set_attribute('cache_action', 'fingerprint_only') + try: + cache_manager = GeminiContextCacheManager(self.api_client) + cache_metadata = await cache_manager.handle_context_caching( + llm_request + ) + if cache_metadata: + if cache_metadata.cache_name: + span.set_attribute('cache_action', 'active_cache') + span.set_attribute('cache_name', cache_metadata.cache_name) + else: + span.set_attribute('cache_action', 'fingerprint_only') + except Exception as e: + logger.warning( + 'Failed to handle context caching, proceeding without cache: %s', + e, + ) + cache_metadata = None + cache_manager = None logger.info( 'Sending out request, model: %s, backend: %s, stream: %s', diff --git a/tests/unittests/agents/test_gemini_context_cache_manager.py b/tests/unittests/agents/test_gemini_context_cache_manager.py index 0350166846e..e1bf7bf8b81 100644 --- a/tests/unittests/agents/test_gemini_context_cache_manager.py +++ b/tests/unittests/agents/test_gemini_context_cache_manager.py @@ -344,6 +344,53 @@ async def test_completed_turn_grows_cacheable_prefix(self): assert create_config.contents == [first_user, first_model] assert next_request.contents == [next_user] + async def test_vertex_ai_skips_cache_below_32768_token_minimum(self): + """Vertex AI skips an explicit cache below its 32,768-token floor even for Gemini 3.5.""" + mock_client = AsyncMock(spec=Client) + mock_client.vertexai = True + manager = GeminiContextCacheManager(mock_client) + + llm_request = self.create_llm_request(contents_count=0) + llm_request.model = "gemini-3.5-flash-lite" + llm_request.config.system_instruction = "x" * 32_000 + llm_request.cacheable_contents_token_count = 8_000 + llm_request.cache_metadata = CacheMetadata( + fingerprint=manager._generate_cache_fingerprint(llm_request, 0), + contents_count=0, + ) + + result = await manager.handle_context_caching(llm_request) + + assert result is not None + assert result.cache_name is None + manager.genai_client.aio.caches.create.assert_not_called() + + async def test_vertex_ai_creates_cache_above_32768_token_minimum(self): + """Vertex AI creates an explicit cache above its 32,768-token floor for Gemini 3.5.""" + mock_client = AsyncMock(spec=Client) + mock_client.vertexai = True + manager = GeminiContextCacheManager(mock_client) + + llm_request = self.create_llm_request(contents_count=0) + llm_request.model = "gemini-3.5-flash-lite" + llm_request.config.system_instruction = "x" * 140_000 + llm_request.cacheable_contents_token_count = 35_000 + llm_request.cache_metadata = CacheMetadata( + fingerprint=manager._generate_cache_fingerprint(llm_request, 0), + contents_count=0, + ) + cached_content = AsyncMock() + cached_content.name = "projects/test/locations/us/cachedContents/vertex-35" + manager.genai_client.aio.caches.create = AsyncMock( + return_value=cached_content + ) + + result = await manager.handle_context_caching(llm_request) + + assert result is not None + assert result.cache_name == "projects/test/locations/us/cachedContents/vertex-35" + manager.genai_client.aio.caches.create.assert_awaited_once() + async def test_gemini_25_creates_cache_above_2048_token_minimum(self): """Gemini 2.5 creates an explicit cache above its 2,048-token floor.""" llm_request = self.create_llm_request(contents_count=0)