diff --git a/src/google/adk/models/gemini_context_cache_manager.py b/src/google/adk/models/gemini_context_cache_manager.py index bbe0d0677e..f7f51206b4 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 dd27d42c1b..782e794a2d 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 0350166846..e1bf7bf8b8 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)