diff --git a/src/google/adk/telemetry/_adk_attributes.py b/src/google/adk/telemetry/_adk_attributes.py index 16fb1774..1f27b1a7 100644 --- a/src/google/adk/telemetry/_adk_attributes.py +++ b/src/google/adk/telemetry/_adk_attributes.py @@ -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' +) diff --git a/src/google/adk/telemetry/tracing.py b/src/google/adk/telemetry/tracing.py index 4361c0ea..cd919ea3 100644 --- a/src/google/adk/telemetry/tracing.py +++ b/src/google/adk/telemetry/tracing.py @@ -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 diff --git a/tests/unittests/telemetry/test_spans.py b/tests/unittests/telemetry/test_spans.py index ff92d70d..5fa8e9a3 100644 --- a/tests/unittests/telemetry/test_spans.py +++ b/tests/unittests/telemetry/test_spans.py @@ -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