From f762e290b987055951ec73f1c3debd84284e0b77 Mon Sep 17 00:00:00 2001 From: DineshThumma9 Date: Thu, 8 Jan 2026 21:02:09 +0530 Subject: [PATCH 1/5] fix:Use JSON string instead of dic for span attributes when content capture is disabled --- tests/unittests/telemetry/test_spans.py | 83 ++++++++++++++++--------- 1 file changed, 55 insertions(+), 28 deletions(-) diff --git a/tests/unittests/telemetry/test_spans.py b/tests/unittests/telemetry/test_spans.py index ed7a588b28..2d5ac2ecb7 100644 --- a/tests/unittests/telemetry/test_spans.py +++ b/tests/unittests/telemetry/test_spans.py @@ -484,20 +484,31 @@ async def test_call_llm_disabling_request_response_content( # Act trace_call_llm(invocation_context, 'test_event_id', llm_request, llm_response) - # Assert + # Assert - Check attributes are set to JSON string '{}' not dict {} + llm_request_calls = [ + call + for call in mock_span_fixture.set_attribute.call_args_list + if call.args[0] == 'gcp.vertex.agent.llm_request' + ] assert ( - 'gcp.vertex.agent.llm_request', - '{}', - ) in ( - call_obj.args - for call_obj in mock_span_fixture.set_attribute.call_args_list + len(llm_request_calls) == 1 + ), "Expected 'gcp.vertex.agent.llm_request' to be set exactly once" + assert llm_request_calls[0].args[1] == '{}', ( + "Expected JSON string '{}' for llm_request when content capture is" + f' disabled, got {llm_request_calls[0].args[1]!r}' ) + + llm_response_calls = [ + call + for call in mock_span_fixture.set_attribute.call_args_list + if call.args[0] == 'gcp.vertex.agent.llm_response' + ] assert ( - 'gcp.vertex.agent.llm_response', - '{}', - ) in ( - call_obj.args - for call_obj in mock_span_fixture.set_attribute.call_args_list + len(llm_response_calls) == 1 + ), "Expected 'gcp.vertex.agent.llm_response' to be set exactly once" + assert llm_response_calls[0].args[1] == '{}', ( + "Expected JSON string '{}' for llm_response when content capture is" + f' disabled, got {llm_response_calls[0].args[1]!r}' ) @@ -543,20 +554,31 @@ def test_trace_tool_call_disabling_request_response_content( function_response_event=mock_event_fixture, ) - # Assert + # Assert - Check attributes are set to JSON string '{}' not dict {} + tool_args_calls = [ + call + for call in mock_span_fixture.set_attribute.call_args_list + if call.args[0] == 'gcp.vertex.agent.tool_call_args' + ] assert ( - 'gcp.vertex.agent.tool_call_args', - '{}', - ) in ( - call_obj.args - for call_obj in mock_span_fixture.set_attribute.call_args_list + len(tool_args_calls) == 1 + ), "Expected 'gcp.vertex.agent.tool_call_args' to be set exactly once" + assert tool_args_calls[0].args[1] == '{}', ( + "Expected JSON string '{}' for tool_call_args when content capture is" + f' disabled, got {tool_args_calls[0].args[1]!r}' ) + + tool_response_calls = [ + call + for call in mock_span_fixture.set_attribute.call_args_list + if call.args[0] == 'gcp.vertex.agent.tool_response' + ] assert ( - 'gcp.vertex.agent.tool_response', - '{}', - ) in ( - call_obj.args - for call_obj in mock_span_fixture.set_attribute.call_args_list + len(tool_response_calls) == 1 + ), "Expected 'gcp.vertex.agent.tool_response' to be set exactly once" + assert tool_response_calls[0].args[1] == '{}', ( + "Expected JSON string '{}' for tool_response when content capture is" + f' disabled, got {tool_response_calls[0].args[1]!r}' ) @@ -584,13 +606,18 @@ def test_trace_merged_tool_disabling_request_response_content( function_response_event=mock_event_fixture, ) - # Assert + # Assert - Check attribute is set to JSON string '{}' not dict {} + tool_response_calls = [ + call + for call in mock_span_fixture.set_attribute.call_args_list + if call.args[0] == 'gcp.vertex.agent.tool_response' + ] assert ( - 'gcp.vertex.agent.tool_response', - '{}', - ) in ( - call_obj.args - for call_obj in mock_span_fixture.set_attribute.call_args_list + len(tool_response_calls) == 1 + ), "Expected 'gcp.vertex.agent.tool_response' to be set exactly once" + assert tool_response_calls[0].args[1] == '{}', ( + "Expected JSON string '{}' for tool_response when content capture is" + f' disabled, got {tool_response_calls[0].args[1]!r}' ) From 13f75e94ad964553f5d0a48ce5cf0a31e90be510 Mon Sep 17 00:00:00 2001 From: DineshThumma9 Date: Thu, 8 Jan 2026 22:17:05 +0530 Subject: [PATCH 2/5] Trigger CLA Check From 62e463be83ed7f2987a31c8477bdde7f7667933a Mon Sep 17 00:00:00 2001 From: DineshThumma9 Date: Thu, 8 Jan 2026 22:35:24 +0530 Subject: [PATCH 3/5] refactored:Added helper function and improved maintainability --- src/google/adk/telemetry/tracing.py | 812 ++++++++++++------------ tests/unittests/telemetry/test_spans.py | 83 +-- 2 files changed, 446 insertions(+), 449 deletions(-) diff --git a/src/google/adk/telemetry/tracing.py b/src/google/adk/telemetry/tracing.py index d54f75173b..31bb3605b2 100644 --- a/src/google/adk/telemetry/tracing.py +++ b/src/google/adk/telemetry/tracing.py @@ -37,20 +37,46 @@ from opentelemetry import _logs from opentelemetry import trace from opentelemetry._logs import LogRecord -from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import GEN_AI_AGENT_DESCRIPTION -from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import GEN_AI_AGENT_NAME -from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import GEN_AI_CONVERSATION_ID -from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import GEN_AI_OPERATION_NAME -from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import GEN_AI_REQUEST_MODEL -from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import GEN_AI_RESPONSE_FINISH_REASONS +from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import ( + GEN_AI_AGENT_DESCRIPTION, +) +from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import ( + GEN_AI_AGENT_NAME, +) +from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import ( + GEN_AI_CONVERSATION_ID, +) +from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import ( + GEN_AI_OPERATION_NAME, +) +from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import ( + GEN_AI_REQUEST_MODEL, +) +from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import ( + GEN_AI_RESPONSE_FINISH_REASONS, +) from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import GEN_AI_SYSTEM -from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import GEN_AI_TOOL_CALL_ID -from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import GEN_AI_TOOL_DESCRIPTION -from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import GEN_AI_TOOL_NAME -from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import GEN_AI_TOOL_TYPE -from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import GEN_AI_USAGE_INPUT_TOKENS -from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import GEN_AI_USAGE_OUTPUT_TOKENS -from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import GenAiSystemValues +from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import ( + GEN_AI_TOOL_CALL_ID, +) +from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import ( + GEN_AI_TOOL_DESCRIPTION, +) +from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import ( + GEN_AI_TOOL_NAME, +) +from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import ( + GEN_AI_TOOL_TYPE, +) +from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import ( + GEN_AI_USAGE_INPUT_TOKENS, +) +from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import ( + GEN_AI_USAGE_OUTPUT_TOKENS, +) +from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import ( + GenAiSystemValues, +) from opentelemetry.semconv.schemas import Schemas from opentelemetry.trace import Span from opentelemetry.util.types import AnyValue @@ -62,89 +88,91 @@ # 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: - from ..agents.base_agent import BaseAgent - from ..agents.invocation_context import InvocationContext - from ..events.event import Event - from ..models.llm_request import LlmRequest - from ..models.llm_response import LlmResponse - from ..tools.base_tool import BaseTool + from ..agents.base_agent import BaseAgent + from ..agents.invocation_context import InvocationContext + from ..events.event import Event + from ..models.llm_request import LlmRequest + from ..models.llm_response import LlmResponse + 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: - """Convert any Python object to a JSON-serializable type or string. + """Convert any Python object to a JSON-serializable type or string. - Args: - obj: The object to serialize. + Args: + obj: The object to serialize. - Returns: - The JSON-serialized object string or if the object cannot be serialized. - """ + Returns: + The JSON-serialized object string or if the object cannot be serialized. + """ - try: - # Try direct JSON serialization first - return json.dumps( - obj, ensure_ascii=False, default=lambda o: '' - ) - except (TypeError, OverflowError): - return '' + try: + # Try direct JSON serialization first + return json.dumps( + obj, ensure_ascii=False, default=lambda o: "" + ) + except (TypeError, OverflowError): + return "" def trace_agent_invocation( span: trace.Span, agent: BaseAgent, ctx: InvocationContext ) -> None: - """Sets span attributes immediately available on agent invocation according to OTEL semconv version 1.37. + """Sets span attributes immediately available on agent invocation according to OTEL semconv version 1.37. - Args: - span: Span on which attributes are set. - agent: Agent from which attributes are gathered. - ctx: InvocationContext from which attributes are gathered. + Args: + span: Span on which attributes are set. + agent: Agent from which attributes are gathered. + ctx: InvocationContext from which attributes are gathered. - Inference related fields are not set, due to their planned removal from invoke_agent span: - https://github.com/open-telemetry/semantic-conventions/issues/2632 + Inference related fields are not set, due to their planned removal from invoke_agent span: + https://github.com/open-telemetry/semantic-conventions/issues/2632 - `gen_ai.agent.id` is not set because currently it's unclear what attributes this field should have, specifically: - - In which scope should it be unique (globally, given project, given agentic flow, given deployment). - - Should it be unchanging between deployments, and how this should this be achieved. + `gen_ai.agent.id` is not set because currently it's unclear what attributes this field should have, specifically: + - In which scope should it be unique (globally, given project, given agentic flow, given deployment). + - Should it be unchanging between deployments, and how this should this be achieved. - `gen_ai.data_source.id` is not set because it's not available. - Closest type which could contain this information is types.GroundingMetadata, which does not have an ID. + `gen_ai.data_source.id` is not set because it's not available. + Closest type which could contain this information is types.GroundingMetadata, which does not have an ID. - `server.*` attributes are not set pending confirmation from aabmass. - """ + `server.*` attributes are not set pending confirmation from aabmass. + """ - # Required - span.set_attribute(GEN_AI_OPERATION_NAME, 'invoke_agent') + # Required + span.set_attribute(GEN_AI_OPERATION_NAME, "invoke_agent") - # Conditionally Required - span.set_attribute(GEN_AI_AGENT_DESCRIPTION, agent.description) + # Conditionally Required + span.set_attribute(GEN_AI_AGENT_DESCRIPTION, agent.description) - span.set_attribute(GEN_AI_AGENT_NAME, agent.name) - span.set_attribute(GEN_AI_CONVERSATION_ID, ctx.session.id) + span.set_attribute(GEN_AI_AGENT_NAME, agent.name) + span.set_attribute(GEN_AI_CONVERSATION_ID, ctx.session.id) def trace_tool_call( @@ -152,112 +180,112 @@ def trace_tool_call( args: dict[str, Any], function_response_event: Event | None, ): - """Traces tool call. + """Traces tool call. - Args: - tool: The tool that was called. - args: The arguments to the tool call. - function_response_event: The event with the function response details. - """ - span = trace.get_current_span() + Args: + tool: The tool that was called. + args: The arguments to the tool call. + function_response_event: The event with the function response details. + """ + 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) + span.set_attribute(GEN_AI_TOOL_DESCRIPTION, tool.description) + span.set_attribute(GEN_AI_TOOL_NAME, tool.name) - # e.g. FunctionTool - span.set_attribute(GEN_AI_TOOL_TYPE, tool.__class__.__name__) + # e.g. FunctionTool + span.set_attribute(GEN_AI_TOOL_TYPE, tool.__class__.__name__) - # 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', '{}') + # Setting empty llm request and response (as UI expect these) while not + # applicable for tool_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', - _safe_json_serialize(args), - ) - else: - span.set_attribute('gcp.vertex.agent.tool_call_args', '{}') - - # Tracing tool response - tool_call_id = '' - tool_response = '' - if ( - function_response_event is not None - and function_response_event.content is not None - and function_response_event.content.parts - ): - response_parts = function_response_event.content.parts - function_response = response_parts[0].function_response - if function_response is not None: - if function_response.id is not None: - tool_call_id = function_response.id - if function_response.response is not None: - tool_response = function_response.response - - span.set_attribute(GEN_AI_TOOL_CALL_ID, tool_call_id) - - if not isinstance(tool_response, dict): - tool_response = {'result': tool_response} - if function_response_event is not None: - 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', - _safe_json_serialize(tool_response), - ) - else: - span.set_attribute('gcp.vertex.agent.tool_response', '{}') + if _should_add_request_response_to_spans(): + span.set_attribute( + "gcp.vertex.agent.tool_call_args", + _safe_json_serialize(args), + ) + else: + span.set_attribute("gcp.vertex.agent.tool_call_args", EMPTY_JSON_STRING) + + # Tracing tool response + tool_call_id = "" + tool_response = "" + if ( + function_response_event is not None + and function_response_event.content is not None + and function_response_event.content.parts + ): + response_parts = function_response_event.content.parts + function_response = response_parts[0].function_response + if function_response is not None: + if function_response.id is not None: + tool_call_id = function_response.id + if function_response.response is not None: + tool_response = function_response.response + + span.set_attribute(GEN_AI_TOOL_CALL_ID, tool_call_id) + + if not isinstance(tool_response, dict): + tool_response = {"result": tool_response} + if function_response_event is not None: + 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", + _safe_json_serialize(tool_response), + ) + else: + span.set_attribute("gcp.vertex.agent.tool_response", EMPTY_JSON_STRING) def trace_merged_tool_calls( response_event_id: str, function_response_event: Event, ): - """Traces merged tool call events. + """Traces merged tool call events. - Calling this function is not needed for telemetry purposes. This is provided - for preventing /debug/trace requests (typically sent by web UI). + Calling this function is not needed for telemetry purposes. This is provided + for preventing /debug/trace requests (typically sent by web UI). - Args: - response_event_id: The ID of the response event. - function_response_event: The merged response event. - """ + Args: + response_event_id: The ID of the response event. + function_response_event: The merged response event. + """ - span = trace.get_current_span() + 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_TOOL_CALL_ID, response_event_id) + 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) - 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 = '' + # 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) + 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 = "" - if _should_add_request_response_to_spans(): + if _should_add_request_response_to_spans(): + span.set_attribute( + "gcp.vertex.agent.tool_response", + function_response_event_json, + ) + else: + 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", EMPTY_JSON_STRING) span.set_attribute( - 'gcp.vertex.agent.tool_response', - function_response_event_json, + "gcp.vertex.agent.llm_response", + EMPTY_JSON_STRING, ) - else: - span.set_attribute('gcp.vertex.agent.tool_response', '{}') - # 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', - '{}', - ) def trace_call_llm( @@ -267,82 +295,80 @@ def trace_call_llm( llm_response: LlmResponse, span: Span | None = None, ): - """Traces a call to the LLM. - - This function records details about the LLM request and response as - attributes on the current OpenTelemetry span. - - Args: - invocation_context: The invocation context for the current agent run. - event_id: The ID of the event. - llm_request: The LLM request object. - llm_response: The LLM response object. - """ - 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( - 'gcp.vertex.agent.invocation_id', invocation_context.invocation_id - ) - span.set_attribute( - 'gcp.vertex.agent.session_id', invocation_context.session.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', - _safe_json_serialize(_build_llm_request_for_trace(llm_request)), - ) - else: - span.set_attribute('gcp.vertex.agent.llm_request', '{}') - # 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', - llm_request.config.top_p, - ) - if llm_request.config.max_output_tokens: - span.set_attribute( - '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 = '' - - if _should_add_request_response_to_spans(): + """Traces a call to the LLM. + + This function records details about the LLM request and response as + attributes on the current OpenTelemetry span. + + Args: + invocation_context: The invocation context for the current agent run. + event_id: The ID of the event. + llm_request: The LLM request object. + llm_response: The LLM response object. + """ + 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( - 'gcp.vertex.agent.llm_response', - llm_response_json, + "gcp.vertex.agent.invocation_id", invocation_context.invocation_id ) - else: - span.set_attribute('gcp.vertex.agent.llm_response', '{}') + span.set_attribute("gcp.vertex.agent.session_id", invocation_context.session.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", + _safe_json_serialize(_build_llm_request_for_trace(llm_request)), + ) + else: + 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", + llm_request.config.top_p, + ) + if llm_request.config.max_output_tokens: + span.set_attribute( + "gen_ai.request.max_tokens", + llm_request.config.max_output_tokens, + ) - if llm_response.usage_metadata is not None: - span.set_attribute( - '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', - llm_response.usage_metadata.candidates_token_count, - ) - if llm_response.finish_reason: try: - finish_reason_str = llm_response.finish_reason.value.lower() - except AttributeError: - finish_reason_str = str(llm_response.finish_reason).lower() - span.set_attribute( - 'gen_ai.response.finish_reasons', - [finish_reason_str], - ) + llm_response_json = llm_response.model_dump_json(exclude_none=True) + except Exception: # pylint: disable=broad-exception-caught + llm_response_json = "" + + if _should_add_request_response_to_spans(): + span.set_attribute( + "gcp.vertex.agent.llm_response", + llm_response_json, + ) + else: + 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", + 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", + llm_response.usage_metadata.candidates_token_count, + ) + if llm_response.finish_reason: + try: + finish_reason_str = llm_response.finish_reason.value.lower() + except AttributeError: + finish_reason_str = str(llm_response.finish_reason).lower() + span.set_attribute( + "gen_ai.response.finish_reasons", + [finish_reason_str], + ) def trace_send_data( @@ -350,67 +376,69 @@ def trace_send_data( event_id: str, data: list[types.Content], ): - """Traces the sending of data to the agent. - - This function records details about the data sent to the agent as - attributes on the current OpenTelemetry span. - - Args: - invocation_context: The invocation context for the current agent run. - event_id: The ID of the event. - data: A list of content objects. - """ - span = trace.get_current_span() - span.set_attribute( - 'gcp.vertex.agent.invocation_id', invocation_context.invocation_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(): + """Traces the sending of data to the agent. + + This function records details about the data sent to the agent as + attributes on the current OpenTelemetry span. + + Args: + invocation_context: The invocation context for the current agent run. + event_id: The ID of the event. + data: A list of content objects. + """ + span = trace.get_current_span() span.set_attribute( - 'gcp.vertex.agent.data', - _safe_json_serialize([ - types.Content(role=content.role, parts=content.parts).model_dump( - exclude_none=True, mode='json' - ) - for content in data - ]), + "gcp.vertex.agent.invocation_id", invocation_context.invocation_id ) - else: - span.set_attribute('gcp.vertex.agent.data', '{}') + 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", + _safe_json_serialize( + [ + types.Content(role=content.role, parts=content.parts).model_dump( + exclude_none=True, mode="json" + ) + for content in data + ] + ), + ) + else: + span.set_attribute("gcp.vertex.agent.data", EMPTY_JSON_STRING) def _build_llm_request_for_trace(llm_request: LlmRequest) -> dict[str, Any]: - """Builds a dictionary representation of the LLM request for tracing. - - This function prepares a dictionary representation of the LlmRequest - object, suitable for inclusion in a trace. It excludes fields that cannot - be serialized (e.g., function pointers) and avoids sending bytes data. - - Args: - llm_request: The LlmRequest object. - - Returns: - A dictionary representation of the LLM request. - """ - # 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' - ), - '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( - types.Content(role=content.role, parts=parts).model_dump( - exclude_none=True, mode='json' + """Builds a dictionary representation of the LLM request for tracing. + + This function prepares a dictionary representation of the LlmRequest + object, suitable for inclusion in a trace. It excludes fields that cannot + be serialized (e.g., function pointers) and avoids sending bytes data. + + Args: + llm_request: The LlmRequest object. + + Returns: + A dictionary representation of the LLM request. + """ + # 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" + ), + "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( + types.Content(role=content.role, parts=parts).model_dump( + exclude_none=True, mode="json" + ) ) - ) - return result + return result # Defaults to true for now to preserve backward compatibility. @@ -418,10 +446,10 @@ def _build_llm_request_for_trace(llm_request: LlmRequest) -> dict[str, Any]: # a deprecation of request/response content in spans by switching the default # 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') - return not disabled_via_env_var + disabled_via_env_var = os.getenv( + ADK_CAPTURE_MESSAGE_CONTENT_IN_SPANS, "true" + ).lower() in ("false", "0") + return not disabled_via_env_var @contextmanager @@ -430,81 +458,81 @@ def use_generate_content_span( invocation_context: InvocationContext, model_response_event: Event, ) -> Iterator[Span | None]: - """Context manager encompassing `generate_content {model.name}` span. - - When an external library for inference instrumentation is installed (e.g. opentelemetry-instrumentation-google-genai), - span creation is delegated to said library. - """ - - common_attributes = { - GEN_AI_CONVERSATION_ID: invocation_context.session.id, - 'gcp.vertex.agent.event_id': model_response_event.id, - } - if ( - _is_gemini_agent(invocation_context.agent) - and _instrumented_with_opentelemetry_instrumentation_google_genai() - ): - yield None - else: - with _use_native_generate_content_span( - llm_request=llm_request, - common_attributes=common_attributes, - ) as span: - yield span + """Context manager encompassing `generate_content {model.name}` span. + + When an external library for inference instrumentation is installed (e.g. opentelemetry-instrumentation-google-genai), + span creation is delegated to said library. + """ + + common_attributes = { + GEN_AI_CONVERSATION_ID: invocation_context.session.id, + "gcp.vertex.agent.event_id": model_response_event.id, + } + if ( + _is_gemini_agent(invocation_context.agent) + and _instrumented_with_opentelemetry_instrumentation_google_genai() + ): + yield None + else: + with _use_native_generate_content_span( + llm_request=llm_request, + common_attributes=common_attributes, + ) as span: + yield span def _should_log_prompt_response_content() -> bool: - return os.getenv( - OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT, '' - ).lower() in ('1', 'true') + return os.getenv( + OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT, "" + ).lower() in ("1", "true") def _serialize_content(content: types.ContentUnion) -> AnyValue: - if isinstance(content, BaseModel): - return content.model_dump() - if isinstance(content, str): - return content - if isinstance(content, list): - return [_serialize_content(part) for part in content] - return _safe_json_serialize(content) + if isinstance(content, BaseModel): + return content.model_dump() + if isinstance(content, str): + return content + if isinstance(content, list): + return [_serialize_content(part) for part in content] + return _safe_json_serialize(content) def _serialize_content_with_elision( content: types.ContentUnion | None, ) -> AnyValue: - if not _should_log_prompt_response_content(): - return USER_CONTENT_ELIDED - if content is None: - return None - return _serialize_content(content) + if not _should_log_prompt_response_content(): + return USER_CONTENT_ELIDED + if content is None: + return None + return _serialize_content(content) 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): - if ( - 'opentelemetry/instrumentation/google_genai' - in maybe_wrapped_function.__code__.co_filename - ): - return True - maybe_wrapped_function = wrapped # pyright: ignore[reportAny] + maybe_wrapped_function = Models.generate_content + print(f"{Models.generate_content.__code__.co_filename=}") + while wrapped := getattr(maybe_wrapped_function, "__wrapped__", None): + if ( + "opentelemetry/instrumentation/google_genai" + in maybe_wrapped_function.__code__.co_filename + ): + return True + maybe_wrapped_function = wrapped # pyright: ignore[reportAny] - return False + return False def _is_gemini_agent(agent: BaseAgent) -> bool: - from ..agents.llm_agent import LlmAgent + from ..agents.llm_agent import LlmAgent - if not isinstance(agent, LlmAgent): - return False + if not isinstance(agent, LlmAgent): + return False - if isinstance(agent.model, str): - return is_gemini_model(agent.model) + if isinstance(agent.model, str): + return is_gemini_model(agent.model) - from ..models.google_llm import Gemini + from ..models.google_llm import Gemini - return isinstance(agent.model, Gemini) + return isinstance(agent.model, Gemini) @contextmanager @@ -512,77 +540,77 @@ def _use_native_generate_content_span( llm_request: LlmRequest, common_attributes: Mapping[str, AttributeValue], ) -> Iterator[Span]: - with tracer.start_as_current_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_attributes(common_attributes) + with tracer.start_as_current_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_attributes(common_attributes) + + otel_logger.emit( + LogRecord( + event_name="gen_ai.system.message", + body={ + "content": _serialize_content_with_elision( + llm_request.config.system_instruction + ) + }, + attributes={GEN_AI_SYSTEM: _guess_gemini_system_name()}, + ) + ) + + for content in llm_request.contents: + otel_logger.emit( + LogRecord( + event_name="gen_ai.user.message", + body={"content": _serialize_content_with_elision(content)}, + attributes={GEN_AI_SYSTEM: _guess_gemini_system_name()}, + ) + ) + + yield span + + +def trace_generate_content_result(span: Span | None, llm_response: LlmResponse): + """Trace result of the inference in generate_content span.""" + + if span is None: + return + + if llm_response.partial: + return + + if finish_reason := llm_response.finish_reason: + span.set_attribute(GEN_AI_RESPONSE_FINISH_REASONS, [finish_reason.lower()]) + if usage_metadata := llm_response.usage_metadata: + if usage_metadata.prompt_token_count is not None: + span.set_attribute( + GEN_AI_USAGE_INPUT_TOKENS, usage_metadata.prompt_token_count + ) + if usage_metadata.candidates_token_count is not None: + span.set_attribute( + GEN_AI_USAGE_OUTPUT_TOKENS, usage_metadata.candidates_token_count + ) otel_logger.emit( LogRecord( - event_name='gen_ai.system.message', + event_name="gen_ai.choice", body={ - 'content': _serialize_content_with_elision( - llm_request.config.system_instruction - ) - }, + "content": _serialize_content_with_elision(llm_response.content), + "index": 0, # ADK always returns a single candidate + } + | {"finish_reason": llm_response.finish_reason.value} + if llm_response.finish_reason is not None + else {}, attributes={GEN_AI_SYSTEM: _guess_gemini_system_name()}, ) ) - for content in llm_request.contents: - otel_logger.emit( - LogRecord( - event_name='gen_ai.user.message', - body={'content': _serialize_content_with_elision(content)}, - attributes={GEN_AI_SYSTEM: _guess_gemini_system_name()}, - ) - ) - - yield span - - -def trace_generate_content_result(span: Span | None, llm_response: LlmResponse): - """Trace result of the inference in generate_content span.""" - - if span is None: - return - - if llm_response.partial: - return - - if finish_reason := llm_response.finish_reason: - span.set_attribute(GEN_AI_RESPONSE_FINISH_REASONS, [finish_reason.lower()]) - if usage_metadata := llm_response.usage_metadata: - if usage_metadata.prompt_token_count is not None: - span.set_attribute( - GEN_AI_USAGE_INPUT_TOKENS, usage_metadata.prompt_token_count - ) - if usage_metadata.candidates_token_count is not None: - span.set_attribute( - GEN_AI_USAGE_OUTPUT_TOKENS, usage_metadata.candidates_token_count - ) - - otel_logger.emit( - LogRecord( - event_name='gen_ai.choice', - body={ - 'content': _serialize_content_with_elision(llm_response.content), - 'index': 0, # ADK always returns a single candidate - } - | {'finish_reason': llm_response.finish_reason.value} - if llm_response.finish_reason is not None - else {}, - attributes={GEN_AI_SYSTEM: _guess_gemini_system_name()}, - ) - ) - def _guess_gemini_system_name() -> str: - return ( - GenAiSystemValues.VERTEX_AI.name.lower() - if os.getenv('GOOGLE_GENAI_USE_VERTEXAI', '').lower() in ('true', '1') - else GenAiSystemValues.GEMINI.name.lower() - ) + return ( + GenAiSystemValues.VERTEX_AI.name.lower() + 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 2d5ac2ecb7..e1fa144bf5 100644 --- a/tests/unittests/telemetry/test_spans.py +++ b/tests/unittests/telemetry/test_spans.py @@ -93,6 +93,22 @@ 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 is" + f' 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.""" @@ -485,30 +501,11 @@ async def test_call_llm_disabling_request_response_content( trace_call_llm(invocation_context, 'test_event_id', llm_request, llm_response) # Assert - Check attributes are set to JSON string '{}' not dict {} - llm_request_calls = [ - call - for call in mock_span_fixture.set_attribute.call_args_list - if call.args[0] == 'gcp.vertex.agent.llm_request' - ] - assert ( - len(llm_request_calls) == 1 - ), "Expected 'gcp.vertex.agent.llm_request' to be set exactly once" - assert llm_request_calls[0].args[1] == '{}', ( - "Expected JSON string '{}' for llm_request when content capture is" - f' disabled, got {llm_request_calls[0].args[1]!r}' + _assert_span_attribute_set_to_empty_json( + mock_span_fixture, 'gcp.vertex.agent.llm_request' ) - - llm_response_calls = [ - call - for call in mock_span_fixture.set_attribute.call_args_list - if call.args[0] == 'gcp.vertex.agent.llm_response' - ] - assert ( - len(llm_response_calls) == 1 - ), "Expected 'gcp.vertex.agent.llm_response' to be set exactly once" - assert llm_response_calls[0].args[1] == '{}', ( - "Expected JSON string '{}' for llm_response when content capture is" - f' disabled, got {llm_response_calls[0].args[1]!r}' + _assert_span_attribute_set_to_empty_json( + mock_span_fixture, 'gcp.vertex.agent.llm_response' ) @@ -555,30 +552,11 @@ def test_trace_tool_call_disabling_request_response_content( ) # Assert - Check attributes are set to JSON string '{}' not dict {} - tool_args_calls = [ - call - for call in mock_span_fixture.set_attribute.call_args_list - if call.args[0] == 'gcp.vertex.agent.tool_call_args' - ] - assert ( - len(tool_args_calls) == 1 - ), "Expected 'gcp.vertex.agent.tool_call_args' to be set exactly once" - assert tool_args_calls[0].args[1] == '{}', ( - "Expected JSON string '{}' for tool_call_args when content capture is" - f' disabled, got {tool_args_calls[0].args[1]!r}' + _assert_span_attribute_set_to_empty_json( + mock_span_fixture, 'gcp.vertex.agent.tool_call_args' ) - - tool_response_calls = [ - call - for call in mock_span_fixture.set_attribute.call_args_list - if call.args[0] == 'gcp.vertex.agent.tool_response' - ] - assert ( - len(tool_response_calls) == 1 - ), "Expected 'gcp.vertex.agent.tool_response' to be set exactly once" - assert tool_response_calls[0].args[1] == '{}', ( - "Expected JSON string '{}' for tool_response when content capture is" - f' disabled, got {tool_response_calls[0].args[1]!r}' + _assert_span_attribute_set_to_empty_json( + mock_span_fixture, 'gcp.vertex.agent.tool_response' ) @@ -607,17 +585,8 @@ def test_trace_merged_tool_disabling_request_response_content( ) # Assert - Check attribute is set to JSON string '{}' not dict {} - tool_response_calls = [ - call - for call in mock_span_fixture.set_attribute.call_args_list - if call.args[0] == 'gcp.vertex.agent.tool_response' - ] - assert ( - len(tool_response_calls) == 1 - ), "Expected 'gcp.vertex.agent.tool_response' to be set exactly once" - assert tool_response_calls[0].args[1] == '{}', ( - "Expected JSON string '{}' for tool_response when content capture is" - f' disabled, got {tool_response_calls[0].args[1]!r}' + _assert_span_attribute_set_to_empty_json( + mock_span_fixture, 'gcp.vertex.agent.tool_response' ) From 3fe52a4750d17904f505d0c197e5d2e0b9e50ac2 Mon Sep 17 00:00:00 2001 From: DineshThumma9 Date: Fri, 9 Jan 2026 10:32:13 +0530 Subject: [PATCH 4/5] fixed formatting --- tests/unittests/telemetry/test_spans.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/tests/unittests/telemetry/test_spans.py b/tests/unittests/telemetry/test_spans.py index e1fa144bf5..7718c89024 100644 --- a/tests/unittests/telemetry/test_spans.py +++ b/tests/unittests/telemetry/test_spans.py @@ -100,12 +100,10 @@ def _assert_span_attribute_set_to_empty_json(mock_span, attribute_name: str): 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 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 is" - f' disabled, got {calls[0].args[1]!r}' + f"Expected JSON string '{{}}' for {attribute_name} when content capture" + f' is disabled, got {calls[0].args[1]!r}' ) From b73be84e8f73e3633ad783132911c090f1117839 Mon Sep 17 00:00:00 2001 From: Thumma Dinesh Date: Fri, 23 Jan 2026 09:59:57 +0530 Subject: [PATCH 5/5] style: Apply code formatting to tracing.py --- src/google/adk/telemetry/tracing.py | 798 ++++++++++++++-------------- 1 file changed, 386 insertions(+), 412 deletions(-) diff --git a/src/google/adk/telemetry/tracing.py b/src/google/adk/telemetry/tracing.py index 31bb3605b2..a5ae411470 100644 --- a/src/google/adk/telemetry/tracing.py +++ b/src/google/adk/telemetry/tracing.py @@ -37,46 +37,20 @@ from opentelemetry import _logs from opentelemetry import trace from opentelemetry._logs import LogRecord -from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import ( - GEN_AI_AGENT_DESCRIPTION, -) -from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import ( - GEN_AI_AGENT_NAME, -) -from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import ( - GEN_AI_CONVERSATION_ID, -) -from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import ( - GEN_AI_OPERATION_NAME, -) -from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import ( - GEN_AI_REQUEST_MODEL, -) -from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import ( - GEN_AI_RESPONSE_FINISH_REASONS, -) +from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import GEN_AI_AGENT_DESCRIPTION +from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import GEN_AI_AGENT_NAME +from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import GEN_AI_CONVERSATION_ID +from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import GEN_AI_OPERATION_NAME +from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import GEN_AI_REQUEST_MODEL +from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import GEN_AI_RESPONSE_FINISH_REASONS from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import GEN_AI_SYSTEM -from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import ( - GEN_AI_TOOL_CALL_ID, -) -from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import ( - GEN_AI_TOOL_DESCRIPTION, -) -from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import ( - GEN_AI_TOOL_NAME, -) -from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import ( - GEN_AI_TOOL_TYPE, -) -from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import ( - GEN_AI_USAGE_INPUT_TOKENS, -) -from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import ( - GEN_AI_USAGE_OUTPUT_TOKENS, -) -from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import ( - GenAiSystemValues, -) +from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import GEN_AI_TOOL_CALL_ID +from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import GEN_AI_TOOL_DESCRIPTION +from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import GEN_AI_TOOL_NAME +from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import GEN_AI_TOOL_TYPE +from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import GEN_AI_USAGE_INPUT_TOKENS +from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import GEN_AI_USAGE_OUTPUT_TOKENS +from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import GenAiSystemValues from opentelemetry.semconv.schemas import Schemas from opentelemetry.trace import Span from opentelemetry.util.types import AnyValue @@ -101,12 +75,12 @@ # Needed to avoid circular imports if TYPE_CHECKING: - from ..agents.base_agent import BaseAgent - from ..agents.invocation_context import InvocationContext - from ..events.event import Event - from ..models.llm_request import LlmRequest - from ..models.llm_response import LlmResponse - from ..tools.base_tool import BaseTool + from ..agents.base_agent import BaseAgent + from ..agents.invocation_context import InvocationContext + from ..events.event import Event + from ..models.llm_request import LlmRequest + from ..models.llm_response import LlmResponse + from ..tools.base_tool import BaseTool tracer = trace.get_tracer( instrumenting_module_name="gcp.vertex.agent", @@ -124,55 +98,55 @@ def _safe_json_serialize(obj) -> str: - """Convert any Python object to a JSON-serializable type or string. + """Convert any Python object to a JSON-serializable type or string. - Args: - obj: The object to serialize. + Args: + obj: The object to serialize. - Returns: - The JSON-serialized object string or if the object cannot be serialized. - """ + Returns: + The JSON-serialized object string or if the object cannot be serialized. + """ - try: - # Try direct JSON serialization first - return json.dumps( - obj, ensure_ascii=False, default=lambda o: "" - ) - except (TypeError, OverflowError): - return "" + try: + # Try direct JSON serialization first + return json.dumps( + obj, ensure_ascii=False, default=lambda o: "" + ) + except (TypeError, OverflowError): + return "" def trace_agent_invocation( span: trace.Span, agent: BaseAgent, ctx: InvocationContext ) -> None: - """Sets span attributes immediately available on agent invocation according to OTEL semconv version 1.37. + """Sets span attributes immediately available on agent invocation according to OTEL semconv version 1.37. - Args: - span: Span on which attributes are set. - agent: Agent from which attributes are gathered. - ctx: InvocationContext from which attributes are gathered. + Args: + span: Span on which attributes are set. + agent: Agent from which attributes are gathered. + ctx: InvocationContext from which attributes are gathered. - Inference related fields are not set, due to their planned removal from invoke_agent span: - https://github.com/open-telemetry/semantic-conventions/issues/2632 + Inference related fields are not set, due to their planned removal from invoke_agent span: + https://github.com/open-telemetry/semantic-conventions/issues/2632 - `gen_ai.agent.id` is not set because currently it's unclear what attributes this field should have, specifically: - - In which scope should it be unique (globally, given project, given agentic flow, given deployment). - - Should it be unchanging between deployments, and how this should this be achieved. + `gen_ai.agent.id` is not set because currently it's unclear what attributes this field should have, specifically: + - In which scope should it be unique (globally, given project, given agentic flow, given deployment). + - Should it be unchanging between deployments, and how this should this be achieved. - `gen_ai.data_source.id` is not set because it's not available. - Closest type which could contain this information is types.GroundingMetadata, which does not have an ID. + `gen_ai.data_source.id` is not set because it's not available. + Closest type which could contain this information is types.GroundingMetadata, which does not have an ID. - `server.*` attributes are not set pending confirmation from aabmass. - """ + `server.*` attributes are not set pending confirmation from aabmass. + """ - # Required - span.set_attribute(GEN_AI_OPERATION_NAME, "invoke_agent") + # Required + span.set_attribute(GEN_AI_OPERATION_NAME, "invoke_agent") - # Conditionally Required - span.set_attribute(GEN_AI_AGENT_DESCRIPTION, agent.description) + # Conditionally Required + span.set_attribute(GEN_AI_AGENT_DESCRIPTION, agent.description) - span.set_attribute(GEN_AI_AGENT_NAME, agent.name) - span.set_attribute(GEN_AI_CONVERSATION_ID, ctx.session.id) + span.set_attribute(GEN_AI_AGENT_NAME, agent.name) + span.set_attribute(GEN_AI_CONVERSATION_ID, ctx.session.id) def trace_tool_call( @@ -180,112 +154,112 @@ def trace_tool_call( args: dict[str, Any], function_response_event: Event | None, ): - """Traces tool call. + """Traces tool call. - Args: - tool: The tool that was called. - args: The arguments to the tool call. - function_response_event: The event with the function response details. - """ - span = trace.get_current_span() + Args: + tool: The tool that was called. + args: The arguments to the tool call. + function_response_event: The event with the function response details. + """ + 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) + span.set_attribute(GEN_AI_TOOL_DESCRIPTION, tool.description) + span.set_attribute(GEN_AI_TOOL_NAME, tool.name) - # e.g. FunctionTool - span.set_attribute(GEN_AI_TOOL_TYPE, tool.__class__.__name__) + # e.g. FunctionTool + span.set_attribute(GEN_AI_TOOL_TYPE, tool.__class__.__name__) - # Setting empty llm request and response (as UI expect these) while not - # applicable for tool_response. - span.set_attribute("gcp.vertex.agent.llm_request", EMPTY_JSON_STRING) - span.set_attribute("gcp.vertex.agent.llm_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", 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", - _safe_json_serialize(args), - ) - else: - span.set_attribute("gcp.vertex.agent.tool_call_args", EMPTY_JSON_STRING) - - # Tracing tool response - tool_call_id = "" - tool_response = "" - if ( - function_response_event is not None - and function_response_event.content is not None - and function_response_event.content.parts - ): - response_parts = function_response_event.content.parts - function_response = response_parts[0].function_response - if function_response is not None: - if function_response.id is not None: - tool_call_id = function_response.id - if function_response.response is not None: - tool_response = function_response.response - - span.set_attribute(GEN_AI_TOOL_CALL_ID, tool_call_id) - - if not isinstance(tool_response, dict): - tool_response = {"result": tool_response} - if function_response_event is not None: - 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", - _safe_json_serialize(tool_response), - ) - else: - span.set_attribute("gcp.vertex.agent.tool_response", EMPTY_JSON_STRING) + if _should_add_request_response_to_spans(): + span.set_attribute( + "gcp.vertex.agent.tool_call_args", + _safe_json_serialize(args), + ) + else: + span.set_attribute("gcp.vertex.agent.tool_call_args", EMPTY_JSON_STRING) + + # Tracing tool response + tool_call_id = "" + tool_response = "" + if ( + function_response_event is not None + and function_response_event.content is not None + and function_response_event.content.parts + ): + response_parts = function_response_event.content.parts + function_response = response_parts[0].function_response + if function_response is not None: + if function_response.id is not None: + tool_call_id = function_response.id + if function_response.response is not None: + tool_response = function_response.response + + span.set_attribute(GEN_AI_TOOL_CALL_ID, tool_call_id) + + if not isinstance(tool_response, dict): + tool_response = {"result": tool_response} + if function_response_event is not None: + 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", + _safe_json_serialize(tool_response), + ) + else: + span.set_attribute("gcp.vertex.agent.tool_response", EMPTY_JSON_STRING) def trace_merged_tool_calls( response_event_id: str, function_response_event: Event, ): - """Traces merged tool call events. + """Traces merged tool call events. - Calling this function is not needed for telemetry purposes. This is provided - for preventing /debug/trace requests (typically sent by web UI). + Calling this function is not needed for telemetry purposes. This is provided + for preventing /debug/trace requests (typically sent by web UI). - Args: - response_event_id: The ID of the response event. - function_response_event: The merged response event. - """ + Args: + response_event_id: The ID of the response event. + function_response_event: The merged response event. + """ - span = trace.get_current_span() + 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_TOOL_CALL_ID, response_event_id) + 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) - 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 = "" + # 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) + 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 = "" - if _should_add_request_response_to_spans(): - span.set_attribute( - "gcp.vertex.agent.tool_response", - function_response_event_json, - ) - else: - 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", EMPTY_JSON_STRING) + if _should_add_request_response_to_spans(): span.set_attribute( - "gcp.vertex.agent.llm_response", - EMPTY_JSON_STRING, + "gcp.vertex.agent.tool_response", + function_response_event_json, ) + else: + 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", EMPTY_JSON_STRING) + span.set_attribute( + "gcp.vertex.agent.llm_response", + EMPTY_JSON_STRING, + ) def trace_call_llm( @@ -295,80 +269,82 @@ def trace_call_llm( llm_response: LlmResponse, span: Span | None = None, ): - """Traces a call to the LLM. - - This function records details about the LLM request and response as - attributes on the current OpenTelemetry span. - - Args: - invocation_context: The invocation context for the current agent run. - event_id: The ID of the event. - llm_request: The LLM request object. - llm_response: The LLM response object. - """ - 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) + """Traces a call to the LLM. + + This function records details about the LLM request and response as + attributes on the current OpenTelemetry span. + + Args: + invocation_context: The invocation context for the current agent run. + event_id: The ID of the event. + llm_request: The LLM request object. + llm_response: The LLM response object. + """ + 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( + "gcp.vertex.agent.invocation_id", invocation_context.invocation_id + ) + span.set_attribute( + "gcp.vertex.agent.session_id", invocation_context.session.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.invocation_id", invocation_context.invocation_id + "gcp.vertex.agent.llm_request", + _safe_json_serialize(_build_llm_request_for_trace(llm_request)), ) - span.set_attribute("gcp.vertex.agent.session_id", invocation_context.session.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", - _safe_json_serialize(_build_llm_request_for_trace(llm_request)), - ) - else: - 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", - llm_request.config.top_p, - ) - if llm_request.config.max_output_tokens: - span.set_attribute( - "gen_ai.request.max_tokens", - llm_request.config.max_output_tokens, - ) + else: + 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", + llm_request.config.top_p, + ) + if llm_request.config.max_output_tokens: + span.set_attribute( + "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 = "" + + if _should_add_request_response_to_spans(): + span.set_attribute( + "gcp.vertex.agent.llm_response", + llm_response_json, + ) + else: + 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", + 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", + llm_response.usage_metadata.candidates_token_count, + ) + if llm_response.finish_reason: try: - llm_response_json = llm_response.model_dump_json(exclude_none=True) - except Exception: # pylint: disable=broad-exception-caught - llm_response_json = "" - - if _should_add_request_response_to_spans(): - span.set_attribute( - "gcp.vertex.agent.llm_response", - llm_response_json, - ) - else: - 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", - 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", - llm_response.usage_metadata.candidates_token_count, - ) - if llm_response.finish_reason: - try: - finish_reason_str = llm_response.finish_reason.value.lower() - except AttributeError: - finish_reason_str = str(llm_response.finish_reason).lower() - span.set_attribute( - "gen_ai.response.finish_reasons", - [finish_reason_str], - ) + finish_reason_str = llm_response.finish_reason.value.lower() + except AttributeError: + finish_reason_str = str(llm_response.finish_reason).lower() + span.set_attribute( + "gen_ai.response.finish_reasons", + [finish_reason_str], + ) def trace_send_data( @@ -376,69 +352,67 @@ def trace_send_data( event_id: str, data: list[types.Content], ): - """Traces the sending of data to the agent. - - This function records details about the data sent to the agent as - attributes on the current OpenTelemetry span. - - Args: - invocation_context: The invocation context for the current agent run. - event_id: The ID of the event. - data: A list of content objects. - """ - span = trace.get_current_span() + """Traces the sending of data to the agent. + + This function records details about the data sent to the agent as + attributes on the current OpenTelemetry span. + + Args: + invocation_context: The invocation context for the current agent run. + event_id: The ID of the event. + data: A list of content objects. + """ + span = trace.get_current_span() + span.set_attribute( + "gcp.vertex.agent.invocation_id", invocation_context.invocation_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.invocation_id", invocation_context.invocation_id + "gcp.vertex.agent.data", + _safe_json_serialize([ + types.Content(role=content.role, parts=content.parts).model_dump( + exclude_none=True, mode="json" + ) + for content in data + ]), ) - 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", - _safe_json_serialize( - [ - types.Content(role=content.role, parts=content.parts).model_dump( - exclude_none=True, mode="json" - ) - for content in data - ] - ), - ) - else: - span.set_attribute("gcp.vertex.agent.data", EMPTY_JSON_STRING) + else: + span.set_attribute("gcp.vertex.agent.data", EMPTY_JSON_STRING) def _build_llm_request_for_trace(llm_request: LlmRequest) -> dict[str, Any]: - """Builds a dictionary representation of the LLM request for tracing. - - This function prepares a dictionary representation of the LlmRequest - object, suitable for inclusion in a trace. It excludes fields that cannot - be serialized (e.g., function pointers) and avoids sending bytes data. - - Args: - llm_request: The LlmRequest object. - - Returns: - A dictionary representation of the LLM request. - """ - # 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" - ), - "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( - types.Content(role=content.role, parts=parts).model_dump( - exclude_none=True, mode="json" - ) + """Builds a dictionary representation of the LLM request for tracing. + + This function prepares a dictionary representation of the LlmRequest + object, suitable for inclusion in a trace. It excludes fields that cannot + be serialized (e.g., function pointers) and avoids sending bytes data. + + Args: + llm_request: The LlmRequest object. + + Returns: + A dictionary representation of the LLM request. + """ + # 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" + ), + "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( + types.Content(role=content.role, parts=parts).model_dump( + exclude_none=True, mode="json" ) - return result + ) + return result # Defaults to true for now to preserve backward compatibility. @@ -446,10 +420,10 @@ def _build_llm_request_for_trace(llm_request: LlmRequest) -> dict[str, Any]: # a deprecation of request/response content in spans by switching the default # 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") - return not disabled_via_env_var + disabled_via_env_var = os.getenv( + ADK_CAPTURE_MESSAGE_CONTENT_IN_SPANS, "true" + ).lower() in ("false", "0") + return not disabled_via_env_var @contextmanager @@ -458,81 +432,81 @@ def use_generate_content_span( invocation_context: InvocationContext, model_response_event: Event, ) -> Iterator[Span | None]: - """Context manager encompassing `generate_content {model.name}` span. - - When an external library for inference instrumentation is installed (e.g. opentelemetry-instrumentation-google-genai), - span creation is delegated to said library. - """ - - common_attributes = { - GEN_AI_CONVERSATION_ID: invocation_context.session.id, - "gcp.vertex.agent.event_id": model_response_event.id, - } - if ( - _is_gemini_agent(invocation_context.agent) - and _instrumented_with_opentelemetry_instrumentation_google_genai() - ): - yield None - else: - with _use_native_generate_content_span( - llm_request=llm_request, - common_attributes=common_attributes, - ) as span: - yield span + """Context manager encompassing `generate_content {model.name}` span. + + When an external library for inference instrumentation is installed (e.g. opentelemetry-instrumentation-google-genai), + span creation is delegated to said library. + """ + + common_attributes = { + GEN_AI_CONVERSATION_ID: invocation_context.session.id, + "gcp.vertex.agent.event_id": model_response_event.id, + } + if ( + _is_gemini_agent(invocation_context.agent) + and _instrumented_with_opentelemetry_instrumentation_google_genai() + ): + yield None + else: + with _use_native_generate_content_span( + llm_request=llm_request, + common_attributes=common_attributes, + ) as span: + yield span def _should_log_prompt_response_content() -> bool: - return os.getenv( - OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT, "" - ).lower() in ("1", "true") + return os.getenv( + OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT, "" + ).lower() in ("1", "true") def _serialize_content(content: types.ContentUnion) -> AnyValue: - if isinstance(content, BaseModel): - return content.model_dump() - if isinstance(content, str): - return content - if isinstance(content, list): - return [_serialize_content(part) for part in content] - return _safe_json_serialize(content) + if isinstance(content, BaseModel): + return content.model_dump() + if isinstance(content, str): + return content + if isinstance(content, list): + return [_serialize_content(part) for part in content] + return _safe_json_serialize(content) def _serialize_content_with_elision( content: types.ContentUnion | None, ) -> AnyValue: - if not _should_log_prompt_response_content(): - return USER_CONTENT_ELIDED - if content is None: - return None - return _serialize_content(content) + if not _should_log_prompt_response_content(): + return USER_CONTENT_ELIDED + if content is None: + return None + return _serialize_content(content) 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): - if ( - "opentelemetry/instrumentation/google_genai" - in maybe_wrapped_function.__code__.co_filename - ): - return True - maybe_wrapped_function = wrapped # pyright: ignore[reportAny] + maybe_wrapped_function = Models.generate_content + print(f"{Models.generate_content.__code__.co_filename=}") + while wrapped := getattr(maybe_wrapped_function, "__wrapped__", None): + if ( + "opentelemetry/instrumentation/google_genai" + in maybe_wrapped_function.__code__.co_filename + ): + return True + maybe_wrapped_function = wrapped # pyright: ignore[reportAny] - return False + return False def _is_gemini_agent(agent: BaseAgent) -> bool: - from ..agents.llm_agent import LlmAgent + from ..agents.llm_agent import LlmAgent - if not isinstance(agent, LlmAgent): - return False + if not isinstance(agent, LlmAgent): + return False - if isinstance(agent.model, str): - return is_gemini_model(agent.model) + if isinstance(agent.model, str): + return is_gemini_model(agent.model) - from ..models.google_llm import Gemini + from ..models.google_llm import Gemini - return isinstance(agent.model, Gemini) + return isinstance(agent.model, Gemini) @contextmanager @@ -540,77 +514,77 @@ def _use_native_generate_content_span( llm_request: LlmRequest, common_attributes: Mapping[str, AttributeValue], ) -> Iterator[Span]: - with tracer.start_as_current_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_attributes(common_attributes) - - otel_logger.emit( - LogRecord( - event_name="gen_ai.system.message", - body={ - "content": _serialize_content_with_elision( - llm_request.config.system_instruction - ) - }, - attributes={GEN_AI_SYSTEM: _guess_gemini_system_name()}, - ) - ) - - for content in llm_request.contents: - otel_logger.emit( - LogRecord( - event_name="gen_ai.user.message", - body={"content": _serialize_content_with_elision(content)}, - attributes={GEN_AI_SYSTEM: _guess_gemini_system_name()}, - ) - ) - - yield span - - -def trace_generate_content_result(span: Span | None, llm_response: LlmResponse): - """Trace result of the inference in generate_content span.""" - - if span is None: - return - - if llm_response.partial: - return - - if finish_reason := llm_response.finish_reason: - span.set_attribute(GEN_AI_RESPONSE_FINISH_REASONS, [finish_reason.lower()]) - if usage_metadata := llm_response.usage_metadata: - if usage_metadata.prompt_token_count is not None: - span.set_attribute( - GEN_AI_USAGE_INPUT_TOKENS, usage_metadata.prompt_token_count - ) - if usage_metadata.candidates_token_count is not None: - span.set_attribute( - GEN_AI_USAGE_OUTPUT_TOKENS, usage_metadata.candidates_token_count - ) + with tracer.start_as_current_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_attributes(common_attributes) otel_logger.emit( LogRecord( - event_name="gen_ai.choice", + event_name="gen_ai.system.message", body={ - "content": _serialize_content_with_elision(llm_response.content), - "index": 0, # ADK always returns a single candidate - } - | {"finish_reason": llm_response.finish_reason.value} - if llm_response.finish_reason is not None - else {}, + "content": _serialize_content_with_elision( + llm_request.config.system_instruction + ) + }, attributes={GEN_AI_SYSTEM: _guess_gemini_system_name()}, ) ) + for content in llm_request.contents: + otel_logger.emit( + LogRecord( + event_name="gen_ai.user.message", + body={"content": _serialize_content_with_elision(content)}, + attributes={GEN_AI_SYSTEM: _guess_gemini_system_name()}, + ) + ) + + yield span + + +def trace_generate_content_result(span: Span | None, llm_response: LlmResponse): + """Trace result of the inference in generate_content span.""" + + if span is None: + return + + if llm_response.partial: + return + + if finish_reason := llm_response.finish_reason: + span.set_attribute(GEN_AI_RESPONSE_FINISH_REASONS, [finish_reason.lower()]) + if usage_metadata := llm_response.usage_metadata: + if usage_metadata.prompt_token_count is not None: + span.set_attribute( + GEN_AI_USAGE_INPUT_TOKENS, usage_metadata.prompt_token_count + ) + if usage_metadata.candidates_token_count is not None: + span.set_attribute( + GEN_AI_USAGE_OUTPUT_TOKENS, usage_metadata.candidates_token_count + ) + + otel_logger.emit( + LogRecord( + event_name="gen_ai.choice", + body={ + "content": _serialize_content_with_elision(llm_response.content), + "index": 0, # ADK always returns a single candidate + } + | {"finish_reason": llm_response.finish_reason.value} + if llm_response.finish_reason is not None + else {}, + attributes={GEN_AI_SYSTEM: _guess_gemini_system_name()}, + ) + ) + def _guess_gemini_system_name() -> str: - return ( - GenAiSystemValues.VERTEX_AI.name.lower() - if os.getenv("GOOGLE_GENAI_USE_VERTEXAI", "").lower() in ("true", "1") - else GenAiSystemValues.GEMINI.name.lower() - ) + return ( + GenAiSystemValues.VERTEX_AI.name.lower() + if os.getenv("GOOGLE_GENAI_USE_VERTEXAI", "").lower() in ("true", "1") + else GenAiSystemValues.GEMINI.name.lower() + )