github-featured/google--adk-python
2
0
Fork 0
Code Issues Pull Requests Actions 7 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

feat: record context cache state on the LLM call span

Browse Source 
Co-authored-by: George Weale <gweale@google.com>
PiperOrigin-RevId: 968671743
This commit is contained in:
George Weale
2026-08-21 13:49:28 -07:00
committed by Copybara-Service
parent 26110c7559
commit c0614d6580
3 changed files with 257 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/google/adk/telemetry/_adk_attributes.py
+11
View File
@@ -33,3 +33,14 @@ ADK_EXPERIMENTAL_SKILL_ADDITIONAL_TOOLS = (
)
ADK_EXPERIMENTAL_SKILL_SOURCE_URI = 'adk.experimental.skill.source.uri'
ADK_EXPERIMENTAL_SKILL_RESOURCE_PATH = 'adk.experimental.skill.resource.path'
ADK_EXPERIMENTAL_CONTEXT_CACHE_HIT = 'adk.experimental.context_cache.hit'
ADK_EXPERIMENTAL_CONTEXT_CACHE_FINGERPRINT = (
'adk.experimental.context_cache.fingerprint'
)
ADK_EXPERIMENTAL_CONTEXT_CACHE_CONTENTS_COUNT = (
'adk.experimental.context_cache.contents_count'
)
ADK_EXPERIMENTAL_CONTEXT_CACHE_INVOCATIONS_USED = (
'adk.experimental.context_cache.invocations_used'
)
src/google/adk/telemetry/tracing.py
+39
View File
@@ -65,6 +65,10 @@ from .. import version
from ..utils.env_utils import is_enterprise_mode_enabled
from ..utils.model_name_utils import extract_model_name
from ..utils.model_name_utils import is_gemini_model
from ._adk_attributes import ADK_EXPERIMENTAL_CONTEXT_CACHE_CONTENTS_COUNT
from ._adk_attributes import ADK_EXPERIMENTAL_CONTEXT_CACHE_FINGERPRINT
from ._adk_attributes import ADK_EXPERIMENTAL_CONTEXT_CACHE_HIT
from ._adk_attributes import ADK_EXPERIMENTAL_CONTEXT_CACHE_INVOCATIONS_USED
from ._experimental_semconv import maybe_log_completion_details
from ._experimental_semconv import set_operation_details_attributes_from_request
from ._experimental_semconv import set_operation_details_attributes_from_response
@@ -100,6 +104,7 @@ if TYPE_CHECKING:
from ..agents.base_agent import BaseAgent
from ..agents.invocation_context import InvocationContext
from ..events.event import Event
from ..models.cache_metadata import CacheMetadata
from ..models.llm_request import LlmRequest
from ..models.llm_response import LlmResponse
from ..tools.base_tool import BaseTool
@@ -353,6 +358,32 @@ def _set_usage_metadata_attributes(
span.set_attributes(TokenUsage(usage_metadata).to_attributes())
def _set_context_cache_attributes(
span: Span,
cache_metadata: CacheMetadata | None,
telemetry_config: TelemetryConfig,
) -> None:
"""Records context cache state on the given span."""
if cache_metadata is None:
return
# The fingerprint is a content hash, so these attributes stay behind the
# experimental opt-in rather than landing on every span by default.
if not telemetry_config.should_emit_experimental_telemetry:
return
attributes: dict[str, AttributeValue] = {
ADK_EXPERIMENTAL_CONTEXT_CACHE_HIT: cache_metadata.cache_name is not None,
ADK_EXPERIMENTAL_CONTEXT_CACHE_FINGERPRINT: cache_metadata.fingerprint,
ADK_EXPERIMENTAL_CONTEXT_CACHE_CONTENTS_COUNT: (
cache_metadata.contents_count
),
}
if cache_metadata.invocations_used is not None:
attributes[ADK_EXPERIMENTAL_CONTEXT_CACHE_INVOCATIONS_USED] = (
cache_metadata.invocations_used
)
span.set_attributes(attributes)
def trace_call_llm(
invocation_context: InvocationContext,
event_id: str,
@@ -441,6 +472,9 @@ def trace_call_llm(
span.set_attribute("gcp.vertex.agent.llm_response", "{}")
_set_usage_metadata_attributes(span, llm_response.usage_metadata)
_set_context_cache_attributes(
span, getattr(llm_response, "cache_metadata", None), telemetry_config
)
if llm_response.finish_reason:
try:
finish_reason_str = llm_response.finish_reason.value.lower()
@@ -962,6 +996,11 @@ def trace_inference_result(
if finish_reason := llm_response.finish_reason:
span.set_attribute(GEN_AI_RESPONSE_FINISH_REASONS, [finish_reason.lower()])
_set_usage_metadata_attributes(span, llm_response.usage_metadata)
# Callers outside adk pass their own response objects here, which are only
# required to carry the fields this function already read.
_set_context_cache_attributes(
span, getattr(llm_response, "cache_metadata", None), telemetry_config
)
if telemetry_config.should_use_experimental_genai_semconv and isinstance(
gc_span, GenerateContentSpan
tests/unittests/telemetry/test_spans.py
+207 -2
View File
@@ -24,10 +24,16 @@ from google.adk.agents.run_config import RunConfig
from google.adk.errors.tool_execution_error import ToolErrorType
from google.adk.errors.tool_execution_error import ToolExecutionError
from google.adk.events.event import Event
from google.adk.models.cache_metadata import CacheMetadata
from google.adk.models.llm_request import LlmRequest
from google.adk.models.llm_response import LlmResponse
from google.adk.sessions.in_memory_session_service import InMemorySessionService
from google.adk.telemetry import TelemetryConfig
from google.adk.telemetry import tracing
from google.adk.telemetry._adk_attributes import ADK_EXPERIMENTAL_CONTEXT_CACHE_CONTENTS_COUNT
from google.adk.telemetry._adk_attributes import ADK_EXPERIMENTAL_CONTEXT_CACHE_FINGERPRINT
from google.adk.telemetry._adk_attributes import ADK_EXPERIMENTAL_CONTEXT_CACHE_HIT
from google.adk.telemetry._adk_attributes import ADK_EXPERIMENTAL_CONTEXT_CACHE_INVOCATIONS_USED
from google.adk.telemetry._experimental_semconv import _safe_json_serialize_no_whitespaces
from google.adk.telemetry.tracing import _use_extra_generate_content_attributes
from google.adk.telemetry.tracing import ADK_CAPTURE_MESSAGE_CONTENT_IN_SPANS
@@ -102,7 +108,9 @@ def mock_event_fixture():
async def _create_invocation_context(
agent: LlmAgent, state: Optional[dict[str, object]] = None
agent: LlmAgent,
state: Optional[dict[str, object]] = None,
run_config: Optional[RunConfig] = None,
) -> InvocationContext:
session_service = InMemorySessionService()
session = await session_service.create_session(
@@ -113,7 +121,7 @@ async def _create_invocation_context(
agent=agent,
session=session,
session_service=session_service,
run_config=RunConfig(),
run_config=run_config or RunConfig(),
)
return invocation_context
@@ -288,6 +296,203 @@ async def test_trace_call_llm_with_no_usage_metadata(
)
_EXPERIMENTAL_TELEMETRY_ON = RunConfig(
telemetry=TelemetryConfig(adk_experimental_telemetry_opt_in=True)
)
@pytest.mark.asyncio
async def test_trace_call_llm_with_active_context_cache(
monkeypatch, mock_span_fixture
):
"""Test trace_call_llm records the state of an active context cache."""
monkeypatch.setattr(
'opentelemetry.trace.get_current_span', lambda: mock_span_fixture
)
agent = LlmAgent(name='test_agent')
invocation_context = await _create_invocation_context(
agent, run_config=_EXPERIMENTAL_TELEMETRY_ON
)
llm_request = LlmRequest(
model='gemini-pro',
contents=[
types.Content(role='user', parts=[types.Part(text='Hello')]),
],
)
llm_response = LlmResponse(
turn_complete=True,
finish_reason=types.FinishReason.STOP,
cache_metadata=CacheMetadata(
cache_name='projects/p/locations/l/cachedContents/c',
expire_time=1893456000.0,
fingerprint='fp-123',
invocations_used=3,
contents_count=5,
),
)
trace_call_llm(invocation_context, 'test_event_id', llm_request, llm_response)
mock_span_fixture.set_attributes.assert_any_call({
ADK_EXPERIMENTAL_CONTEXT_CACHE_HIT: True,
ADK_EXPERIMENTAL_CONTEXT_CACHE_FINGERPRINT: 'fp-123',
ADK_EXPERIMENTAL_CONTEXT_CACHE_CONTENTS_COUNT: 5,
ADK_EXPERIMENTAL_CONTEXT_CACHE_INVOCATIONS_USED: 3,
})
@pytest.mark.asyncio
async def test_trace_call_llm_with_fingerprint_only_context_cache(
monkeypatch, mock_span_fixture
):
"""Test trace_call_llm records a miss when no cache is active."""
monkeypatch.setattr(
'opentelemetry.trace.get_current_span', lambda: mock_span_fixture
)
agent = LlmAgent(name='test_agent')
invocation_context = await _create_invocation_context(
agent, run_config=_EXPERIMENTAL_TELEMETRY_ON
)
llm_request = LlmRequest(
model='gemini-pro',
contents=[
types.Content(role='user', parts=[types.Part(text='Hello')]),
],
)
llm_response = LlmResponse(
turn_complete=True,
finish_reason=types.FinishReason.STOP,
cache_metadata=CacheMetadata(fingerprint='fp-123', contents_count=8),
)
trace_call_llm(invocation_context, 'test_event_id', llm_request, llm_response)
# The exact dict also pins that invocations_used is left out when unset.
mock_span_fixture.set_attributes.assert_any_call({
ADK_EXPERIMENTAL_CONTEXT_CACHE_HIT: False,
ADK_EXPERIMENTAL_CONTEXT_CACHE_FINGERPRINT: 'fp-123',
ADK_EXPERIMENTAL_CONTEXT_CACHE_CONTENTS_COUNT: 8,
})
@pytest.mark.asyncio
async def test_trace_call_llm_omits_context_cache_without_opt_in(
monkeypatch, mock_span_fixture
):
"""Test context cache attributes stay off unless experimental is opted in."""
monkeypatch.setattr(
'opentelemetry.trace.get_current_span', lambda: mock_span_fixture
)
agent = LlmAgent(name='test_agent')
invocation_context = await _create_invocation_context(agent)
llm_request = LlmRequest(
model='gemini-pro',
contents=[
types.Content(role='user', parts=[types.Part(text='Hello')]),
],
)
llm_response = LlmResponse(
turn_complete=True,
finish_reason=types.FinishReason.STOP,
cache_metadata=CacheMetadata(
cache_name='projects/p/locations/l/cachedContents/c',
expire_time=1893456000.0,
fingerprint='fp-123',
invocations_used=3,
contents_count=5,
),
)
trace_call_llm(invocation_context, 'test_event_id', llm_request, llm_response)
# The span is still traced; only the experimental attributes are withheld.
mock_span_fixture.set_attribute.assert_any_call(
'gen_ai.system', 'gcp.vertex.agent'
)
set_keys = [
call.args[0] for call in mock_span_fixture.set_attribute.call_args_list
]
for call in mock_span_fixture.set_attributes.call_args_list:
set_keys.extend(call.args[0])
assert not [key for key in set_keys if key.startswith('adk.experimental.')]
@pytest.mark.asyncio
async def test_trace_inference_result_with_context_cache(mock_span_fixture):
"""Test the generate_content span also carries context cache state."""
agent = LlmAgent(name='test_agent')
invocation_context = await _create_invocation_context(
agent, run_config=_EXPERIMENTAL_TELEMETRY_ON
)
llm_response = LlmResponse(
turn_complete=True,
finish_reason=types.FinishReason.STOP,
cache_metadata=CacheMetadata(
cache_name='projects/p/locations/l/cachedContents/c',
expire_time=1893456000.0,
fingerprint='fp-123',
invocations_used=3,
contents_count=5,
),
)
trace_inference_result(invocation_context, mock_span_fixture, llm_response)
mock_span_fixture.set_attributes.assert_any_call({
ADK_EXPERIMENTAL_CONTEXT_CACHE_HIT: True,
ADK_EXPERIMENTAL_CONTEXT_CACHE_FINGERPRINT: 'fp-123',
ADK_EXPERIMENTAL_CONTEXT_CACHE_CONTENTS_COUNT: 5,
ADK_EXPERIMENTAL_CONTEXT_CACHE_INVOCATIONS_USED: 3,
})
@pytest.mark.asyncio
async def test_trace_inference_result_omits_context_cache_without_opt_in(
mock_span_fixture,
):
"""Test the generate_content span withholds cache state without opt-in."""
agent = LlmAgent(name='test_agent')
invocation_context = await _create_invocation_context(agent)
llm_response = LlmResponse(
turn_complete=True,
finish_reason=types.FinishReason.STOP,
cache_metadata=CacheMetadata(fingerprint='fp-123', contents_count=5),
)
trace_inference_result(invocation_context, mock_span_fixture, llm_response)
set_keys = []
for call in mock_span_fixture.set_attributes.call_args_list:
set_keys.extend(call.args[0])
assert not [key for key in set_keys if key.startswith('adk.experimental.')]
@pytest.mark.asyncio
async def test_trace_inference_result_allows_a_response_without_cache_metadata(
mock_span_fixture,
):
"""Test a caller's own response object without cache_metadata is accepted."""
class ResponseWithoutCacheMetadata:
"""Stands in for a caller's response type outside adk."""
partial = False
finish_reason = types.FinishReason.STOP
usage_metadata = None
content = None
model_version = 'gemini-pro'
agent = LlmAgent(name='test_agent')
invocation_context = await _create_invocation_context(agent)
trace_inference_result(
invocation_context, mock_span_fixture, ResponseWithoutCacheMetadata()
)
@pytest.mark.asyncio
async def test_trace_call_llm_with_binary_content(
monkeypatch, mock_span_fixture