Skip to content

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
google:mainfrom
Gabriel-Mesq:fix/vertex-context-cache-threshold
Open

fix(models): enforce 32k token floor for Vertex AI context caching and add graceful fallback#6810
Gabriel-Mesq wants to merge 1 commit into
google:mainfrom
Gabriel-Mesq:fix/vertex-context-cache-threshold

Conversation

@Gabriel-Mesq

Copy link
Copy Markdown

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, GeminiContextCacheManager only checks model name prefixes against AI Studio thresholds, causing it to issue caches.create requests to Vertex AI with prompts between 4,096 and 32,768 tokens. Vertex AI rejects these with 400 INVALID_ARGUMENT, and because google_llm.py lacked exception handling around context caching, the unhandled error terminated the entire LLM generation and aborted active response streams.

Changes Made

  1. Backend Floor Enforcement (src/google/adk/models/gemini_context_cache_manager.py):

    • Added _VERTEX_AI_MIN_CACHE_TOKENS = 32768.
    • Updated _minimum_cache_tokens to accept is_vertex=bool(self.genai_client.vertexai) and enforce the 32,768 token floor when running on Vertex AI.
  2. Graceful Fallback (src/google/adk/models/google_llm.py):

    • Wrapped cache_manager.handle_context_caching() in a try...except block in Gemini.generate_content_async.
    • If cache initialization fails for any reason (token floor, quota, transient backend issue), it logs a warning (logger.warning) and proceeds with regular generation without breaking active streams.
  3. Unit Tests (tests/unittests/agents/test_gemini_context_cache_manager.py):

    • Added test_vertex_ai_skips_cache_below_32768_token_minimum verifying Vertex AI skips cache creation below 32k tokens (tested with gemini-3.5-flash-lite).
    • Added test_vertex_ai_creates_cache_above_32768_token_minimum verifying Vertex AI successfully creates cache when above 32k tokens.

Testing Plan

Unit Tests Executed

pytest tests/unittests/agents/test_gemini_context_cache_manager.py tests/unittests/models/test_google_llm.py

…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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants