mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
Xiang (Sean) ZhouandCopybara-Service
c5b976b306
before this change, we estimate the token count of the contents to cache and use it to compare with the threshold user set. but that's not precise , so we use the actual prompt token count of previous llm request. We won't create cache for the very initial request PiperOrigin-RevId: 814484840
592 lines
21 KiB
Python
592 lines
21 KiB
Python
# Copyright 2025 Google LLC
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
"""Tests for GeminiContextCacheManager."""
|
|
|
|
import time
|
|
from unittest.mock import AsyncMock
|
|
from unittest.mock import MagicMock
|
|
from unittest.mock import patch
|
|
|
|
from google.adk.agents.context_cache_config import ContextCacheConfig
|
|
from google.adk.models.cache_metadata import CacheMetadata
|
|
from google.adk.models.gemini_context_cache_manager import GeminiContextCacheManager
|
|
from google.adk.models.llm_request import LlmRequest
|
|
from google.adk.models.llm_response import LlmResponse
|
|
from google.genai import Client
|
|
from google.genai import types
|
|
import pytest
|
|
|
|
|
|
class TestGeminiContextCacheManager:
|
|
"""Test suite for GeminiContextCacheManager."""
|
|
|
|
def setup_method(self):
|
|
"""Set up test fixtures."""
|
|
mock_client = AsyncMock(spec=Client)
|
|
self.manager = GeminiContextCacheManager(mock_client)
|
|
self.cache_config = ContextCacheConfig(
|
|
cache_intervals=10,
|
|
ttl_seconds=1800,
|
|
min_tokens=0, # Allow caching for tests
|
|
)
|
|
|
|
def create_llm_request(self, cache_metadata=None, contents_count=3):
|
|
"""Helper to create test LlmRequest."""
|
|
contents = []
|
|
for i in range(contents_count):
|
|
contents.append(
|
|
types.Content(
|
|
role="user", parts=[types.Part(text=f"Test message {i}")]
|
|
)
|
|
)
|
|
|
|
# Create tools for testing fingerprinting
|
|
tools = [
|
|
types.Tool(
|
|
function_declarations=[
|
|
types.FunctionDeclaration(
|
|
name="test_tool",
|
|
description="A test tool",
|
|
parameters=types.Schema(
|
|
type=types.Type.OBJECT,
|
|
properties={
|
|
"param": types.Schema(type=types.Type.STRING)
|
|
},
|
|
),
|
|
)
|
|
]
|
|
)
|
|
]
|
|
|
|
tool_config = types.ToolConfig(
|
|
function_calling_config=types.FunctionCallingConfig(mode="AUTO")
|
|
)
|
|
|
|
return LlmRequest(
|
|
model="gemini-2.0-flash",
|
|
contents=contents,
|
|
config=types.GenerateContentConfig(
|
|
system_instruction="Test instruction",
|
|
tools=tools,
|
|
tool_config=tool_config,
|
|
),
|
|
cache_config=self.cache_config,
|
|
cache_metadata=cache_metadata,
|
|
)
|
|
|
|
def create_cache_metadata(
|
|
self, invocations_used=0, expired=False, cached_contents_count=3
|
|
):
|
|
"""Helper to create test CacheMetadata."""
|
|
current_time = time.time()
|
|
expire_time = current_time - 300 if expired else current_time + 1800
|
|
|
|
return CacheMetadata(
|
|
cache_name="projects/test/locations/us-central1/cachedContents/test123",
|
|
expire_time=expire_time,
|
|
fingerprint="test_fingerprint",
|
|
invocations_used=invocations_used,
|
|
cached_contents_count=cached_contents_count,
|
|
created_at=current_time - 600,
|
|
)
|
|
|
|
def test_init(self):
|
|
"""Test manager initialization."""
|
|
mock_client = MagicMock(spec=Client)
|
|
manager = GeminiContextCacheManager(mock_client)
|
|
assert manager is not None
|
|
assert manager.genai_client == mock_client
|
|
|
|
async def test_handle_context_caching_new_cache(self):
|
|
"""Test handling context caching with no existing cache."""
|
|
# Setup mocks
|
|
mock_cached_content = AsyncMock()
|
|
mock_cached_content.name = (
|
|
"projects/test/locations/us-central1/cachedContents/new123"
|
|
)
|
|
self.manager.genai_client.aio.caches.create = AsyncMock(
|
|
return_value=mock_cached_content
|
|
)
|
|
|
|
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(
|
|
self.manager, "_generate_cache_fingerprint", return_value="test_fp"
|
|
):
|
|
result = await self.manager.handle_context_caching(llm_request)
|
|
|
|
end_time = time.time()
|
|
|
|
assert result is not None
|
|
# Verify new cache metadata is created with fresh values
|
|
assert (
|
|
result.cache_name
|
|
== "projects/test/locations/us-central1/cachedContents/new123"
|
|
)
|
|
assert result.invocations_used == 1 # New cache starts with 1 invocation
|
|
assert result.fingerprint == "test_fp"
|
|
|
|
# Verify timestamps are recent (within test execution time)
|
|
assert start_time <= result.created_at <= end_time
|
|
assert result.expire_time > time.time() # Should be in the future
|
|
|
|
# Verify cache creation was called
|
|
self.manager.genai_client.aio.caches.create.assert_called_once()
|
|
|
|
async def test_handle_context_caching_valid_existing_cache(self):
|
|
"""Test handling context caching with valid existing cache."""
|
|
|
|
# Create request with existing valid cache
|
|
existing_cache = self.create_cache_metadata(invocations_used=5)
|
|
llm_request = self.create_llm_request(cache_metadata=existing_cache)
|
|
|
|
with patch.object(self.manager, "_is_cache_valid", return_value=True):
|
|
result = await self.manager.handle_context_caching(llm_request)
|
|
|
|
assert result is not None
|
|
# Verify that existing cache metadata is preserved (copied)
|
|
assert result.cache_name == existing_cache.cache_name
|
|
assert (
|
|
result.invocations_used == existing_cache.invocations_used
|
|
) # Should preserve original invocations_used
|
|
assert (
|
|
result.expire_time == existing_cache.expire_time
|
|
) # Should preserve original expire_time
|
|
assert (
|
|
result.fingerprint == existing_cache.fingerprint
|
|
) # Should preserve original fingerprint
|
|
assert (
|
|
result.created_at == existing_cache.created_at
|
|
) # Should preserve original created_at
|
|
|
|
# Verify it's a copy, not the same object
|
|
assert result is not existing_cache
|
|
|
|
# Should not create new cache
|
|
self.manager.genai_client.aio.caches.create.assert_not_called()
|
|
|
|
async def test_handle_context_caching_invalid_existing_cache(self):
|
|
"""Test handling context caching with invalid existing cache."""
|
|
# Setup mocks
|
|
mock_cached_content = AsyncMock()
|
|
mock_cached_content.name = (
|
|
"projects/test/locations/us-central1/cachedContents/new456"
|
|
)
|
|
self.manager.genai_client.aio.caches.create = AsyncMock(
|
|
return_value=mock_cached_content
|
|
)
|
|
|
|
# Create request with invalid existing cache
|
|
existing_cache = self.create_cache_metadata(
|
|
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),
|
|
patch.object(self.manager, "cleanup_cache") as mock_cleanup,
|
|
patch.object(
|
|
self.manager, "_generate_cache_fingerprint", return_value="new_fp"
|
|
),
|
|
):
|
|
|
|
result = await self.manager.handle_context_caching(llm_request)
|
|
|
|
assert result is not None
|
|
assert (
|
|
result.cache_name
|
|
== "projects/test/locations/us-central1/cachedContents/new456"
|
|
)
|
|
mock_cleanup.assert_called_once_with(existing_cache.cache_name)
|
|
self.manager.genai_client.aio.caches.create.assert_called_once()
|
|
|
|
async def test_is_cache_valid_fingerprint_mismatch(self):
|
|
"""Test cache validation with fingerprint mismatch."""
|
|
cache_metadata = self.create_cache_metadata()
|
|
llm_request = self.create_llm_request(cache_metadata=cache_metadata)
|
|
|
|
with patch.object(
|
|
self.manager,
|
|
"_generate_cache_fingerprint",
|
|
return_value="different_fingerprint",
|
|
):
|
|
result = await self.manager._is_cache_valid(llm_request)
|
|
|
|
assert result is False
|
|
|
|
async def test_is_cache_valid_expired_cache(self):
|
|
"""Test cache validation with expired cache."""
|
|
cache_metadata = self.create_cache_metadata(expired=True)
|
|
llm_request = self.create_llm_request(cache_metadata=cache_metadata)
|
|
|
|
with patch.object(
|
|
self.manager,
|
|
"_generate_cache_fingerprint",
|
|
return_value="test_fingerprint",
|
|
):
|
|
result = await self.manager._is_cache_valid(llm_request)
|
|
|
|
assert result is False
|
|
|
|
async def test_is_cache_valid_cache_intervals_exceeded(self):
|
|
"""Test cache validation with max invocations exceeded."""
|
|
cache_metadata = self.create_cache_metadata(
|
|
invocations_used=15
|
|
) # Exceeds cache_intervals=10
|
|
llm_request = self.create_llm_request(cache_metadata=cache_metadata)
|
|
|
|
with patch.object(
|
|
self.manager,
|
|
"_generate_cache_fingerprint",
|
|
return_value="test_fingerprint",
|
|
):
|
|
result = await self.manager._is_cache_valid(llm_request)
|
|
|
|
assert result is False
|
|
|
|
async def test_is_cache_valid_all_checks_pass(self):
|
|
"""Test cache validation when all checks pass."""
|
|
cache_metadata = self.create_cache_metadata(
|
|
invocations_used=5
|
|
) # Within cache_intervals=10
|
|
llm_request = self.create_llm_request(cache_metadata=cache_metadata)
|
|
|
|
with patch.object(
|
|
self.manager,
|
|
"_generate_cache_fingerprint",
|
|
return_value="test_fingerprint",
|
|
):
|
|
result = await self.manager._is_cache_valid(llm_request)
|
|
|
|
assert result is True
|
|
|
|
async def test_cleanup_cache(self):
|
|
"""Test cache cleanup functionality."""
|
|
cache_name = "projects/test/locations/us-central1/cachedContents/test123"
|
|
|
|
await self.manager.cleanup_cache(cache_name)
|
|
|
|
self.manager.genai_client.aio.caches.delete.assert_called_once_with(
|
|
name=cache_name
|
|
)
|
|
|
|
def test_generate_cache_fingerprint(self):
|
|
"""Test cache fingerprint generation includes tools and tool_config."""
|
|
llm_request = self.create_llm_request()
|
|
cache_contents_count = 2 # Cache all but last content
|
|
|
|
fingerprint1 = self.manager._generate_cache_fingerprint(
|
|
llm_request, cache_contents_count
|
|
)
|
|
fingerprint2 = self.manager._generate_cache_fingerprint(
|
|
llm_request, cache_contents_count
|
|
)
|
|
|
|
# Same request should generate same fingerprint
|
|
assert fingerprint1 == fingerprint2
|
|
assert isinstance(fingerprint1, str)
|
|
assert len(fingerprint1) > 0
|
|
|
|
# Test that tool_config and tools are included in fingerprint
|
|
# Create request without tools/tool_config
|
|
llm_request_no_tools = LlmRequest(
|
|
model="gemini-2.0-flash",
|
|
contents=[types.Content(role="user", parts=[types.Part(text="Test")])],
|
|
config=types.GenerateContentConfig(
|
|
system_instruction="Test instruction"
|
|
),
|
|
cache_config=self.cache_config,
|
|
)
|
|
|
|
fingerprint_no_tools = self.manager._generate_cache_fingerprint(
|
|
llm_request_no_tools, cache_contents_count
|
|
)
|
|
|
|
# Should be different from request with tools
|
|
assert fingerprint1 != fingerprint_no_tools
|
|
|
|
def test_generate_cache_fingerprint_different_requests(self):
|
|
"""Test that different requests generate different fingerprints."""
|
|
llm_request1 = self.create_llm_request()
|
|
|
|
llm_request2 = LlmRequest(
|
|
model="gemini-2.0-flash",
|
|
contents=[
|
|
types.Content(
|
|
role="user", parts=[types.Part(text="Different message")]
|
|
)
|
|
],
|
|
config=types.GenerateContentConfig(
|
|
system_instruction="Different instruction"
|
|
),
|
|
cache_config=self.cache_config,
|
|
)
|
|
|
|
cache_contents_count = 2
|
|
fingerprint1 = self.manager._generate_cache_fingerprint(
|
|
llm_request1, cache_contents_count
|
|
)
|
|
fingerprint2 = self.manager._generate_cache_fingerprint(
|
|
llm_request2, cache_contents_count
|
|
)
|
|
|
|
assert fingerprint1 != fingerprint2
|
|
|
|
def test_generate_cache_fingerprint_tool_config_variations(self):
|
|
"""Test that different tool configs generate different fingerprints."""
|
|
# Request with AUTO mode
|
|
llm_request_auto = self.create_llm_request()
|
|
|
|
# Request with NONE mode
|
|
tool_config_none = types.ToolConfig(
|
|
function_calling_config=types.FunctionCallingConfig(mode="NONE")
|
|
)
|
|
|
|
llm_request_none = LlmRequest(
|
|
model="gemini-2.0-flash",
|
|
contents=[types.Content(role="user", parts=[types.Part(text="Test")])],
|
|
config=types.GenerateContentConfig(
|
|
system_instruction="Test instruction",
|
|
tools=llm_request_auto.config.tools,
|
|
tool_config=tool_config_none,
|
|
),
|
|
cache_config=self.cache_config,
|
|
)
|
|
|
|
cache_contents_count = 2
|
|
fingerprint_auto = self.manager._generate_cache_fingerprint(
|
|
llm_request_auto, cache_contents_count
|
|
)
|
|
fingerprint_none = self.manager._generate_cache_fingerprint(
|
|
llm_request_none, cache_contents_count
|
|
)
|
|
|
|
assert fingerprint_auto != fingerprint_none
|
|
|
|
async def test_populate_cache_metadata_in_response_no_invocations_increment(
|
|
self,
|
|
):
|
|
"""Test that populate_cache_metadata_in_response doesn't increment invocations_used."""
|
|
# Create mock response with usage metadata
|
|
usage_metadata = MagicMock()
|
|
usage_metadata.cached_content_token_count = 800
|
|
usage_metadata.prompt_token_count = 1000
|
|
|
|
llm_response = MagicMock(spec=LlmResponse)
|
|
llm_response.usage_metadata = usage_metadata
|
|
|
|
cache_metadata = self.create_cache_metadata(invocations_used=3)
|
|
|
|
self.manager.populate_cache_metadata_in_response(
|
|
llm_response, cache_metadata
|
|
)
|
|
|
|
# Verify response metadata preserves the original invocations_used (no increment)
|
|
updated_metadata = llm_response.cache_metadata
|
|
assert (
|
|
updated_metadata.invocations_used == 3
|
|
) # Should preserve original value
|
|
assert updated_metadata.cache_name == cache_metadata.cache_name
|
|
assert updated_metadata.fingerprint == cache_metadata.fingerprint
|
|
assert updated_metadata.expire_time == cache_metadata.expire_time
|
|
assert updated_metadata.created_at == cache_metadata.created_at
|
|
|
|
async def test_populate_cache_metadata_no_usage_metadata(self):
|
|
"""Test populating cache metadata when no usage metadata."""
|
|
llm_response = MagicMock(spec=LlmResponse)
|
|
llm_response.usage_metadata = None
|
|
|
|
cache_metadata = self.create_cache_metadata(invocations_used=3)
|
|
|
|
self.manager.populate_cache_metadata_in_response(
|
|
llm_response, cache_metadata
|
|
)
|
|
|
|
# Should still create metadata even without usage info
|
|
updated_metadata = llm_response.cache_metadata
|
|
assert (
|
|
updated_metadata.invocations_used == 3
|
|
) # Should preserve original value
|
|
assert updated_metadata.cache_name == cache_metadata.cache_name
|
|
|
|
async def test_create_new_cache_with_proper_ttl(self):
|
|
"""Test that new cache is created with proper TTL."""
|
|
mock_cached_content = AsyncMock()
|
|
mock_cached_content.name = (
|
|
"projects/test/locations/us-central1/cachedContents/test123"
|
|
)
|
|
self.manager.genai_client.aio.caches.create = AsyncMock(
|
|
return_value=mock_cached_content
|
|
)
|
|
|
|
llm_request = self.create_llm_request()
|
|
|
|
cache_contents_count = max(0, len(llm_request.contents) - 1)
|
|
|
|
with patch.object(
|
|
self.manager, "_generate_cache_fingerprint", return_value="test_fp"
|
|
):
|
|
await self.manager._create_gemini_cache(llm_request, cache_contents_count)
|
|
|
|
# Verify cache creation call includes TTL
|
|
create_call = self.manager.genai_client.aio.caches.create.call_args
|
|
assert create_call is not None
|
|
cache_config = create_call[1]["config"]
|
|
assert cache_config.ttl == "1800s" # From cache_config
|
|
|
|
def test_all_but_last_content_caching(self):
|
|
"""Test that cache content counting works correctly."""
|
|
# Test with multiple contents
|
|
llm_request_multi = self.create_llm_request(contents_count=5)
|
|
|
|
# Test cache contents count calculation
|
|
cache_contents_count = max(0, len(llm_request_multi.contents) - 1)
|
|
|
|
assert cache_contents_count == 4 # 5 contents, so cache 4 contents
|
|
|
|
# Test with single content
|
|
llm_request_single = self.create_llm_request(contents_count=1)
|
|
single_cache_contents_count = max(0, len(llm_request_single.contents) - 1)
|
|
|
|
assert single_cache_contents_count == 0 # Single content, cache 0 contents
|
|
|
|
def test_edge_cases(self):
|
|
"""Test various edge cases."""
|
|
# Test with None cache_config
|
|
llm_request_no_config = LlmRequest(
|
|
model="gemini-2.0-flash",
|
|
contents=[types.Content(role="user", parts=[types.Part(text="Test")])],
|
|
config=types.GenerateContentConfig(system_instruction="Test"),
|
|
cache_config=None,
|
|
)
|
|
|
|
# Should handle gracefully
|
|
cache_contents_count = 2
|
|
fingerprint = self.manager._generate_cache_fingerprint(
|
|
llm_request_no_config, cache_contents_count
|
|
)
|
|
assert isinstance(fingerprint, str)
|
|
|
|
# Test with empty contents
|
|
llm_request_empty = LlmRequest(
|
|
model="gemini-2.0-flash",
|
|
contents=[],
|
|
config=types.GenerateContentConfig(system_instruction="Test"),
|
|
cache_config=self.cache_config,
|
|
)
|
|
|
|
empty_cache_contents_count = 0
|
|
fingerprint = self.manager._generate_cache_fingerprint(
|
|
llm_request_empty, empty_cache_contents_count
|
|
)
|
|
assert isinstance(fingerprint, str)
|
|
|
|
def test_parameter_types_enforcement(self):
|
|
"""Test that method calls with correct parameter types work properly."""
|
|
# Create proper objects
|
|
usage_metadata = MagicMock()
|
|
usage_metadata.cached_content_token_count = 500
|
|
usage_metadata.prompt_token_count = 1000
|
|
|
|
llm_response = MagicMock(spec=LlmResponse)
|
|
llm_response.usage_metadata = usage_metadata
|
|
|
|
cache_metadata = self.create_cache_metadata(invocations_used=3)
|
|
|
|
# This should work fine (correct types and order)
|
|
self.manager.populate_cache_metadata_in_response(
|
|
llm_response, cache_metadata
|
|
)
|
|
updated_metadata = llm_response.cache_metadata
|
|
assert updated_metadata.invocations_used == 3 # No increment in this method
|
|
|
|
# Document expected types for integration tests
|
|
assert isinstance(cache_metadata, CacheMetadata)
|
|
assert hasattr(
|
|
llm_response, "usage_metadata"
|
|
) # LlmResponse should have this
|
|
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()
|