feat: agent run span

Author: feng-95

veadk/tracing/telemetry/attributes/extractors/llm_attributes_extractors.py

@@ -561,7 +561,7 @@ def llm_gen_ai_request_functions(params: LLMAttributesParams) -> ExtractorRespon
     "gen_ai.is_streaming": llm_gen_ai_is_streaming,
     # -> 1.4. span kind
     "gen_ai.operation.name": llm_gen_ai_operation_name,
-    "gen_ai.span.kind": llm_gen_ai_span_kind, # apmplus required
+    "gen_ai.span.kind": llm_gen_ai_span_kind,  # apmplus required
     # -> 1.5. inputs
     "gen_ai.prompt": llm_gen_ai_prompt,
     # -> 1.6. outputs

veadk/tracing/telemetry/attributes/extractors/tool_attributes_extractors.py

@@ -76,7 +76,7 @@ def tool_gen_ai_tool_output(params: ToolAttributesParams) -> ExtractorResponse:
     "gen_ai.tool.output": tool_gen_ai_tool_output,  # TLS required
     "cozeloop.input": tool_gen_ai_tool_input,  # CozeLoop required
     "cozeloop.output": tool_gen_ai_tool_output,  # CozeLoop required
-    "gen_ai.span.kind": tool_gen_ai_span_kind, # apmplus required
-    "gen_ai.input": tool_gen_ai_tool_input, # apmplus required
-    "gen_ai.output": tool_gen_ai_tool_output, # apmplus required
+    "gen_ai.span.kind": tool_gen_ai_span_kind,  # apmplus required
+    "gen_ai.input": tool_gen_ai_tool_input,  # apmplus required
+    "gen_ai.output": tool_gen_ai_tool_output,  # apmplus required
 }

veadk/tracing/telemetry/exporters/inmemory_exporter.py

@@ -77,18 +77,20 @@ def __init__(self, exporter: _InMemoryExporter) -> None:
         super().__init__(exporter)
 
     def on_start(self, span, parent_context) -> None:
-        if span.name.startswith("invocation") or span.name.startswith("invoke") :
+        if span.name.startswith("invocation"):
             span.set_attribute("gen_ai.operation.name", "chain")
             span.set_attribute("gen_ai.span.kind", "workflow")
             span.set_attribute("gen_ai.usage.total_tokens", 0)
             ctx = set_value("invocation_span_instance", span, context=parent_context)
-            # suppress instrumentation for llm from apmplus, such as openai
-            ctx = set_value("suppress_language_model_instrumentation", True, context=ctx)
+            # suppress instrumentation for llm to avoid auto instrument from apmplus, such as openai
+            ctx = set_value(
+                "suppress_language_model_instrumentation", True, context=ctx
+            )
 
             token = attach(ctx)  # mount context on `invocation` root span in Google ADK
             setattr(span, "_invocation_token", token)  # for later detach
 
-        if span.name.startswith("agent_run"):
+        if span.name.startswith("agent_run") or span.name.startswith("invoke_agent"):
             span.set_attribute("gen_ai.operation.name", "agent")
             span.set_attribute("gen_ai.span.kind", "agent")
 

veadk/tracing/telemetry/telemetry.py

@@ -20,10 +20,9 @@
 from google.adk.models.llm_response import LlmResponse
 from google.adk.tools import BaseTool
 from opentelemetry import trace
-from opentelemetry.context import get_value, set_value
+from opentelemetry.context import get_value
 from opentelemetry.sdk.trace import Span, _Span
 
-from veadk.utils.misc import safe_json_serialize
 from veadk.tracing.telemetry.attributes.attributes import ATTRIBUTES
 from veadk.tracing.telemetry.attributes.extractors.types import (
     ExtractorResponse,
@@ -93,8 +92,7 @@ def _set_agent_input_attribute(
     user_content = invocation_context.user_content
     if user_content and user_content.parts:
         # set gen_ai.input attribute required by APMPlus
-        # TODO: 优化 gen_ai.input，目前无法序列化
-        span.set_attribute("gen_ai.input", safe_json_serialize(user_content.parts))
+        span.set_attribute("gen_ai.input", user_content.model_dump(exclude_none=True))
 
         span.add_event(
             "gen_ai.user.message",
@@ -132,15 +130,9 @@ def _set_agent_input_attribute(
 
 def _set_agent_output_attribute(span: Span, llm_response: LlmResponse) -> None:
     content = llm_response.content
-    print(f"### response:  {llm_response}")
-    print(f"### output_transcription:  {llm_response.output_transcription}")
-    print(f"### input_transcription:  {llm_response.input_transcription}")
-
-
     if content and content.parts:
         # set gen_ai.output attribute required by APMPlus
-        span.set_attribute("gen_ai.output", safe_json_serialize(content.parts))
-        print("### parts: ", content.parts)
+        span.set_attribute("gen_ai.output", content.model_dump(exclude_none=True))
         for idx, part in enumerate(content.parts):
             if part.text:
                 span.add_event(
@@ -163,7 +155,7 @@ def set_common_attributes_on_model_span(
         invocation_span: Span = get_value("invocation_span_instance")  # type: ignore
         agent_run_span: Span = get_value("agent_run_span_instance")  # type: ignore
 
-        if invocation_span and (invocation_span.name.startswith("invocation") or invocation_span.name.startswith("invoke")):
+        if invocation_span and invocation_span.name.startswith("invocation"):
             _set_agent_input_attribute(invocation_span, invocation_context)
             _set_agent_output_attribute(invocation_span, llm_response)
             for attr_name, attr_extractor in common_attributes.items():
@@ -187,10 +179,14 @@ def set_common_attributes_on_model_span(
             )  # we can ignore this warning, cause we manually set the attribute to int before
             invocation_span.set_attribute(
                 # record input/output token usage?
-                "gen_ai.usage.total_tokens", accumulated_total_token_usage
+                "gen_ai.usage.total_tokens",
+                accumulated_total_token_usage,
             )
 
-        if agent_run_span and agent_run_span.name.startswith("agent_run"):
+        if agent_run_span and (
+            agent_run_span.name.startswith("agent_run")
+            or agent_run_span.name.startswith("invoke_agent")
+        ):
             _set_agent_input_attribute(agent_run_span, invocation_context)
             _set_agent_output_attribute(agent_run_span, llm_response)
             for attr_name, attr_extractor in common_attributes.items():