fix(models): enforce 32k token floor for Vertex AI context caching and add graceful fallback - #6810
Open
Gabriel-Mesq wants to merge 1 commit into
Open
Conversation
…d 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
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description of the Change
When using ADK with Google Cloud Vertex AI (
vertexai=True), context caching requires a minimum token floor of 32,768 tokens (unlike Google AI Studio which supports lower thresholds depending on the model, e.g. 4096 tokens for Gemini 3 / 3.5).Currently,
GeminiContextCacheManageronly checks model name prefixes against AI Studio thresholds, causing it to issuecaches.createrequests to Vertex AI with prompts between 4,096 and 32,768 tokens. Vertex AI rejects these with400 INVALID_ARGUMENT, and becausegoogle_llm.pylacked exception handling around context caching, the unhandled error terminated the entire LLM generation and aborted active response streams.Changes Made
Backend Floor Enforcement (
src/google/adk/models/gemini_context_cache_manager.py):_VERTEX_AI_MIN_CACHE_TOKENS = 32768._minimum_cache_tokensto acceptis_vertex=bool(self.genai_client.vertexai)and enforce the 32,768 token floor when running on Vertex AI.Graceful Fallback (
src/google/adk/models/google_llm.py):cache_manager.handle_context_caching()in atry...exceptblock inGemini.generate_content_async.logger.warning) and proceeds with regular generation without breaking active streams.Unit Tests (
tests/unittests/agents/test_gemini_context_cache_manager.py):test_vertex_ai_skips_cache_below_32768_token_minimumverifying Vertex AI skips cache creation below 32k tokens (tested withgemini-3.5-flash-lite).test_vertex_ai_creates_cache_above_32768_token_minimumverifying Vertex AI successfully creates cache when above 32k tokens.Testing Plan
Unit Tests Executed