diff --git a/src/google/adk/flows/llm_flows/context_cache_processor.py b/src/google/adk/flows/llm_flows/context_cache_processor.py index b274b3ab..a9e9eb76 100644 --- a/src/google/adk/flows/llm_flows/context_cache_processor.py +++ b/src/google/adk/flows/llm_flows/context_cache_processor.py @@ -62,9 +62,11 @@ class ContextCacheRequestProcessor(BaseLlmRequestProcessor): # Set cache config to request llm_request.cache_config = invocation_context.context_cache_config - # Find latest cache metadata from session events - latest_cache_metadata = self._find_latest_cache_metadata( - invocation_context, agent.name, invocation_context.invocation_id + # Find latest cache metadata and previous token count from session events + latest_cache_metadata, previous_token_count = ( + self._find_cache_info_from_events( + invocation_context, agent.name, invocation_context.invocation_id + ) ) if latest_cache_metadata: @@ -77,51 +79,78 @@ class ContextCacheRequestProcessor(BaseLlmRequestProcessor): latest_cache_metadata.cached_contents_count, ) + if previous_token_count is not None: + llm_request.cacheable_contents_token_count = previous_token_count + logger.debug( + 'Found previous prompt token count for agent %s: %d', + agent.name, + previous_token_count, + ) + logger.debug('Context caching enabled for agent %s', agent.name) # This processor yields no events return yield # AsyncGenerator requires a yield in function body - def _find_latest_cache_metadata( + def _find_cache_info_from_events( self, invocation_context: 'InvocationContext', agent_name: str, current_invocation_id: str, - ) -> Optional[CacheMetadata]: - """Find the latest cache metadata from session events. + ) -> tuple[Optional[CacheMetadata], Optional[int]]: + """Find cache metadata and previous token count from session events. Args: invocation_context: Context containing session with events - agent_name: Name of agent to find cache metadata for + agent_name: Name of agent to find cache info for current_invocation_id: Current invocation ID to compare for increment Returns: - Latest cache metadata for the agent (with updated invocations_used - if needed), or None if not found + Tuple of (cache_metadata, previous_prompt_token_count) + cache_metadata: Latest cache metadata with updated invocations_used if needed + previous_prompt_token_count: Most recent prompt token count from LLM response """ if not invocation_context.session or not invocation_context.session.events: - return None + return None, None + + cache_metadata = None + previous_token_count = None # Search events from most recent to oldest using index traversal events = invocation_context.session.events for i in range(len(events) - 1, -1, -1): event = events[i] - if event.cache_metadata is not None and event.author == agent_name: - - cache_metadata = event.cache_metadata + if event.author != agent_name: + continue + # Look for cache metadata (only in actual LLM response events) + if cache_metadata is None and event.cache_metadata is not None: # Check if this is a different invocation - increment invocations_used if event.invocation_id and event.invocation_id != current_invocation_id: # Different invocation - increment invocations_used - return cache_metadata.model_copy( - update={'invocations_used': cache_metadata.invocations_used + 1} + cache_metadata = event.cache_metadata.model_copy( + update={ + 'invocations_used': event.cache_metadata.invocations_used + 1 + } ) else: # Same invocation or no invocation_id - return as-is - return cache_metadata + cache_metadata = event.cache_metadata - return None + # Look for previous prompt token count (from actual LLM response events) + if ( + previous_token_count is None + and event.usage_metadata + and event.usage_metadata.prompt_token_count is not None + ): + previous_token_count = event.usage_metadata.prompt_token_count + + # Stop early if we found both pieces of information + if cache_metadata is not None and previous_token_count is not None: + break + + return cache_metadata, previous_token_count # Create processor instance for use in flows diff --git a/src/google/adk/models/gemini_context_cache_manager.py b/src/google/adk/models/gemini_context_cache_manager.py index 31eb1af8..5c1aeaf5 100644 --- a/src/google/adk/models/gemini_context_cache_manager.py +++ b/src/google/adk/models/gemini_context_cache_manager.py @@ -257,12 +257,21 @@ class GeminiContextCacheManager: Returns: Cache metadata if successful, None otherwise """ - # Estimate token count for minimum cache size check - estimated_tokens = self._estimate_request_tokens(llm_request) - if estimated_tokens < llm_request.cache_config.min_tokens: + # Check if we have token count from previous response for cache size validation + if llm_request.cacheable_contents_token_count is None: logger.info( - "Request too small for caching (%d < %d tokens)", - estimated_tokens, + "No previous token count available, skipping cache creation for" + " initial request" + ) + return None + + if ( + llm_request.cacheable_contents_token_count + < llm_request.cache_config.min_tokens + ): + logger.info( + "Previous request too small for caching (%d < %d tokens)", + llm_request.cacheable_contents_token_count, llm_request.cache_config.min_tokens, ) return None diff --git a/src/google/adk/models/llm_request.py b/src/google/adk/models/llm_request.py index 814e0dbe..04a61fd9 100644 --- a/src/google/adk/models/llm_request.py +++ b/src/google/adk/models/llm_request.py @@ -88,6 +88,9 @@ class LlmRequest(BaseModel): cache_metadata: Optional[CacheMetadata] = None """Cache metadata from previous requests, used for cache management.""" + cacheable_contents_token_count: Optional[int] = None + """Token count from previous request's prompt, used for cache size validation.""" + def append_instructions( self, instructions: Union[list[str], types.Content] ) -> list[types.Content]: diff --git a/tests/unittests/agents/test_gemini_context_cache_manager.py b/tests/unittests/agents/test_gemini_context_cache_manager.py index ff1d1c27..23b678cc 100644 --- a/tests/unittests/agents/test_gemini_context_cache_manager.py +++ b/tests/unittests/agents/test_gemini_context_cache_manager.py @@ -121,6 +121,9 @@ class TestGeminiContextCacheManager: ) llm_request = self.create_llm_request() + llm_request.cacheable_contents_token_count = ( + 2048 # Add token count for cache creation + ) start_time = time.time() with patch.object( @@ -194,6 +197,9 @@ class TestGeminiContextCacheManager: invocations_used=15 ) # Exceeds cache_intervals llm_request = self.create_llm_request(cache_metadata=existing_cache) + llm_request.cacheable_contents_token_count = ( + 2048 # Add token count for cache creation + ) with ( patch.object(self.manager, "_is_cache_valid", return_value=False), @@ -521,3 +527,65 @@ class TestGeminiContextCacheManager: assert not hasattr( cache_metadata, "usage_metadata" ) # CacheMetadata should NOT have this + + def create_llm_request_with_token_count( + self, token_count=None, cache_metadata=None + ): + """Helper to create LlmRequest with cacheable_contents_token_count.""" + llm_request = self.create_llm_request(cache_metadata=cache_metadata) + llm_request.cacheable_contents_token_count = token_count + return llm_request + + async def test_cache_creation_with_sufficient_token_count(self): + """Test cache creation succeeds when token count meets minimum.""" + # Setup mocks + mock_cached_content = AsyncMock() + mock_cached_content.name = ( + "projects/test/locations/us-central1/cachedContents/token123" + ) + self.manager.genai_client.aio.caches.create = AsyncMock( + return_value=mock_cached_content + ) + + # Create request with sufficient token count + llm_request = self.create_llm_request_with_token_count(token_count=2048) + + with patch.object( + self.manager, "_generate_cache_fingerprint", return_value="test_fp" + ): + result = await self.manager.handle_context_caching(llm_request) + + # Should succeed in creating cache + assert result is not None + assert result.cache_name == mock_cached_content.name + self.manager.genai_client.aio.caches.create.assert_called_once() + + async def test_cache_creation_with_insufficient_token_count(self): + """Test cache creation fails when token count is below minimum.""" + # Set higher minimum token requirement + self.manager.cache_config = ContextCacheConfig( + cache_intervals=10, + ttl_seconds=1800, + min_tokens=2048, + ) + + # Create request with insufficient token count + llm_request = self.create_llm_request_with_token_count(token_count=1024) + llm_request.cache_config = self.manager.cache_config + + result = await self.manager.handle_context_caching(llm_request) + + # Should not create cache + assert result is None + self.manager.genai_client.aio.caches.create.assert_not_called() + + async def test_cache_creation_without_token_count(self): + """Test cache creation is skipped when no token count is available.""" + # Create request without token count (initial request) + llm_request = self.create_llm_request_with_token_count(token_count=None) + + result = await self.manager.handle_context_caching(llm_request) + + # Should skip cache creation for initial request + assert result is None + self.manager.genai_client.aio.caches.create.assert_not_called() diff --git a/tests/unittests/flows/llm_flows/test_context_cache_processor.py b/tests/unittests/flows/llm_flows/test_context_cache_processor.py index 31ee270d..026bdc85 100644 --- a/tests/unittests/flows/llm_flows/test_context_cache_processor.py +++ b/tests/unittests/flows/llm_flows/test_context_cache_processor.py @@ -452,3 +452,195 @@ class TestContextCacheRequestProcessor: assert llm_request.cache_config == self.cache_config assert llm_request.cache_metadata is not None assert llm_request.cache_metadata.invocations_used == 11 # 10 + 1 + + async def test_cacheable_contents_token_count_extraction(self): + """Test that previous prompt token count is extracted and set.""" + agent = LlmAgent(name="test_agent") + + # Create event with usage metadata + event_with_tokens = Event( + author="test_agent", + usage_metadata=types.UsageMetadata( + prompt_token_count=1024, + response_token_count=256, + total_token_count=1280, + ), + ) + + events = [event_with_tokens] + + invocation_context = self.create_invocation_context( + agent, + context_cache_config=self.cache_config, + session_events=events, + ) + + llm_request = LlmRequest( + model="gemini-2.0-flash", + contents=[ + types.Content( + role="user", + parts=[types.Part(text="Hello")], + ) + ], + ) + + async for event in self.processor.run_async( + invocation_context, llm_request + ): + pass + + # Should extract token count from the event + assert llm_request.cacheable_contents_token_count == 1024 + + async def test_cacheable_contents_token_count_no_usage_metadata(self): + """Test when no usage metadata is available.""" + agent = LlmAgent(name="test_agent") + + events = [ + Event(author="test_agent", usage_metadata=None), + Event(author="other_agent"), + ] + + invocation_context = self.create_invocation_context( + agent, + context_cache_config=self.cache_config, + session_events=events, + ) + + llm_request = LlmRequest( + model="gemini-2.0-flash", + contents=[ + types.Content( + role="user", + parts=[types.Part(text="Hello")], + ) + ], + ) + + async for event in self.processor.run_async( + invocation_context, llm_request + ): + pass + + # Should not set token count when no usage metadata + assert llm_request.cacheable_contents_token_count is None + + async def test_cacheable_contents_token_count_agent_filtering(self): + """Test that token count is filtered by agent name.""" + agent = LlmAgent(name="target_agent") + + events = [ + Event( + author="other_agent", + usage_metadata=types.UsageMetadata(prompt_token_count=2048), + ), + Event( + author="target_agent", + usage_metadata=types.UsageMetadata(prompt_token_count=1024), + ), + ] + + invocation_context = self.create_invocation_context( + agent, + context_cache_config=self.cache_config, + session_events=events, + ) + + llm_request = LlmRequest( + model="gemini-2.0-flash", + contents=[ + types.Content( + role="user", + parts=[types.Part(text="Hello")], + ) + ], + ) + + async for event in self.processor.run_async( + invocation_context, llm_request + ): + pass + + # Should use target_agent's token count, not other_agent's + assert llm_request.cacheable_contents_token_count == 1024 + + async def test_cacheable_contents_token_count_latest_selected(self): + """Test that the most recent token count is selected.""" + agent = LlmAgent(name="test_agent") + + events = [ + Event( + author="test_agent", + usage_metadata=types.UsageMetadata(prompt_token_count=512), + ), + Event( + author="test_agent", + usage_metadata=types.UsageMetadata(prompt_token_count=1024), + ), + ] + + invocation_context = self.create_invocation_context( + agent, + context_cache_config=self.cache_config, + session_events=events, + ) + + llm_request = LlmRequest( + model="gemini-2.0-flash", + contents=[ + types.Content( + role="user", + parts=[types.Part(text="Hello")], + ) + ], + ) + + async for event in self.processor.run_async( + invocation_context, llm_request + ): + pass + + # Should use the latest (most recent) token count + assert llm_request.cacheable_contents_token_count == 1024 + + async def test_cache_metadata_and_token_count_both_found(self): + """Test that both cache metadata and token count are found in single pass.""" + agent = LlmAgent(name="test_agent") + cache_metadata = self.create_cache_metadata(invocations_used=5) + + events = [ + Event( + author="test_agent", + cache_metadata=cache_metadata, + usage_metadata=types.UsageMetadata(prompt_token_count=1024), + invocation_id="previous_invocation", + ), + ] + + invocation_context = self.create_invocation_context( + agent, + context_cache_config=self.cache_config, + session_events=events, + invocation_id="current_invocation", + ) + + llm_request = LlmRequest( + model="gemini-2.0-flash", + contents=[ + types.Content( + role="user", + parts=[types.Part(text="Hello")], + ) + ], + ) + + async for event in self.processor.run_async( + invocation_context, llm_request + ): + pass + + # Should find both cache metadata and token count + assert llm_request.cache_metadata is not None + assert llm_request.cache_metadata.invocations_used == 6 # 5 + 1 + assert llm_request.cacheable_contents_token_count == 1024