mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
feat: support context caching
1. add a context cache config in app level which will apply to all agents in the app 2. pass on cache config through invocation context to llm_reqeust 3. store cache metadata in llm_response 4. lookup old cache metadata from latest event for reusing old cache 5. create new cache if old cache cannot be reused PiperOrigin-RevId: 809158578
This commit is contained in:
Xiang (Sean) Zhou
committed by
Copybara-Service
Copybara-Service
parent
13a95c463d
commit
c66245a3b8
@@ -0,0 +1,178 @@
|
||||
# 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 ContextCacheConfig."""
|
||||
|
||||
from google.adk.agents.context_cache_config import ContextCacheConfig
|
||||
from pydantic import ValidationError
|
||||
import pytest
|
||||
|
||||
|
||||
class TestContextCacheConfig:
|
||||
"""Test suite for ContextCacheConfig."""
|
||||
|
||||
def test_default_values(self):
|
||||
"""Test that default values are set correctly."""
|
||||
config = ContextCacheConfig()
|
||||
|
||||
assert config.cache_intervals == 10
|
||||
assert config.ttl_seconds == 1800 # 30 minutes
|
||||
assert config.min_tokens == 0
|
||||
|
||||
def test_custom_values(self):
|
||||
"""Test creating config with custom values."""
|
||||
config = ContextCacheConfig(
|
||||
cache_intervals=15, ttl_seconds=3600, min_tokens=1024
|
||||
)
|
||||
|
||||
assert config.cache_intervals == 15
|
||||
assert config.ttl_seconds == 3600
|
||||
assert config.min_tokens == 1024
|
||||
|
||||
def test_cache_intervals_validation(self):
|
||||
"""Test cache_intervals validation constraints."""
|
||||
# Valid range
|
||||
config = ContextCacheConfig(cache_intervals=1)
|
||||
assert config.cache_intervals == 1
|
||||
|
||||
config = ContextCacheConfig(cache_intervals=100)
|
||||
assert config.cache_intervals == 100
|
||||
|
||||
# Invalid: too low
|
||||
with pytest.raises(ValidationError) as exc_info:
|
||||
ContextCacheConfig(cache_intervals=0)
|
||||
assert "greater than or equal to 1" in str(exc_info.value)
|
||||
|
||||
# Invalid: too high
|
||||
with pytest.raises(ValidationError) as exc_info:
|
||||
ContextCacheConfig(cache_intervals=101)
|
||||
assert "less than or equal to 100" in str(exc_info.value)
|
||||
|
||||
def test_ttl_seconds_validation(self):
|
||||
"""Test ttl_seconds validation constraints."""
|
||||
# Valid range
|
||||
config = ContextCacheConfig(ttl_seconds=1)
|
||||
assert config.ttl_seconds == 1
|
||||
|
||||
config = ContextCacheConfig(ttl_seconds=86400) # 24 hours
|
||||
assert config.ttl_seconds == 86400
|
||||
|
||||
# Invalid: zero or negative
|
||||
with pytest.raises(ValidationError) as exc_info:
|
||||
ContextCacheConfig(ttl_seconds=0)
|
||||
assert "greater than 0" in str(exc_info.value)
|
||||
|
||||
with pytest.raises(ValidationError) as exc_info:
|
||||
ContextCacheConfig(ttl_seconds=-1)
|
||||
assert "greater than 0" in str(exc_info.value)
|
||||
|
||||
def test_min_tokens_validation(self):
|
||||
"""Test min_tokens validation constraints."""
|
||||
# Valid values
|
||||
config = ContextCacheConfig(min_tokens=0)
|
||||
assert config.min_tokens == 0
|
||||
|
||||
config = ContextCacheConfig(min_tokens=1024)
|
||||
assert config.min_tokens == 1024
|
||||
|
||||
# Invalid: negative
|
||||
with pytest.raises(ValidationError) as exc_info:
|
||||
ContextCacheConfig(min_tokens=-1)
|
||||
assert "greater than or equal to 0" in str(exc_info.value)
|
||||
|
||||
def test_ttl_string_property(self):
|
||||
"""Test ttl_string property returns correct format."""
|
||||
config = ContextCacheConfig(ttl_seconds=1800)
|
||||
assert config.ttl_string == "1800s"
|
||||
|
||||
config = ContextCacheConfig(ttl_seconds=3600)
|
||||
assert config.ttl_string == "3600s"
|
||||
|
||||
def test_str_representation(self):
|
||||
"""Test string representation for logging."""
|
||||
config = ContextCacheConfig(
|
||||
cache_intervals=15, ttl_seconds=3600, min_tokens=1024
|
||||
)
|
||||
|
||||
expected = (
|
||||
"ContextCacheConfig(cache_intervals=15, ttl=3600s, min_tokens=1024)"
|
||||
)
|
||||
assert str(config) == expected
|
||||
|
||||
def test_str_representation_defaults(self):
|
||||
"""Test string representation with default values."""
|
||||
config = ContextCacheConfig()
|
||||
|
||||
expected = "ContextCacheConfig(cache_intervals=10, ttl=1800s, min_tokens=0)"
|
||||
assert str(config) == expected
|
||||
|
||||
def test_pydantic_model_validation(self):
|
||||
"""Test that Pydantic model validation works correctly."""
|
||||
# Test extra fields are forbidden
|
||||
with pytest.raises(ValidationError) as exc_info:
|
||||
ContextCacheConfig(cache_intervals=10, extra_field="not_allowed")
|
||||
assert "extra" in str(exc_info.value).lower()
|
||||
|
||||
def test_field_descriptions(self):
|
||||
"""Test that fields have proper descriptions."""
|
||||
config = ContextCacheConfig()
|
||||
schema = config.model_json_schema()
|
||||
|
||||
assert "cache_intervals" in schema["properties"]
|
||||
assert (
|
||||
"Maximum number of invocations"
|
||||
in schema["properties"]["cache_intervals"]["description"]
|
||||
)
|
||||
|
||||
assert "ttl_seconds" in schema["properties"]
|
||||
assert (
|
||||
"Time-to-live for cache"
|
||||
in schema["properties"]["ttl_seconds"]["description"]
|
||||
)
|
||||
|
||||
assert "min_tokens" in schema["properties"]
|
||||
assert (
|
||||
"Minimum estimated request tokens"
|
||||
in schema["properties"]["min_tokens"]["description"]
|
||||
)
|
||||
|
||||
def test_immutability_config(self):
|
||||
"""Test that the model config is set correctly."""
|
||||
config = ContextCacheConfig()
|
||||
assert config.model_config["extra"] == "forbid"
|
||||
|
||||
def test_realistic_scenarios(self):
|
||||
"""Test realistic configuration scenarios."""
|
||||
# Quick caching for development
|
||||
dev_config = ContextCacheConfig(
|
||||
cache_intervals=5, ttl_seconds=600, min_tokens=0 # 10 minutes
|
||||
)
|
||||
assert dev_config.cache_intervals == 5
|
||||
assert dev_config.ttl_seconds == 600
|
||||
|
||||
# Production caching
|
||||
prod_config = ContextCacheConfig(
|
||||
cache_intervals=20, ttl_seconds=7200, min_tokens=2048 # 2 hours
|
||||
)
|
||||
assert prod_config.cache_intervals == 20
|
||||
assert prod_config.ttl_seconds == 7200
|
||||
assert prod_config.min_tokens == 2048
|
||||
|
||||
# Conservative caching
|
||||
conservative_config = ContextCacheConfig(
|
||||
cache_intervals=3, ttl_seconds=300, min_tokens=4096 # 5 minutes
|
||||
)
|
||||
assert conservative_config.cache_intervals == 3
|
||||
assert conservative_config.ttl_seconds == 300
|
||||
assert conservative_config.min_tokens == 4096
|
||||
File diff suppressed because it is too large
Load Diff
@@ -15,6 +15,7 @@
|
||||
from unittest.mock import Mock
|
||||
|
||||
from google.adk.agents.base_agent import BaseAgent
|
||||
from google.adk.agents.context_cache_config import ContextCacheConfig
|
||||
from google.adk.apps.app import App
|
||||
from google.adk.plugins.base_plugin import BasePlugin
|
||||
|
||||
@@ -38,3 +39,82 @@ class TestApp:
|
||||
assert app.name == "test_app"
|
||||
assert app.root_agent == mock_agent
|
||||
assert app.plugins == [mock_plugin]
|
||||
|
||||
def test_app_initialization_without_cache_config(self):
|
||||
"""Test that the app is initialized correctly without context cache config."""
|
||||
mock_agent = Mock(spec=BaseAgent)
|
||||
app = App(name="test_app", root_agent=mock_agent)
|
||||
assert app.name == "test_app"
|
||||
assert app.root_agent == mock_agent
|
||||
assert app.context_cache_config is None
|
||||
|
||||
def test_app_initialization_with_cache_config(self):
|
||||
"""Test that the app is initialized correctly with context cache config."""
|
||||
mock_agent = Mock(spec=BaseAgent)
|
||||
cache_config = ContextCacheConfig(
|
||||
cache_intervals=15, ttl_seconds=3600, min_tokens=1024
|
||||
)
|
||||
|
||||
app = App(
|
||||
name="test_app",
|
||||
root_agent=mock_agent,
|
||||
context_cache_config=cache_config,
|
||||
)
|
||||
|
||||
assert app.name == "test_app"
|
||||
assert app.root_agent == mock_agent
|
||||
assert app.context_cache_config == cache_config
|
||||
assert app.context_cache_config.cache_intervals == 15
|
||||
assert app.context_cache_config.ttl_seconds == 3600
|
||||
assert app.context_cache_config.min_tokens == 1024
|
||||
|
||||
def test_app_with_all_components(self):
|
||||
"""Test app with all components: agent, plugins, and cache config."""
|
||||
mock_agent = Mock(spec=BaseAgent)
|
||||
mock_plugin = Mock(spec=BasePlugin)
|
||||
cache_config = ContextCacheConfig(
|
||||
cache_intervals=20, ttl_seconds=7200, min_tokens=2048
|
||||
)
|
||||
|
||||
app = App(
|
||||
name="full_test_app",
|
||||
root_agent=mock_agent,
|
||||
plugins=[mock_plugin],
|
||||
context_cache_config=cache_config,
|
||||
)
|
||||
|
||||
assert app.name == "full_test_app"
|
||||
assert app.root_agent == mock_agent
|
||||
assert app.plugins == [mock_plugin]
|
||||
assert app.context_cache_config == cache_config
|
||||
|
||||
def test_app_cache_config_defaults(self):
|
||||
"""Test that cache config has proper defaults when created."""
|
||||
mock_agent = Mock(spec=BaseAgent)
|
||||
cache_config = ContextCacheConfig() # Use defaults
|
||||
|
||||
app = App(
|
||||
name="default_cache_app",
|
||||
root_agent=mock_agent,
|
||||
context_cache_config=cache_config,
|
||||
)
|
||||
|
||||
assert app.context_cache_config.cache_intervals == 10 # Default
|
||||
assert app.context_cache_config.ttl_seconds == 1800 # Default 30 minutes
|
||||
assert app.context_cache_config.min_tokens == 0 # Default
|
||||
|
||||
def test_app_context_cache_config_is_optional(self):
|
||||
"""Test that context_cache_config is truly optional."""
|
||||
mock_agent = Mock(spec=BaseAgent)
|
||||
|
||||
# Should work without context_cache_config
|
||||
app = App(name="no_cache_app", root_agent=mock_agent)
|
||||
assert app.context_cache_config is None
|
||||
|
||||
# Should work with explicit None
|
||||
app = App(
|
||||
name="explicit_none_app",
|
||||
root_agent=mock_agent,
|
||||
context_cache_config=None,
|
||||
)
|
||||
assert app.context_cache_config is None
|
||||
|
||||
@@ -0,0 +1,454 @@
|
||||
# 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 ContextCacheRequestProcessor."""
|
||||
|
||||
import time
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from google.adk.agents.context_cache_config import ContextCacheConfig
|
||||
from google.adk.agents.invocation_context import InvocationContext
|
||||
from google.adk.agents.llm_agent import LlmAgent
|
||||
from google.adk.events.event import Event
|
||||
from google.adk.flows.llm_flows.context_cache_processor import ContextCacheRequestProcessor
|
||||
from google.adk.models.cache_metadata import CacheMetadata
|
||||
from google.adk.models.llm_request import LlmRequest
|
||||
from google.adk.sessions.base_session_service import BaseSessionService
|
||||
from google.adk.sessions.session import Session
|
||||
from google.genai import types
|
||||
import pytest
|
||||
|
||||
|
||||
class TestContextCacheRequestProcessor:
|
||||
"""Test suite for ContextCacheRequestProcessor."""
|
||||
|
||||
def setup_method(self):
|
||||
"""Set up test fixtures."""
|
||||
self.processor = ContextCacheRequestProcessor()
|
||||
self.cache_config = ContextCacheConfig(
|
||||
cache_intervals=10, ttl_seconds=1800, min_tokens=1024
|
||||
)
|
||||
|
||||
def create_invocation_context(
|
||||
self,
|
||||
agent,
|
||||
context_cache_config=None,
|
||||
session_events=None,
|
||||
invocation_id="test_invocation",
|
||||
):
|
||||
"""Helper to create InvocationContext."""
|
||||
mock_session = Session(
|
||||
id="test_session",
|
||||
app_name="test_app",
|
||||
user_id="test_user",
|
||||
events=session_events or [],
|
||||
)
|
||||
|
||||
mock_session_service = MagicMock(spec=BaseSessionService)
|
||||
|
||||
return InvocationContext(
|
||||
agent=agent,
|
||||
session=mock_session,
|
||||
session_service=mock_session_service,
|
||||
context_cache_config=context_cache_config,
|
||||
invocation_id=invocation_id,
|
||||
)
|
||||
|
||||
def create_cache_metadata(
|
||||
self, invocations_used=1, cache_name="test-cache", cached_contents_count=3
|
||||
):
|
||||
"""Helper to create CacheMetadata."""
|
||||
return CacheMetadata(
|
||||
cache_name=(
|
||||
f"projects/test/locations/us-central1/cachedContents/{cache_name}"
|
||||
),
|
||||
expire_time=time.time() + 1800,
|
||||
fingerprint="test_fingerprint",
|
||||
invocations_used=invocations_used,
|
||||
cached_contents_count=cached_contents_count,
|
||||
created_at=time.time() - 600,
|
||||
)
|
||||
|
||||
async def test_no_cache_config(self):
|
||||
"""Test processor with no cache config."""
|
||||
agent = LlmAgent(name="test_agent")
|
||||
invocation_context = self.create_invocation_context(
|
||||
agent, context_cache_config=None
|
||||
)
|
||||
|
||||
llm_request = LlmRequest(
|
||||
model="gemini-2.0-flash",
|
||||
contents=[
|
||||
types.Content(
|
||||
role="user",
|
||||
parts=[types.Part(text="Hello")],
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
# Process should complete without adding cache config
|
||||
events = []
|
||||
async for event in self.processor.run_async(
|
||||
invocation_context, llm_request
|
||||
):
|
||||
events.append(event)
|
||||
|
||||
assert len(events) == 0 # No events yielded
|
||||
assert llm_request.cache_config is None
|
||||
|
||||
async def test_with_cache_config_no_session_events(self):
|
||||
"""Test processor with cache config but no session events."""
|
||||
agent = LlmAgent(name="test_agent")
|
||||
invocation_context = self.create_invocation_context(
|
||||
agent, context_cache_config=self.cache_config
|
||||
)
|
||||
|
||||
llm_request = LlmRequest(
|
||||
model="gemini-2.0-flash",
|
||||
contents=[
|
||||
types.Content(
|
||||
role="user",
|
||||
parts=[types.Part(text="Hello")],
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
# Process should add cache config but no metadata
|
||||
events = []
|
||||
async for event in self.processor.run_async(
|
||||
invocation_context, llm_request
|
||||
):
|
||||
events.append(event)
|
||||
|
||||
assert len(events) == 0 # No events yielded
|
||||
assert llm_request.cache_config == self.cache_config
|
||||
assert llm_request.cache_metadata is None
|
||||
|
||||
async def test_with_cache_metadata_same_invocation(self):
|
||||
"""Test processor finds cache metadata from same invocation."""
|
||||
agent = LlmAgent(name="test_agent")
|
||||
cache_metadata = self.create_cache_metadata(invocations_used=5)
|
||||
|
||||
# Event with same invocation ID
|
||||
events = [
|
||||
Event(
|
||||
author="test_agent",
|
||||
cache_metadata=cache_metadata,
|
||||
invocation_id="test_invocation",
|
||||
)
|
||||
]
|
||||
|
||||
invocation_context = self.create_invocation_context(
|
||||
agent,
|
||||
context_cache_config=self.cache_config,
|
||||
session_events=events,
|
||||
invocation_id="test_invocation",
|
||||
)
|
||||
|
||||
llm_request = LlmRequest(
|
||||
model="gemini-2.0-flash",
|
||||
contents=[
|
||||
types.Content(
|
||||
role="user",
|
||||
parts=[types.Part(text="Hello")],
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
# Process should add cache config and metadata (same invocation, no increment)
|
||||
async for event in self.processor.run_async(
|
||||
invocation_context, llm_request
|
||||
):
|
||||
pass
|
||||
|
||||
assert llm_request.cache_config == self.cache_config
|
||||
assert llm_request.cache_metadata == cache_metadata
|
||||
assert llm_request.cache_metadata.invocations_used == 5 # No increment
|
||||
|
||||
async def test_with_cache_metadata_different_invocation(self):
|
||||
"""Test processor finds cache metadata from different invocation."""
|
||||
agent = LlmAgent(name="test_agent")
|
||||
cache_metadata = self.create_cache_metadata(invocations_used=5)
|
||||
|
||||
# Event with different invocation ID
|
||||
events = [
|
||||
Event(
|
||||
author="test_agent",
|
||||
cache_metadata=cache_metadata,
|
||||
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")],
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
# Process should add cache config and increment invocations_used
|
||||
async for event in self.processor.run_async(
|
||||
invocation_context, llm_request
|
||||
):
|
||||
pass
|
||||
|
||||
assert llm_request.cache_config == self.cache_config
|
||||
assert llm_request.cache_metadata is not None
|
||||
assert llm_request.cache_metadata.invocations_used == 6 # Incremented
|
||||
|
||||
async def test_cache_metadata_agent_filtering(self):
|
||||
"""Test that cache metadata is filtered by agent name."""
|
||||
agent = LlmAgent(name="target_agent")
|
||||
target_cache = self.create_cache_metadata(
|
||||
invocations_used=3, cache_name="target"
|
||||
)
|
||||
other_cache = self.create_cache_metadata(
|
||||
invocations_used=7, cache_name="other"
|
||||
)
|
||||
|
||||
events = [
|
||||
Event(
|
||||
author="other_agent",
|
||||
cache_metadata=other_cache,
|
||||
invocation_id="other_invocation",
|
||||
),
|
||||
Event(
|
||||
author="target_agent",
|
||||
cache_metadata=target_cache,
|
||||
invocation_id="target_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")],
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
# Should only use target_agent's cache metadata
|
||||
async for event in self.processor.run_async(
|
||||
invocation_context, llm_request
|
||||
):
|
||||
pass
|
||||
|
||||
assert llm_request.cache_metadata is not None
|
||||
assert llm_request.cache_metadata.cache_name == target_cache.cache_name
|
||||
assert llm_request.cache_metadata.invocations_used == 4 # target_cache + 1
|
||||
|
||||
async def test_latest_cache_metadata_selected(self):
|
||||
"""Test that the latest cache metadata is selected."""
|
||||
agent = LlmAgent(name="test_agent")
|
||||
older_cache = self.create_cache_metadata(
|
||||
invocations_used=2, cache_name="older"
|
||||
)
|
||||
newer_cache = self.create_cache_metadata(
|
||||
invocations_used=5, cache_name="newer"
|
||||
)
|
||||
|
||||
# Events in chronological order (older first)
|
||||
events = [
|
||||
Event(
|
||||
author="test_agent",
|
||||
cache_metadata=older_cache,
|
||||
invocation_id="older_invocation",
|
||||
),
|
||||
Event(
|
||||
author="test_agent",
|
||||
cache_metadata=newer_cache,
|
||||
invocation_id="newer_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")],
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
# Should use the newer (latest) cache metadata
|
||||
async for event in self.processor.run_async(
|
||||
invocation_context, llm_request
|
||||
):
|
||||
pass
|
||||
|
||||
assert llm_request.cache_metadata is not None
|
||||
assert llm_request.cache_metadata.cache_name == newer_cache.cache_name
|
||||
assert llm_request.cache_metadata.invocations_used == 6 # newer_cache + 1
|
||||
|
||||
async def test_no_cache_metadata_events(self):
|
||||
"""Test when session has events but no cache metadata."""
|
||||
agent = LlmAgent(name="test_agent")
|
||||
|
||||
events = [
|
||||
Event(author="test_agent", cache_metadata=None),
|
||||
Event(author="other_agent", cache_metadata=None),
|
||||
]
|
||||
|
||||
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")],
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
# Should add cache config but no metadata
|
||||
async for event in self.processor.run_async(
|
||||
invocation_context, llm_request
|
||||
):
|
||||
pass
|
||||
|
||||
assert llm_request.cache_config == self.cache_config
|
||||
assert llm_request.cache_metadata is None
|
||||
|
||||
async def test_empty_session(self):
|
||||
"""Test with empty session."""
|
||||
agent = LlmAgent(name="test_agent")
|
||||
|
||||
invocation_context = self.create_invocation_context(
|
||||
agent,
|
||||
context_cache_config=self.cache_config,
|
||||
session_events=[],
|
||||
)
|
||||
|
||||
llm_request = LlmRequest(
|
||||
model="gemini-2.0-flash",
|
||||
contents=[
|
||||
types.Content(
|
||||
role="user",
|
||||
parts=[types.Part(text="Hello")],
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
# Should add cache config but no metadata
|
||||
async for event in self.processor.run_async(
|
||||
invocation_context, llm_request
|
||||
):
|
||||
pass
|
||||
|
||||
assert llm_request.cache_config == self.cache_config
|
||||
assert llm_request.cache_metadata is None
|
||||
|
||||
async def test_processor_yields_no_events(self):
|
||||
"""Test that processor yields no events."""
|
||||
agent = LlmAgent(name="test_agent")
|
||||
|
||||
invocation_context = self.create_invocation_context(
|
||||
agent, context_cache_config=self.cache_config
|
||||
)
|
||||
|
||||
llm_request = LlmRequest(
|
||||
model="gemini-2.0-flash",
|
||||
contents=[
|
||||
types.Content(
|
||||
role="user",
|
||||
parts=[types.Part(text="Hello")],
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
events = []
|
||||
async for event in self.processor.run_async(
|
||||
invocation_context, llm_request
|
||||
):
|
||||
events.append(event)
|
||||
|
||||
# Processor should never yield events
|
||||
assert len(events) == 0
|
||||
|
||||
async def test_mixed_events_scenario(self):
|
||||
"""Test complex scenario with mixed events."""
|
||||
agent = LlmAgent(name="test_agent")
|
||||
cache_metadata = self.create_cache_metadata(invocations_used=10)
|
||||
|
||||
events = [
|
||||
Event(author="other_agent", cache_metadata=None),
|
||||
Event(author="test_agent", cache_metadata=None), # No cache metadata
|
||||
Event(
|
||||
author="different_agent", cache_metadata=cache_metadata
|
||||
), # Wrong agent
|
||||
Event(
|
||||
author="test_agent",
|
||||
cache_metadata=cache_metadata,
|
||||
invocation_id="prev",
|
||||
),
|
||||
]
|
||||
|
||||
invocation_context = self.create_invocation_context(
|
||||
agent,
|
||||
context_cache_config=self.cache_config,
|
||||
session_events=events,
|
||||
invocation_id="current",
|
||||
)
|
||||
|
||||
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 the test_agent's cache metadata and increment it
|
||||
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
|
||||
@@ -0,0 +1,314 @@
|
||||
# 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 CacheMetadata."""
|
||||
|
||||
import time
|
||||
|
||||
from google.adk.models.cache_metadata import CacheMetadata
|
||||
from pydantic import ValidationError
|
||||
import pytest
|
||||
|
||||
|
||||
class TestCacheMetadata:
|
||||
"""Test suite for CacheMetadata."""
|
||||
|
||||
def test_required_fields(self):
|
||||
"""Test that all required fields must be provided."""
|
||||
# Valid creation with all required fields
|
||||
metadata = CacheMetadata(
|
||||
cache_name="projects/123/locations/us-central1/cachedContents/456",
|
||||
expire_time=time.time() + 1800,
|
||||
fingerprint="abc123",
|
||||
invocations_used=5,
|
||||
cached_contents_count=3,
|
||||
)
|
||||
|
||||
assert (
|
||||
metadata.cache_name
|
||||
== "projects/123/locations/us-central1/cachedContents/456"
|
||||
)
|
||||
assert metadata.expire_time > time.time()
|
||||
assert metadata.fingerprint == "abc123"
|
||||
assert metadata.invocations_used == 5
|
||||
assert metadata.cached_contents_count == 3
|
||||
assert metadata.created_at is None # Optional field
|
||||
|
||||
def test_optional_created_at(self):
|
||||
"""Test that created_at is optional."""
|
||||
current_time = time.time()
|
||||
|
||||
metadata = CacheMetadata(
|
||||
cache_name="projects/123/locations/us-central1/cachedContents/456",
|
||||
expire_time=time.time() + 1800,
|
||||
fingerprint="abc123",
|
||||
invocations_used=3,
|
||||
cached_contents_count=2,
|
||||
created_at=current_time,
|
||||
)
|
||||
|
||||
assert metadata.created_at == current_time
|
||||
|
||||
def test_invocations_used_validation(self):
|
||||
"""Test invocations_used validation constraints."""
|
||||
# Valid: zero or positive
|
||||
metadata = CacheMetadata(
|
||||
cache_name="projects/123/locations/us-central1/cachedContents/456",
|
||||
expire_time=time.time() + 1800,
|
||||
fingerprint="abc123",
|
||||
invocations_used=0,
|
||||
cached_contents_count=1,
|
||||
)
|
||||
assert metadata.invocations_used == 0
|
||||
|
||||
metadata = CacheMetadata(
|
||||
cache_name="projects/123/locations/us-central1/cachedContents/456",
|
||||
expire_time=time.time() + 1800,
|
||||
fingerprint="abc123",
|
||||
invocations_used=10,
|
||||
cached_contents_count=1,
|
||||
)
|
||||
assert metadata.invocations_used == 10
|
||||
|
||||
# Invalid: negative
|
||||
with pytest.raises(ValidationError) as exc_info:
|
||||
CacheMetadata(
|
||||
cache_name="projects/123/locations/us-central1/cachedContents/456",
|
||||
expire_time=time.time() + 1800,
|
||||
fingerprint="abc123",
|
||||
invocations_used=-1,
|
||||
cached_contents_count=1,
|
||||
)
|
||||
assert "greater than or equal to 0" in str(exc_info.value)
|
||||
|
||||
def test_cached_contents_count_validation(self):
|
||||
"""Test cached_contents_count validation constraints."""
|
||||
# Valid: zero or positive
|
||||
metadata = CacheMetadata(
|
||||
cache_name="projects/123/locations/us-central1/cachedContents/456",
|
||||
expire_time=time.time() + 1800,
|
||||
fingerprint="abc123",
|
||||
invocations_used=1,
|
||||
cached_contents_count=0,
|
||||
)
|
||||
assert metadata.cached_contents_count == 0
|
||||
|
||||
metadata = CacheMetadata(
|
||||
cache_name="projects/123/locations/us-central1/cachedContents/456",
|
||||
expire_time=time.time() + 1800,
|
||||
fingerprint="abc123",
|
||||
invocations_used=1,
|
||||
cached_contents_count=10,
|
||||
)
|
||||
assert metadata.cached_contents_count == 10
|
||||
|
||||
# Invalid: negative
|
||||
with pytest.raises(ValidationError) as exc_info:
|
||||
CacheMetadata(
|
||||
cache_name="projects/123/locations/us-central1/cachedContents/456",
|
||||
expire_time=time.time() + 1800,
|
||||
fingerprint="abc123",
|
||||
invocations_used=1,
|
||||
cached_contents_count=-1,
|
||||
)
|
||||
assert "greater than or equal to 0" in str(exc_info.value)
|
||||
|
||||
def test_expire_soon_property(self):
|
||||
"""Test expire_soon property."""
|
||||
# Cache that expires in 10 minutes (should not expire soon)
|
||||
future_time = time.time() + 600 # 10 minutes
|
||||
metadata = CacheMetadata(
|
||||
cache_name="projects/123/locations/us-central1/cachedContents/456",
|
||||
expire_time=future_time,
|
||||
fingerprint="abc123",
|
||||
invocations_used=1,
|
||||
cached_contents_count=1,
|
||||
)
|
||||
assert not metadata.expire_soon
|
||||
|
||||
# Cache that expires in 1 minute (should expire soon)
|
||||
soon_time = time.time() + 60 # 1 minute
|
||||
metadata = CacheMetadata(
|
||||
cache_name="projects/123/locations/us-central1/cachedContents/456",
|
||||
expire_time=soon_time,
|
||||
fingerprint="abc123",
|
||||
invocations_used=1,
|
||||
cached_contents_count=1,
|
||||
)
|
||||
assert metadata.expire_soon
|
||||
|
||||
def test_str_representation(self):
|
||||
"""Test string representation."""
|
||||
current_time = time.time()
|
||||
expire_time = current_time + 1800 # 30 minutes
|
||||
|
||||
metadata = CacheMetadata(
|
||||
cache_name="projects/123/locations/us-central1/cachedContents/test456",
|
||||
expire_time=expire_time,
|
||||
fingerprint="abc123",
|
||||
invocations_used=7,
|
||||
cached_contents_count=4,
|
||||
)
|
||||
|
||||
str_repr = str(metadata)
|
||||
assert "test456" in str_repr # Cache ID
|
||||
assert "used 7 invocations" in str_repr
|
||||
assert "cached 4 contents" in str_repr
|
||||
assert "expires in" in str_repr
|
||||
|
||||
def test_immutability(self):
|
||||
"""Test that CacheMetadata is immutable (frozen)."""
|
||||
metadata = CacheMetadata(
|
||||
cache_name="projects/123/locations/us-central1/cachedContents/456",
|
||||
expire_time=time.time() + 1800,
|
||||
fingerprint="abc123",
|
||||
invocations_used=5,
|
||||
cached_contents_count=3,
|
||||
)
|
||||
|
||||
# Should not be able to modify fields
|
||||
with pytest.raises(ValidationError):
|
||||
metadata.invocations_used = 10
|
||||
|
||||
def test_model_config(self):
|
||||
"""Test that model config is set correctly."""
|
||||
metadata = CacheMetadata(
|
||||
cache_name="projects/123/locations/us-central1/cachedContents/456",
|
||||
expire_time=time.time() + 1800,
|
||||
fingerprint="abc123",
|
||||
invocations_used=5,
|
||||
cached_contents_count=3,
|
||||
)
|
||||
|
||||
assert metadata.model_config["extra"] == "forbid"
|
||||
assert metadata.model_config["frozen"] == True
|
||||
|
||||
def test_field_descriptions(self):
|
||||
"""Test that fields have proper descriptions."""
|
||||
metadata = CacheMetadata(
|
||||
cache_name="projects/123/locations/us-central1/cachedContents/456",
|
||||
expire_time=time.time() + 1800,
|
||||
fingerprint="abc123",
|
||||
invocations_used=5,
|
||||
cached_contents_count=3,
|
||||
)
|
||||
schema = metadata.model_json_schema()
|
||||
|
||||
assert "invocations_used" in schema["properties"]
|
||||
assert (
|
||||
"Number of invocations"
|
||||
in schema["properties"]["invocations_used"]["description"]
|
||||
)
|
||||
|
||||
assert "cached_contents_count" in schema["properties"]
|
||||
assert (
|
||||
"Number of contents"
|
||||
in schema["properties"]["cached_contents_count"]["description"]
|
||||
)
|
||||
|
||||
def test_realistic_cache_scenarios(self):
|
||||
"""Test realistic cache scenarios."""
|
||||
current_time = time.time()
|
||||
|
||||
# Fresh cache
|
||||
fresh_cache = CacheMetadata(
|
||||
cache_name="projects/123/locations/us-central1/cachedContents/fresh123",
|
||||
expire_time=current_time + 1800,
|
||||
fingerprint="fresh_fingerprint",
|
||||
invocations_used=1,
|
||||
cached_contents_count=5,
|
||||
created_at=current_time,
|
||||
)
|
||||
assert fresh_cache.invocations_used == 1
|
||||
assert not fresh_cache.expire_soon
|
||||
|
||||
# Well-used cache
|
||||
used_cache = CacheMetadata(
|
||||
cache_name="projects/123/locations/us-central1/cachedContents/used456",
|
||||
expire_time=current_time + 600,
|
||||
fingerprint="used_fingerprint",
|
||||
invocations_used=8,
|
||||
cached_contents_count=3,
|
||||
created_at=current_time - 1200,
|
||||
)
|
||||
assert used_cache.invocations_used == 8
|
||||
|
||||
# Expiring cache
|
||||
expiring_cache = CacheMetadata(
|
||||
cache_name=(
|
||||
"projects/123/locations/us-central1/cachedContents/expiring789"
|
||||
),
|
||||
expire_time=current_time + 60, # 1 minute
|
||||
fingerprint="expiring_fingerprint",
|
||||
invocations_used=15,
|
||||
cached_contents_count=10,
|
||||
)
|
||||
assert expiring_cache.expire_soon
|
||||
|
||||
def test_cache_name_extraction(self):
|
||||
"""Test cache name ID extraction in string representation."""
|
||||
metadata = CacheMetadata(
|
||||
cache_name=(
|
||||
"projects/123/locations/us-central1/cachedContents/extracted_id"
|
||||
),
|
||||
expire_time=time.time() + 1800,
|
||||
fingerprint="abc123",
|
||||
invocations_used=1,
|
||||
cached_contents_count=2,
|
||||
)
|
||||
|
||||
str_repr = str(metadata)
|
||||
assert "extracted_id" in str_repr
|
||||
|
||||
def test_no_performance_metrics(self):
|
||||
"""Test that performance metrics are not in CacheMetadata."""
|
||||
metadata = CacheMetadata(
|
||||
cache_name="projects/123/locations/us-central1/cachedContents/456",
|
||||
expire_time=time.time() + 1800,
|
||||
fingerprint="abc123",
|
||||
invocations_used=5,
|
||||
cached_contents_count=3,
|
||||
)
|
||||
|
||||
# Verify that token counts are NOT in CacheMetadata
|
||||
# (they should be in LlmResponse.usage_metadata)
|
||||
assert not hasattr(metadata, "cached_tokens")
|
||||
assert not hasattr(metadata, "total_tokens")
|
||||
assert not hasattr(metadata, "prompt_tokens")
|
||||
|
||||
def test_missing_required_fields(self):
|
||||
"""Test validation when required fields are missing."""
|
||||
# Test each required field
|
||||
required_fields = [
|
||||
"cache_name",
|
||||
"expire_time",
|
||||
"fingerprint",
|
||||
"invocations_used",
|
||||
"cached_contents_count",
|
||||
]
|
||||
|
||||
base_args = {
|
||||
"cache_name": "projects/123/locations/us-central1/cachedContents/456",
|
||||
"expire_time": time.time() + 1800,
|
||||
"fingerprint": "abc123",
|
||||
"invocations_used": 1,
|
||||
"cached_contents_count": 2,
|
||||
}
|
||||
|
||||
for field in required_fields:
|
||||
args = base_args.copy()
|
||||
del args[field]
|
||||
|
||||
with pytest.raises(ValidationError):
|
||||
CacheMetadata(**args)
|
||||
@@ -19,6 +19,8 @@ from unittest import mock
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
from google.adk import version as adk_version
|
||||
from google.adk.agents.context_cache_config import ContextCacheConfig
|
||||
from google.adk.models.cache_metadata import CacheMetadata
|
||||
from google.adk.models.gemini_llm_connection import GeminiLlmConnection
|
||||
from google.adk.models.google_llm import _AGENT_ENGINE_TELEMETRY_ENV_VARIABLE_NAME
|
||||
from google.adk.models.google_llm import _AGENT_ENGINE_TELEMETRY_TAG
|
||||
@@ -84,6 +86,37 @@ def llm_request():
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def cache_metadata():
|
||||
import time
|
||||
|
||||
return CacheMetadata(
|
||||
cache_name="projects/test/locations/us-central1/cachedContents/test123",
|
||||
expire_time=time.time() + 3600,
|
||||
fingerprint="test_fingerprint",
|
||||
invocations_used=2,
|
||||
cached_contents_count=3,
|
||||
created_at=time.time() - 600,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def llm_request_with_cache(cache_metadata):
|
||||
return LlmRequest(
|
||||
model="gemini-1.5-flash",
|
||||
contents=[Content(role="user", parts=[Part.from_text(text="Hello")])],
|
||||
config=types.GenerateContentConfig(
|
||||
temperature=0.1,
|
||||
response_modalities=[types.Modality.TEXT],
|
||||
system_instruction="You are a helpful assistant",
|
||||
),
|
||||
cache_config=ContextCacheConfig(
|
||||
cache_intervals=10, ttl_seconds=3600, min_tokens=100
|
||||
),
|
||||
cache_metadata=cache_metadata,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def llm_request_with_computer_use():
|
||||
return LlmRequest(
|
||||
@@ -1600,3 +1633,96 @@ async def test_adapt_computer_use_tool_no_wait():
|
||||
# Verify tools_dict is unchanged
|
||||
assert llm_request.tools_dict == original_tools_dict
|
||||
assert "wait_5_seconds" not in llm_request.tools_dict
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generate_content_async_with_cache_metadata_integration(
|
||||
gemini_llm, llm_request_with_cache, cache_metadata
|
||||
):
|
||||
"""Test integration between Google LLM and cache manager with proper parameter order.
|
||||
|
||||
This test specifically validates that the cache manager's populate_cache_metadata_in_response
|
||||
method is called with the correct parameter order: (llm_response, cache_metadata).
|
||||
|
||||
This test would have caught the parameter order bug where cache_metadata and llm_response
|
||||
were passed in the wrong order, causing 'CacheMetadata' object has no attribute 'usage_metadata' errors.
|
||||
"""
|
||||
|
||||
# Create a mock response with usage metadata including cached tokens
|
||||
generate_content_response = types.GenerateContentResponse(
|
||||
candidates=[
|
||||
types.Candidate(
|
||||
content=Content(
|
||||
role="model",
|
||||
parts=[Part.from_text(text="Hello, how can I help you?")],
|
||||
),
|
||||
finish_reason=types.FinishReason.STOP,
|
||||
)
|
||||
],
|
||||
usage_metadata=types.GenerateContentResponseUsageMetadata(
|
||||
prompt_token_count=1500,
|
||||
candidates_token_count=150,
|
||||
cached_content_token_count=800, # This is the key field that was always 0 due to the bug
|
||||
total_token_count=1650,
|
||||
),
|
||||
)
|
||||
|
||||
with mock.patch.object(gemini_llm, "api_client") as mock_client:
|
||||
# Create a mock coroutine that returns the generate_content_response
|
||||
async def mock_coro():
|
||||
return generate_content_response
|
||||
|
||||
mock_client.aio.models.generate_content.return_value = mock_coro()
|
||||
|
||||
# Mock the cache manager module to verify correct method call
|
||||
with mock.patch(
|
||||
"google.adk.models.gemini_context_cache_manager.GeminiContextCacheManager"
|
||||
) as MockCacheManagerClass:
|
||||
mock_cache_manager = MockCacheManagerClass.return_value
|
||||
# Configure cache manager to handle context caching
|
||||
mock_cache_manager.handle_context_caching = AsyncMock(
|
||||
return_value=cache_metadata
|
||||
)
|
||||
|
||||
responses = [
|
||||
resp
|
||||
async for resp in gemini_llm.generate_content_async(
|
||||
llm_request_with_cache, stream=False
|
||||
)
|
||||
]
|
||||
|
||||
# Verify the response was processed
|
||||
assert len(responses) == 1
|
||||
response = responses[0]
|
||||
assert isinstance(response, LlmResponse)
|
||||
assert response.content.parts[0].text == "Hello, how can I help you?"
|
||||
|
||||
# CRITICAL TEST: Verify populate_cache_metadata_in_response was called with correct parameter order
|
||||
mock_cache_manager.populate_cache_metadata_in_response.assert_called_once()
|
||||
call_args = (
|
||||
mock_cache_manager.populate_cache_metadata_in_response.call_args
|
||||
)
|
||||
|
||||
# The first argument should be the LlmResponse (not CacheMetadata)
|
||||
first_arg = call_args[0][0] # First positional argument
|
||||
second_arg = call_args[0][1] # Second positional argument
|
||||
|
||||
# Verify correct parameter order: (llm_response, cache_metadata)
|
||||
assert isinstance(first_arg, LlmResponse), (
|
||||
f"First parameter should be LlmResponse, got {type(first_arg)}. "
|
||||
"This indicates parameters are in wrong order."
|
||||
)
|
||||
assert isinstance(second_arg, CacheMetadata), (
|
||||
f"Second parameter should be CacheMetadata, got {type(second_arg)}. "
|
||||
"This indicates parameters are in wrong order."
|
||||
)
|
||||
|
||||
# Verify the LlmResponse has the expected usage metadata
|
||||
assert first_arg.usage_metadata is not None
|
||||
assert first_arg.usage_metadata.cached_content_token_count == 800
|
||||
assert first_arg.usage_metadata.prompt_token_count == 1500
|
||||
assert first_arg.usage_metadata.candidates_token_count == 150
|
||||
|
||||
# Verify cache metadata is preserved
|
||||
assert second_arg.cache_name == cache_metadata.cache_name
|
||||
assert second_arg.invocations_used == cache_metadata.invocations_used
|
||||
|
||||
@@ -15,6 +15,7 @@
|
||||
from typing import Optional
|
||||
|
||||
from google.adk.agents.base_agent import BaseAgent
|
||||
from google.adk.agents.context_cache_config import ContextCacheConfig
|
||||
from google.adk.agents.invocation_context import InvocationContext
|
||||
from google.adk.agents.llm_agent import LlmAgent
|
||||
from google.adk.apps.app import App
|
||||
@@ -467,5 +468,191 @@ class TestRunnerWithPlugins:
|
||||
)
|
||||
|
||||
|
||||
class TestRunnerCacheConfig:
|
||||
"""Tests for Runner cache config extraction and handling."""
|
||||
|
||||
def setup_method(self):
|
||||
"""Set up test fixtures."""
|
||||
self.session_service = InMemorySessionService()
|
||||
self.artifact_service = InMemoryArtifactService()
|
||||
self.root_agent = MockLlmAgent("root_agent")
|
||||
|
||||
def test_runner_extracts_cache_config_from_app(self):
|
||||
"""Test that Runner extracts cache config from App."""
|
||||
cache_config = ContextCacheConfig(
|
||||
cache_intervals=15, ttl_seconds=3600, min_tokens=1024
|
||||
)
|
||||
|
||||
app = App(
|
||||
name="test_app",
|
||||
root_agent=self.root_agent,
|
||||
context_cache_config=cache_config,
|
||||
)
|
||||
|
||||
runner = Runner(
|
||||
app=app,
|
||||
session_service=self.session_service,
|
||||
artifact_service=self.artifact_service,
|
||||
)
|
||||
|
||||
assert runner.context_cache_config == cache_config
|
||||
assert runner.context_cache_config.cache_intervals == 15
|
||||
assert runner.context_cache_config.ttl_seconds == 3600
|
||||
assert runner.context_cache_config.min_tokens == 1024
|
||||
|
||||
def test_runner_with_app_without_cache_config(self):
|
||||
"""Test Runner with App that has no cache config."""
|
||||
app = App(
|
||||
name="test_app", root_agent=self.root_agent, context_cache_config=None
|
||||
)
|
||||
|
||||
runner = Runner(
|
||||
app=app,
|
||||
session_service=self.session_service,
|
||||
artifact_service=self.artifact_service,
|
||||
)
|
||||
|
||||
assert runner.context_cache_config is None
|
||||
|
||||
def test_runner_without_app_has_no_cache_config(self):
|
||||
"""Test Runner created without App has no cache config."""
|
||||
runner = Runner(
|
||||
app_name="test_app",
|
||||
agent=self.root_agent,
|
||||
session_service=self.session_service,
|
||||
artifact_service=self.artifact_service,
|
||||
)
|
||||
|
||||
assert runner.context_cache_config is None
|
||||
|
||||
def test_runner_cache_config_passed_to_invocation_context(self):
|
||||
"""Test that cache config is passed to InvocationContext."""
|
||||
cache_config = ContextCacheConfig(
|
||||
cache_intervals=20, ttl_seconds=7200, min_tokens=2048
|
||||
)
|
||||
|
||||
app = App(
|
||||
name="test_app",
|
||||
root_agent=self.root_agent,
|
||||
context_cache_config=cache_config,
|
||||
)
|
||||
|
||||
runner = Runner(
|
||||
app=app,
|
||||
session_service=self.session_service,
|
||||
artifact_service=self.artifact_service,
|
||||
)
|
||||
|
||||
# Create a mock session
|
||||
mock_session = Session(
|
||||
id=TEST_SESSION_ID,
|
||||
app_name=TEST_APP_ID,
|
||||
user_id=TEST_USER_ID,
|
||||
events=[],
|
||||
)
|
||||
|
||||
# Create invocation context using runner's method
|
||||
invocation_context = runner._new_invocation_context(mock_session)
|
||||
|
||||
assert invocation_context.context_cache_config == cache_config
|
||||
assert invocation_context.context_cache_config.cache_intervals == 20
|
||||
|
||||
def test_runner_validate_params_return_order(self):
|
||||
"""Test that _validate_runner_params returns values in correct order."""
|
||||
cache_config = ContextCacheConfig(cache_intervals=25)
|
||||
|
||||
app = App(
|
||||
name="order_test_app",
|
||||
root_agent=self.root_agent,
|
||||
context_cache_config=cache_config,
|
||||
)
|
||||
|
||||
runner = Runner(
|
||||
app=app,
|
||||
session_service=self.session_service,
|
||||
artifact_service=self.artifact_service,
|
||||
)
|
||||
|
||||
# Test the validation method directly
|
||||
app_name, agent, context_cache_config, plugins = (
|
||||
runner._validate_runner_params(app, None, None, None)
|
||||
)
|
||||
|
||||
assert app_name == "order_test_app"
|
||||
assert agent == self.root_agent
|
||||
assert context_cache_config == cache_config
|
||||
assert context_cache_config.cache_intervals == 25
|
||||
assert plugins == []
|
||||
|
||||
def test_runner_validate_params_without_app(self):
|
||||
"""Test _validate_runner_params without App returns None for cache config."""
|
||||
runner = Runner(
|
||||
app_name="test_app",
|
||||
agent=self.root_agent,
|
||||
session_service=self.session_service,
|
||||
artifact_service=self.artifact_service,
|
||||
)
|
||||
|
||||
app_name, agent, context_cache_config, plugins = (
|
||||
runner._validate_runner_params(None, "test_app", self.root_agent, None)
|
||||
)
|
||||
|
||||
assert app_name == "test_app"
|
||||
assert agent == self.root_agent
|
||||
assert context_cache_config is None
|
||||
assert plugins is None
|
||||
|
||||
def test_runner_app_name_and_agent_extracted_correctly(self):
|
||||
"""Test that app_name and agent are correctly extracted from App."""
|
||||
cache_config = ContextCacheConfig()
|
||||
|
||||
app = App(
|
||||
name="extracted_app",
|
||||
root_agent=self.root_agent,
|
||||
context_cache_config=cache_config,
|
||||
)
|
||||
|
||||
runner = Runner(
|
||||
app=app,
|
||||
session_service=self.session_service,
|
||||
artifact_service=self.artifact_service,
|
||||
)
|
||||
|
||||
assert runner.app_name == "extracted_app"
|
||||
assert runner.agent == self.root_agent
|
||||
assert runner.context_cache_config == cache_config
|
||||
|
||||
def test_runner_realistic_cache_config_scenario(self):
|
||||
"""Test realistic scenario with production-like cache config."""
|
||||
# Production cache config
|
||||
production_cache_config = ContextCacheConfig(
|
||||
cache_intervals=30, ttl_seconds=14400, min_tokens=4096 # 4 hours
|
||||
)
|
||||
|
||||
app = App(
|
||||
name="production_app",
|
||||
root_agent=self.root_agent,
|
||||
context_cache_config=production_cache_config,
|
||||
)
|
||||
|
||||
runner = Runner(
|
||||
app=app,
|
||||
session_service=self.session_service,
|
||||
artifact_service=self.artifact_service,
|
||||
)
|
||||
|
||||
# Verify all settings are preserved
|
||||
assert runner.context_cache_config.cache_intervals == 30
|
||||
assert runner.context_cache_config.ttl_seconds == 14400
|
||||
assert runner.context_cache_config.ttl_string == "14400s"
|
||||
assert runner.context_cache_config.min_tokens == 4096
|
||||
|
||||
# Verify string representation
|
||||
expected_str = (
|
||||
"ContextCacheConfig(cache_intervals=30, ttl=14400s, min_tokens=4096)"
|
||||
)
|
||||
assert str(runner.context_cache_config) == expected_str
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__])
|
||||
|
||||
@@ -0,0 +1,450 @@
|
||||
# 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 CachePerformanceAnalyzer."""
|
||||
|
||||
import time
|
||||
from unittest.mock import AsyncMock
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from google.adk.events.event import Event
|
||||
from google.adk.models.cache_metadata import CacheMetadata
|
||||
from google.adk.sessions.base_session_service import BaseSessionService
|
||||
from google.adk.sessions.session import Session
|
||||
from google.adk.utils.cache_performance_analyzer import CachePerformanceAnalyzer
|
||||
from google.genai import types
|
||||
import pytest
|
||||
|
||||
|
||||
class TestCachePerformanceAnalyzer:
|
||||
"""Test suite for CachePerformanceAnalyzer."""
|
||||
|
||||
def setup_method(self):
|
||||
"""Set up test fixtures."""
|
||||
self.mock_session_service = MagicMock(spec=BaseSessionService)
|
||||
self.analyzer = CachePerformanceAnalyzer(self.mock_session_service)
|
||||
|
||||
def create_cache_metadata(
|
||||
self, invocations_used=1, cache_name="test-cache", cached_contents_count=5
|
||||
):
|
||||
"""Helper to create test CacheMetadata."""
|
||||
return CacheMetadata(
|
||||
cache_name=(
|
||||
f"projects/test/locations/us-central1/cachedContents/{cache_name}"
|
||||
),
|
||||
expire_time=time.time() + 1800,
|
||||
fingerprint="test_fingerprint",
|
||||
invocations_used=invocations_used,
|
||||
cached_contents_count=cached_contents_count,
|
||||
created_at=time.time() - 600,
|
||||
)
|
||||
|
||||
def create_mock_usage_metadata(
|
||||
self, prompt_tokens=1000, cached_tokens=500, candidates_tokens=100
|
||||
):
|
||||
"""Helper to create mock usage metadata."""
|
||||
return types.GenerateContentResponseUsageMetadata(
|
||||
prompt_token_count=prompt_tokens,
|
||||
cached_content_token_count=cached_tokens,
|
||||
candidates_token_count=candidates_tokens,
|
||||
total_token_count=prompt_tokens + candidates_tokens,
|
||||
)
|
||||
|
||||
def create_mock_event(
|
||||
self, author="test_agent", cache_metadata=None, usage_metadata=None
|
||||
):
|
||||
"""Helper to create mock event."""
|
||||
event = Event(author=author, cache_metadata=cache_metadata)
|
||||
if usage_metadata:
|
||||
event.usage_metadata = usage_metadata
|
||||
return event
|
||||
|
||||
def test_init(self):
|
||||
"""Test analyzer initialization."""
|
||||
assert self.analyzer.session_service == self.mock_session_service
|
||||
|
||||
async def test_get_agent_cache_history_empty_session(self):
|
||||
"""Test getting cache history from empty session."""
|
||||
mock_session = Session(
|
||||
id="test_session",
|
||||
app_name="test_app",
|
||||
user_id="test_user",
|
||||
events=[],
|
||||
)
|
||||
self.mock_session_service.get_session = AsyncMock(return_value=mock_session)
|
||||
|
||||
result = await self.analyzer._get_agent_cache_history(
|
||||
"test_session", "test_user", "test_app", "test_agent"
|
||||
)
|
||||
|
||||
assert result == []
|
||||
|
||||
async def test_get_agent_cache_history_no_cache_events(self):
|
||||
"""Test getting cache history when no events have cache metadata."""
|
||||
events = [
|
||||
self.create_mock_event(author="test_agent"),
|
||||
self.create_mock_event(author="other_agent"),
|
||||
self.create_mock_event(author="test_agent"),
|
||||
]
|
||||
|
||||
mock_session = Session(
|
||||
id="test_session",
|
||||
app_name="test_app",
|
||||
user_id="test_user",
|
||||
events=events,
|
||||
)
|
||||
self.mock_session_service.get_session = AsyncMock(return_value=mock_session)
|
||||
|
||||
result = await self.analyzer._get_agent_cache_history(
|
||||
"test_session", "test_user", "test_app", "test_agent"
|
||||
)
|
||||
|
||||
assert result == []
|
||||
|
||||
async def test_get_agent_cache_history_specific_agent(self):
|
||||
"""Test getting cache history for specific agent."""
|
||||
cache1 = self.create_cache_metadata(invocations_used=1, cache_name="cache1")
|
||||
cache2 = self.create_cache_metadata(invocations_used=3, cache_name="cache2")
|
||||
cache3 = self.create_cache_metadata(invocations_used=5, cache_name="cache3")
|
||||
|
||||
events = [
|
||||
self.create_mock_event(author="test_agent", cache_metadata=cache1),
|
||||
self.create_mock_event(author="other_agent", cache_metadata=cache2),
|
||||
self.create_mock_event(author="test_agent", cache_metadata=cache3),
|
||||
self.create_mock_event(author="test_agent"), # No cache metadata
|
||||
]
|
||||
|
||||
mock_session = Session(
|
||||
id="test_session",
|
||||
app_name="test_app",
|
||||
user_id="test_user",
|
||||
events=events,
|
||||
)
|
||||
self.mock_session_service.get_session = AsyncMock(return_value=mock_session)
|
||||
|
||||
result = await self.analyzer._get_agent_cache_history(
|
||||
"test_session", "test_user", "test_app", "test_agent"
|
||||
)
|
||||
|
||||
# Should only return cache metadata for test_agent
|
||||
assert len(result) == 2
|
||||
assert result[0] == cache1
|
||||
assert result[1] == cache3
|
||||
|
||||
async def test_get_agent_cache_history_all_agents(self):
|
||||
"""Test getting cache history for all agents."""
|
||||
cache1 = self.create_cache_metadata(invocations_used=1, cache_name="cache1")
|
||||
cache2 = self.create_cache_metadata(invocations_used=3, cache_name="cache2")
|
||||
|
||||
events = [
|
||||
self.create_mock_event(author="agent1", cache_metadata=cache1),
|
||||
self.create_mock_event(author="agent2", cache_metadata=cache2),
|
||||
self.create_mock_event(author="agent1"), # No cache metadata
|
||||
]
|
||||
|
||||
mock_session = Session(
|
||||
id="test_session",
|
||||
app_name="test_app",
|
||||
user_id="test_user",
|
||||
events=events,
|
||||
)
|
||||
self.mock_session_service.get_session = AsyncMock(return_value=mock_session)
|
||||
|
||||
# Pass None for agent_name to get all agents
|
||||
result = await self.analyzer._get_agent_cache_history(
|
||||
"test_session", "test_user", "test_app", None
|
||||
)
|
||||
|
||||
# Should return cache metadata for all agents
|
||||
assert len(result) == 2
|
||||
assert result[0] == cache1
|
||||
assert result[1] == cache2
|
||||
|
||||
async def test_analyze_agent_cache_performance_no_cache_data(self):
|
||||
"""Test analysis with no cache data."""
|
||||
mock_session = Session(
|
||||
id="test_session",
|
||||
app_name="test_app",
|
||||
user_id="test_user",
|
||||
events=[],
|
||||
)
|
||||
self.mock_session_service.get_session = AsyncMock(return_value=mock_session)
|
||||
|
||||
result = await self.analyzer.analyze_agent_cache_performance(
|
||||
"test_session", "test_user", "test_app", "test_agent"
|
||||
)
|
||||
|
||||
assert result["status"] == "no_cache_data"
|
||||
|
||||
async def test_analyze_agent_cache_performance_with_cache_data(self):
|
||||
"""Test comprehensive analysis with cache data and token metrics."""
|
||||
cache1 = self.create_cache_metadata(invocations_used=2, cache_name="cache1")
|
||||
cache2 = self.create_cache_metadata(invocations_used=5, cache_name="cache2")
|
||||
cache3 = self.create_cache_metadata(invocations_used=8, cache_name="cache3")
|
||||
|
||||
usage1 = self.create_mock_usage_metadata(
|
||||
prompt_tokens=1000, cached_tokens=800
|
||||
)
|
||||
usage2 = self.create_mock_usage_metadata(
|
||||
prompt_tokens=1500, cached_tokens=1200
|
||||
)
|
||||
usage3 = self.create_mock_usage_metadata(prompt_tokens=800, cached_tokens=0)
|
||||
|
||||
events = [
|
||||
self.create_mock_event(
|
||||
author="test_agent", cache_metadata=cache1, usage_metadata=usage1
|
||||
),
|
||||
self.create_mock_event(author="other_agent", cache_metadata=cache2),
|
||||
self.create_mock_event(
|
||||
author="test_agent", cache_metadata=cache2, usage_metadata=usage2
|
||||
),
|
||||
self.create_mock_event(
|
||||
author="test_agent", cache_metadata=cache3, usage_metadata=usage3
|
||||
),
|
||||
]
|
||||
|
||||
mock_session = Session(
|
||||
id="test_session",
|
||||
app_name="test_app",
|
||||
user_id="test_user",
|
||||
events=events,
|
||||
)
|
||||
self.mock_session_service.get_session = AsyncMock(return_value=mock_session)
|
||||
|
||||
result = await self.analyzer.analyze_agent_cache_performance(
|
||||
"test_session", "test_user", "test_app", "test_agent"
|
||||
)
|
||||
|
||||
# Basic cache metrics
|
||||
assert result["status"] == "active"
|
||||
assert result["requests_with_cache"] == 3
|
||||
assert result["cache_refreshes"] == 3 # 3 unique cache names
|
||||
assert result["total_invocations"] == 15 # 2 + 5 + 8
|
||||
|
||||
expected_avg_invocations = (2 + 5 + 8) / 3 # 5.0
|
||||
assert result["avg_invocations_used"] == expected_avg_invocations
|
||||
|
||||
# Token metrics
|
||||
assert result["total_prompt_tokens"] == 3300 # 1000 + 1500 + 800
|
||||
assert result["total_cached_tokens"] == 2000 # 800 + 1200 + 0
|
||||
assert result["total_requests"] == 3
|
||||
assert (
|
||||
result["requests_with_cache_hits"] == 2
|
||||
) # Only first two have cached tokens
|
||||
|
||||
# Calculated metrics
|
||||
expected_hit_ratio = (2000 / 3300) * 100 # ~60.6%
|
||||
expected_utilization = (2 / 3) * 100 # ~66.7%
|
||||
expected_avg_cached = 2000 / 3 # ~666.7
|
||||
|
||||
assert abs(result["cache_hit_ratio_percent"] - expected_hit_ratio) < 0.01
|
||||
assert (
|
||||
abs(result["cache_utilization_ratio_percent"] - expected_utilization)
|
||||
< 0.01
|
||||
)
|
||||
assert (
|
||||
abs(result["avg_cached_tokens_per_request"] - expected_avg_cached)
|
||||
< 0.01
|
||||
)
|
||||
|
||||
async def test_analyze_agent_cache_performance_single_cache(self):
|
||||
"""Test analysis with single cache instance."""
|
||||
cache = self.create_cache_metadata(
|
||||
invocations_used=10, cache_name="single_cache"
|
||||
)
|
||||
usage = self.create_mock_usage_metadata(
|
||||
prompt_tokens=2000, cached_tokens=1500
|
||||
)
|
||||
|
||||
events = [
|
||||
self.create_mock_event(
|
||||
author="test_agent", cache_metadata=cache, usage_metadata=usage
|
||||
),
|
||||
]
|
||||
|
||||
mock_session = Session(
|
||||
id="test_session",
|
||||
app_name="test_app",
|
||||
user_id="test_user",
|
||||
events=events,
|
||||
)
|
||||
self.mock_session_service.get_session = AsyncMock(return_value=mock_session)
|
||||
|
||||
result = await self.analyzer.analyze_agent_cache_performance(
|
||||
"test_session", "test_user", "test_app", "test_agent"
|
||||
)
|
||||
|
||||
assert result["status"] == "active"
|
||||
assert result["requests_with_cache"] == 1
|
||||
assert result["avg_invocations_used"] == 10.0
|
||||
assert result["cache_refreshes"] == 1
|
||||
assert result["total_invocations"] == 10
|
||||
assert result["latest_cache"] == cache.cache_name
|
||||
|
||||
# Token metrics for single request
|
||||
assert result["total_prompt_tokens"] == 2000
|
||||
assert result["total_cached_tokens"] == 1500
|
||||
assert result["cache_hit_ratio_percent"] == 75.0 # 1500/2000 * 100
|
||||
assert result["cache_utilization_ratio_percent"] == 100.0 # 1/1 * 100
|
||||
assert result["avg_cached_tokens_per_request"] == 1500.0
|
||||
|
||||
async def test_analyze_agent_cache_performance_no_token_data(self):
|
||||
"""Test analysis when events have no usage_metadata."""
|
||||
cache = self.create_cache_metadata(invocations_used=5)
|
||||
|
||||
events = [
|
||||
self.create_mock_event(author="test_agent", cache_metadata=cache),
|
||||
]
|
||||
|
||||
mock_session = Session(
|
||||
id="test_session",
|
||||
app_name="test_app",
|
||||
user_id="test_user",
|
||||
events=events,
|
||||
)
|
||||
self.mock_session_service.get_session = AsyncMock(return_value=mock_session)
|
||||
|
||||
result = await self.analyzer.analyze_agent_cache_performance(
|
||||
"test_session", "test_user", "test_app", "test_agent"
|
||||
)
|
||||
|
||||
# Should still work but with zero token metrics
|
||||
assert result["status"] == "active"
|
||||
assert result["requests_with_cache"] == 1
|
||||
assert result["total_prompt_tokens"] == 0
|
||||
assert result["total_cached_tokens"] == 0
|
||||
assert result["cache_hit_ratio_percent"] == 0.0
|
||||
assert result["cache_utilization_ratio_percent"] == 0.0
|
||||
assert result["avg_cached_tokens_per_request"] == 0.0
|
||||
|
||||
async def test_analyze_agent_cache_performance_zero_invocations(self):
|
||||
"""Test analysis with zero invocations."""
|
||||
cache = self.create_cache_metadata(
|
||||
invocations_used=0, cache_name="zero_cache"
|
||||
)
|
||||
usage = self.create_mock_usage_metadata(
|
||||
prompt_tokens=1000, cached_tokens=500
|
||||
)
|
||||
|
||||
events = [
|
||||
self.create_mock_event(
|
||||
author="test_agent", cache_metadata=cache, usage_metadata=usage
|
||||
),
|
||||
]
|
||||
|
||||
mock_session = Session(
|
||||
id="test_session",
|
||||
app_name="test_app",
|
||||
user_id="test_user",
|
||||
events=events,
|
||||
)
|
||||
self.mock_session_service.get_session = AsyncMock(return_value=mock_session)
|
||||
|
||||
result = await self.analyzer.analyze_agent_cache_performance(
|
||||
"test_session", "test_user", "test_app", "test_agent"
|
||||
)
|
||||
|
||||
assert result["status"] == "active"
|
||||
assert result["avg_invocations_used"] == 0.0
|
||||
assert result["total_invocations"] == 0
|
||||
|
||||
# Token metrics should still work
|
||||
assert result["total_prompt_tokens"] == 1000
|
||||
assert result["total_cached_tokens"] == 500
|
||||
|
||||
async def test_session_service_integration(self):
|
||||
"""Test integration with session service."""
|
||||
cache_metadata = self.create_cache_metadata(invocations_used=7)
|
||||
|
||||
events = [
|
||||
self.create_mock_event(
|
||||
author="integration_agent", cache_metadata=cache_metadata
|
||||
),
|
||||
]
|
||||
|
||||
mock_session = Session(
|
||||
id="integration_session",
|
||||
app_name="integration_app",
|
||||
user_id="integration_user",
|
||||
events=events,
|
||||
)
|
||||
|
||||
# Configure the mock to return the session
|
||||
self.mock_session_service.get_session = AsyncMock(return_value=mock_session)
|
||||
|
||||
result = await self.analyzer.analyze_agent_cache_performance(
|
||||
"integration_session",
|
||||
"integration_user",
|
||||
"integration_app",
|
||||
"integration_agent",
|
||||
)
|
||||
|
||||
# Verify the session service was called with correct parameters (twice internally)
|
||||
assert self.mock_session_service.get_session.call_count == 2
|
||||
self.mock_session_service.get_session.assert_called_with(
|
||||
session_id="integration_session",
|
||||
app_name="integration_app",
|
||||
user_id="integration_user",
|
||||
)
|
||||
|
||||
assert result["status"] == "active"
|
||||
assert result["requests_with_cache"] == 1
|
||||
|
||||
async def test_mixed_agents_filtering(self):
|
||||
"""Test that analysis correctly filters by agent name."""
|
||||
target_cache = self.create_cache_metadata(
|
||||
invocations_used=3, cache_name="target"
|
||||
)
|
||||
other_cache = self.create_cache_metadata(
|
||||
invocations_used=5, cache_name="other"
|
||||
)
|
||||
|
||||
target_usage = self.create_mock_usage_metadata(
|
||||
prompt_tokens=1000, cached_tokens=800
|
||||
)
|
||||
other_usage = self.create_mock_usage_metadata(
|
||||
prompt_tokens=2000, cached_tokens=1600
|
||||
)
|
||||
|
||||
events = [
|
||||
self.create_mock_event(
|
||||
author="target_agent",
|
||||
cache_metadata=target_cache,
|
||||
usage_metadata=target_usage,
|
||||
),
|
||||
self.create_mock_event(
|
||||
author="other_agent",
|
||||
cache_metadata=other_cache,
|
||||
usage_metadata=other_usage,
|
||||
),
|
||||
self.create_mock_event(author="target_agent"), # No cache data
|
||||
]
|
||||
|
||||
mock_session = Session(
|
||||
id="test_session",
|
||||
app_name="test_app",
|
||||
user_id="test_user",
|
||||
events=events,
|
||||
)
|
||||
self.mock_session_service.get_session = AsyncMock(return_value=mock_session)
|
||||
|
||||
result = await self.analyzer.analyze_agent_cache_performance(
|
||||
"test_session", "test_user", "test_app", "target_agent"
|
||||
)
|
||||
|
||||
# Should only include target_agent's data
|
||||
assert result["requests_with_cache"] == 1
|
||||
assert result["total_invocations"] == 3
|
||||
assert result["total_prompt_tokens"] == 1000 # Only target_agent's tokens
|
||||
assert result["total_cached_tokens"] == 800 # Only target_agent's tokens
|
||||
Reference in New Issue
Block a user