diff --git a/src/google/adk/telemetry/tracing.py b/src/google/adk/telemetry/tracing.py index d54f75173b..a5ae411470 100644 --- a/src/google/adk/telemetry/tracing.py +++ b/src/google/adk/telemetry/tracing.py @@ -62,14 +62,16 @@ # By default some ADK spans include attributes with potential PII data. # This env, when set to false, allows to disable populating those attributes. -ADK_CAPTURE_MESSAGE_CONTENT_IN_SPANS = 'ADK_CAPTURE_MESSAGE_CONTENT_IN_SPANS' +ADK_CAPTURE_MESSAGE_CONTENT_IN_SPANS = "ADK_CAPTURE_MESSAGE_CONTENT_IN_SPANS" # Standard OTEL env variable to enable logging of prompt/response content. OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT = ( - 'OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT' + "OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT" ) -USER_CONTENT_ELIDED = '' +USER_CONTENT_ELIDED = "" + +EMPTY_JSON_STRING = "{}" # Needed to avoid circular imports if TYPE_CHECKING: @@ -81,18 +83,18 @@ from ..tools.base_tool import BaseTool tracer = trace.get_tracer( - instrumenting_module_name='gcp.vertex.agent', + instrumenting_module_name="gcp.vertex.agent", instrumenting_library_version=version.__version__, schema_url=Schemas.V1_36_0.value, ) otel_logger = _logs.get_logger( - instrumenting_module_name='gcp.vertex.agent', + instrumenting_module_name="gcp.vertex.agent", instrumenting_library_version=version.__version__, schema_url=Schemas.V1_36_0.value, ) -logger = logging.getLogger('google_adk.' + __name__) +logger = logging.getLogger("google_adk." + __name__) def _safe_json_serialize(obj) -> str: @@ -108,10 +110,10 @@ def _safe_json_serialize(obj) -> str: try: # Try direct JSON serialization first return json.dumps( - obj, ensure_ascii=False, default=lambda o: '' + obj, ensure_ascii=False, default=lambda o: "" ) except (TypeError, OverflowError): - return '' + return "" def trace_agent_invocation( @@ -138,7 +140,7 @@ def trace_agent_invocation( """ # Required - span.set_attribute(GEN_AI_OPERATION_NAME, 'invoke_agent') + span.set_attribute(GEN_AI_OPERATION_NAME, "invoke_agent") # Conditionally Required span.set_attribute(GEN_AI_AGENT_DESCRIPTION, agent.description) @@ -161,7 +163,7 @@ def trace_tool_call( """ span = trace.get_current_span() - span.set_attribute(GEN_AI_OPERATION_NAME, 'execute_tool') + span.set_attribute(GEN_AI_OPERATION_NAME, "execute_tool") span.set_attribute(GEN_AI_TOOL_DESCRIPTION, tool.description) span.set_attribute(GEN_AI_TOOL_NAME, tool.name) @@ -171,20 +173,20 @@ def trace_tool_call( # Setting empty llm request and response (as UI expect these) while not # applicable for tool_response. - span.set_attribute('gcp.vertex.agent.llm_request', '{}') - span.set_attribute('gcp.vertex.agent.llm_response', '{}') + span.set_attribute("gcp.vertex.agent.llm_request", EMPTY_JSON_STRING) + span.set_attribute("gcp.vertex.agent.llm_response", EMPTY_JSON_STRING) if _should_add_request_response_to_spans(): span.set_attribute( - 'gcp.vertex.agent.tool_call_args', + "gcp.vertex.agent.tool_call_args", _safe_json_serialize(args), ) else: - span.set_attribute('gcp.vertex.agent.tool_call_args', '{}') + span.set_attribute("gcp.vertex.agent.tool_call_args", EMPTY_JSON_STRING) # Tracing tool response - tool_call_id = '' - tool_response = '' + tool_call_id = "" + tool_response = "" if ( function_response_event is not None and function_response_event.content is not None @@ -201,16 +203,16 @@ def trace_tool_call( span.set_attribute(GEN_AI_TOOL_CALL_ID, tool_call_id) if not isinstance(tool_response, dict): - tool_response = {'result': tool_response} + tool_response = {"result": tool_response} if function_response_event is not None: - span.set_attribute('gcp.vertex.agent.event_id', function_response_event.id) + span.set_attribute("gcp.vertex.agent.event_id", function_response_event.id) if _should_add_request_response_to_spans(): span.set_attribute( - 'gcp.vertex.agent.tool_response', + "gcp.vertex.agent.tool_response", _safe_json_serialize(tool_response), ) else: - span.set_attribute('gcp.vertex.agent.tool_response', '{}') + span.set_attribute("gcp.vertex.agent.tool_response", EMPTY_JSON_STRING) def trace_merged_tool_calls( @@ -229,34 +231,34 @@ def trace_merged_tool_calls( span = trace.get_current_span() - span.set_attribute(GEN_AI_OPERATION_NAME, 'execute_tool') - span.set_attribute(GEN_AI_TOOL_NAME, '(merged tools)') - span.set_attribute(GEN_AI_TOOL_DESCRIPTION, '(merged tools)') + span.set_attribute(GEN_AI_OPERATION_NAME, "execute_tool") + span.set_attribute(GEN_AI_TOOL_NAME, "(merged tools)") + span.set_attribute(GEN_AI_TOOL_DESCRIPTION, "(merged tools)") span.set_attribute(GEN_AI_TOOL_CALL_ID, response_event_id) # TODO(b/441461932): See if these are still necessary - span.set_attribute('gcp.vertex.agent.tool_call_args', 'N/A') - span.set_attribute('gcp.vertex.agent.event_id', response_event_id) + span.set_attribute("gcp.vertex.agent.tool_call_args", "N/A") + span.set_attribute("gcp.vertex.agent.event_id", response_event_id) try: function_response_event_json = function_response_event.model_dumps_json( exclude_none=True ) except Exception: # pylint: disable=broad-exception-caught - function_response_event_json = '' + function_response_event_json = "" if _should_add_request_response_to_spans(): span.set_attribute( - 'gcp.vertex.agent.tool_response', + "gcp.vertex.agent.tool_response", function_response_event_json, ) else: - span.set_attribute('gcp.vertex.agent.tool_response', '{}') + span.set_attribute("gcp.vertex.agent.tool_response", EMPTY_JSON_STRING) # Setting empty llm request and response (as UI expect these) while not # applicable for tool_response. - span.set_attribute('gcp.vertex.agent.llm_request', '{}') + span.set_attribute("gcp.vertex.agent.llm_request", EMPTY_JSON_STRING) span.set_attribute( - 'gcp.vertex.agent.llm_response', - '{}', + "gcp.vertex.agent.llm_response", + EMPTY_JSON_STRING, ) @@ -281,57 +283,57 @@ def trace_call_llm( span = span or trace.get_current_span() # Special standard Open Telemetry GenaI attributes that indicate # that this is a span related to a Generative AI system. - span.set_attribute('gen_ai.system', 'gcp.vertex.agent') - span.set_attribute('gen_ai.request.model', llm_request.model) + span.set_attribute("gen_ai.system", "gcp.vertex.agent") + span.set_attribute("gen_ai.request.model", llm_request.model) span.set_attribute( - 'gcp.vertex.agent.invocation_id', invocation_context.invocation_id + "gcp.vertex.agent.invocation_id", invocation_context.invocation_id ) span.set_attribute( - 'gcp.vertex.agent.session_id', invocation_context.session.id + "gcp.vertex.agent.session_id", invocation_context.session.id ) - span.set_attribute('gcp.vertex.agent.event_id', event_id) + span.set_attribute("gcp.vertex.agent.event_id", event_id) # Consider removing once GenAI SDK provides a way to record this info. if _should_add_request_response_to_spans(): span.set_attribute( - 'gcp.vertex.agent.llm_request', + "gcp.vertex.agent.llm_request", _safe_json_serialize(_build_llm_request_for_trace(llm_request)), ) else: - span.set_attribute('gcp.vertex.agent.llm_request', '{}') + span.set_attribute("gcp.vertex.agent.llm_request", EMPTY_JSON_STRING) # Consider removing once GenAI SDK provides a way to record this info. if llm_request.config: if llm_request.config.top_p: span.set_attribute( - 'gen_ai.request.top_p', + "gen_ai.request.top_p", llm_request.config.top_p, ) if llm_request.config.max_output_tokens: span.set_attribute( - 'gen_ai.request.max_tokens', + "gen_ai.request.max_tokens", llm_request.config.max_output_tokens, ) try: llm_response_json = llm_response.model_dump_json(exclude_none=True) except Exception: # pylint: disable=broad-exception-caught - llm_response_json = '' + llm_response_json = "" if _should_add_request_response_to_spans(): span.set_attribute( - 'gcp.vertex.agent.llm_response', + "gcp.vertex.agent.llm_response", llm_response_json, ) else: - span.set_attribute('gcp.vertex.agent.llm_response', '{}') + span.set_attribute("gcp.vertex.agent.llm_response", EMPTY_JSON_STRING) if llm_response.usage_metadata is not None: span.set_attribute( - 'gen_ai.usage.input_tokens', + "gen_ai.usage.input_tokens", llm_response.usage_metadata.prompt_token_count, ) if llm_response.usage_metadata.candidates_token_count is not None: span.set_attribute( - 'gen_ai.usage.output_tokens', + "gen_ai.usage.output_tokens", llm_response.usage_metadata.candidates_token_count, ) if llm_response.finish_reason: @@ -340,7 +342,7 @@ def trace_call_llm( except AttributeError: finish_reason_str = str(llm_response.finish_reason).lower() span.set_attribute( - 'gen_ai.response.finish_reasons', + "gen_ai.response.finish_reasons", [finish_reason_str], ) @@ -362,23 +364,23 @@ def trace_send_data( """ span = trace.get_current_span() span.set_attribute( - 'gcp.vertex.agent.invocation_id', invocation_context.invocation_id + "gcp.vertex.agent.invocation_id", invocation_context.invocation_id ) - span.set_attribute('gcp.vertex.agent.event_id', event_id) + span.set_attribute("gcp.vertex.agent.event_id", event_id) # Once instrumentation is added to the GenAI SDK, consider whether this # information still needs to be recorded by the Agent Development Kit. if _should_add_request_response_to_spans(): span.set_attribute( - 'gcp.vertex.agent.data', + "gcp.vertex.agent.data", _safe_json_serialize([ types.Content(role=content.role, parts=content.parts).model_dump( - exclude_none=True, mode='json' + exclude_none=True, mode="json" ) for content in data ]), ) else: - span.set_attribute('gcp.vertex.agent.data', '{}') + span.set_attribute("gcp.vertex.agent.data", EMPTY_JSON_STRING) def _build_llm_request_for_trace(llm_request: LlmRequest) -> dict[str, Any]: @@ -396,18 +398,18 @@ def _build_llm_request_for_trace(llm_request: LlmRequest) -> dict[str, Any]: """ # Some fields in LlmRequest are function pointers and cannot be serialized. result = { - 'model': llm_request.model, - 'config': llm_request.config.model_dump( - exclude_none=True, exclude='response_schema', mode='json' + "model": llm_request.model, + "config": llm_request.config.model_dump( + exclude_none=True, exclude="response_schema", mode="json" ), - 'contents': [], + "contents": [], } # We do not want to send bytes data to the trace. for content in llm_request.contents: parts = [part for part in content.parts if not part.inline_data] - result['contents'].append( + result["contents"].append( types.Content(role=content.role, parts=parts).model_dump( - exclude_none=True, mode='json' + exclude_none=True, mode="json" ) ) return result @@ -419,8 +421,8 @@ def _build_llm_request_for_trace(llm_request: LlmRequest) -> dict[str, Any]: # to false. def _should_add_request_response_to_spans() -> bool: disabled_via_env_var = os.getenv( - ADK_CAPTURE_MESSAGE_CONTENT_IN_SPANS, 'true' - ).lower() in ('false', '0') + ADK_CAPTURE_MESSAGE_CONTENT_IN_SPANS, "true" + ).lower() in ("false", "0") return not disabled_via_env_var @@ -438,7 +440,7 @@ def use_generate_content_span( common_attributes = { GEN_AI_CONVERSATION_ID: invocation_context.session.id, - 'gcp.vertex.agent.event_id': model_response_event.id, + "gcp.vertex.agent.event_id": model_response_event.id, } if ( _is_gemini_agent(invocation_context.agent) @@ -455,8 +457,8 @@ def use_generate_content_span( def _should_log_prompt_response_content() -> bool: return os.getenv( - OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT, '' - ).lower() in ('1', 'true') + OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT, "" + ).lower() in ("1", "true") def _serialize_content(content: types.ContentUnion) -> AnyValue: @@ -481,10 +483,10 @@ def _serialize_content_with_elision( def _instrumented_with_opentelemetry_instrumentation_google_genai() -> bool: maybe_wrapped_function = Models.generate_content - print(f'{Models.generate_content.__code__.co_filename=}') - while wrapped := getattr(maybe_wrapped_function, '__wrapped__', None): + print(f"{Models.generate_content.__code__.co_filename=}") + while wrapped := getattr(maybe_wrapped_function, "__wrapped__", None): if ( - 'opentelemetry/instrumentation/google_genai' + "opentelemetry/instrumentation/google_genai" in maybe_wrapped_function.__code__.co_filename ): return True @@ -516,15 +518,15 @@ def _use_native_generate_content_span( f"generate_content {llm_request.model or ''}" ) as span: span.set_attribute(GEN_AI_SYSTEM, _guess_gemini_system_name()) - span.set_attribute(GEN_AI_OPERATION_NAME, 'generate_content') - span.set_attribute(GEN_AI_REQUEST_MODEL, llm_request.model or '') + span.set_attribute(GEN_AI_OPERATION_NAME, "generate_content") + span.set_attribute(GEN_AI_REQUEST_MODEL, llm_request.model or "") span.set_attributes(common_attributes) otel_logger.emit( LogRecord( - event_name='gen_ai.system.message', + event_name="gen_ai.system.message", body={ - 'content': _serialize_content_with_elision( + "content": _serialize_content_with_elision( llm_request.config.system_instruction ) }, @@ -535,8 +537,8 @@ def _use_native_generate_content_span( for content in llm_request.contents: otel_logger.emit( LogRecord( - event_name='gen_ai.user.message', - body={'content': _serialize_content_with_elision(content)}, + event_name="gen_ai.user.message", + body={"content": _serialize_content_with_elision(content)}, attributes={GEN_AI_SYSTEM: _guess_gemini_system_name()}, ) ) @@ -567,12 +569,12 @@ def trace_generate_content_result(span: Span | None, llm_response: LlmResponse): otel_logger.emit( LogRecord( - event_name='gen_ai.choice', + event_name="gen_ai.choice", body={ - 'content': _serialize_content_with_elision(llm_response.content), - 'index': 0, # ADK always returns a single candidate + "content": _serialize_content_with_elision(llm_response.content), + "index": 0, # ADK always returns a single candidate } - | {'finish_reason': llm_response.finish_reason.value} + | {"finish_reason": llm_response.finish_reason.value} if llm_response.finish_reason is not None else {}, attributes={GEN_AI_SYSTEM: _guess_gemini_system_name()}, @@ -583,6 +585,6 @@ def trace_generate_content_result(span: Span | None, llm_response: LlmResponse): def _guess_gemini_system_name() -> str: return ( GenAiSystemValues.VERTEX_AI.name.lower() - if os.getenv('GOOGLE_GENAI_USE_VERTEXAI', '').lower() in ('true', '1') + if os.getenv("GOOGLE_GENAI_USE_VERTEXAI", "").lower() in ("true", "1") else GenAiSystemValues.GEMINI.name.lower() ) diff --git a/tests/unittests/telemetry/test_spans.py b/tests/unittests/telemetry/test_spans.py index ed7a588b28..7718c89024 100644 --- a/tests/unittests/telemetry/test_spans.py +++ b/tests/unittests/telemetry/test_spans.py @@ -93,6 +93,20 @@ async def _create_invocation_context( return invocation_context +def _assert_span_attribute_set_to_empty_json(mock_span, attribute_name: str): + """Helper to assert span attribute is set to empty JSON string '{}'.""" + calls = [ + call + for call in mock_span.set_attribute.call_args_list + if call.args[0] == attribute_name + ] + assert len(calls) == 1, f"Expected '{attribute_name}' to be set exactly once" + assert calls[0].args[1] == '{}', ( + f"Expected JSON string '{{}}' for {attribute_name} when content capture" + f' is disabled, got {calls[0].args[1]!r}' + ) + + @pytest.mark.asyncio async def test_trace_agent_invocation(mock_span_fixture): """Test trace_agent_invocation sets span attributes correctly.""" @@ -484,20 +498,12 @@ async def test_call_llm_disabling_request_response_content( # Act trace_call_llm(invocation_context, 'test_event_id', llm_request, llm_response) - # Assert - assert ( - 'gcp.vertex.agent.llm_request', - '{}', - ) in ( - call_obj.args - for call_obj in mock_span_fixture.set_attribute.call_args_list + # Assert - Check attributes are set to JSON string '{}' not dict {} + _assert_span_attribute_set_to_empty_json( + mock_span_fixture, 'gcp.vertex.agent.llm_request' ) - assert ( - 'gcp.vertex.agent.llm_response', - '{}', - ) in ( - call_obj.args - for call_obj in mock_span_fixture.set_attribute.call_args_list + _assert_span_attribute_set_to_empty_json( + mock_span_fixture, 'gcp.vertex.agent.llm_response' ) @@ -543,20 +549,12 @@ def test_trace_tool_call_disabling_request_response_content( function_response_event=mock_event_fixture, ) - # Assert - assert ( - 'gcp.vertex.agent.tool_call_args', - '{}', - ) in ( - call_obj.args - for call_obj in mock_span_fixture.set_attribute.call_args_list + # Assert - Check attributes are set to JSON string '{}' not dict {} + _assert_span_attribute_set_to_empty_json( + mock_span_fixture, 'gcp.vertex.agent.tool_call_args' ) - assert ( - 'gcp.vertex.agent.tool_response', - '{}', - ) in ( - call_obj.args - for call_obj in mock_span_fixture.set_attribute.call_args_list + _assert_span_attribute_set_to_empty_json( + mock_span_fixture, 'gcp.vertex.agent.tool_response' ) @@ -584,13 +582,9 @@ def test_trace_merged_tool_disabling_request_response_content( function_response_event=mock_event_fixture, ) - # Assert - assert ( - 'gcp.vertex.agent.tool_response', - '{}', - ) in ( - call_obj.args - for call_obj in mock_span_fixture.set_attribute.call_args_list + # Assert - Check attribute is set to JSON string '{}' not dict {} + _assert_span_attribute_set_to_empty_json( + mock_span_fixture, 'gcp.vertex.agent.tool_response' )