feat(tracing): support APMPLus session (#218)

* feat: add gen_ai input

* feat: agent run span

* feat: format

Author: feng-95

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

@@ -325,6 +325,10 @@ def llm_gen_ai_operation_name(params: LLMAttributesParams) -> ExtractorResponse:
     return ExtractorResponse(content="chat")
 
 
+def llm_gen_ai_span_kind(params: LLMAttributesParams) -> ExtractorResponse:
+    return ExtractorResponse(content="llm")
+
+
 # def llm_gen_ai_system_message(params: LLMAttributesParams) -> ExtractorResponse:
 #     event_attributes = {
 #         "content": str(params.llm_request.config.system_instruction),
@@ -559,6 +563,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
     # -> 1.5. inputs
     "gen_ai.prompt": llm_gen_ai_prompt,
     # -> 1.6. outputs

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

@@ -23,6 +23,10 @@ def tool_gen_ai_operation_name(params: ToolAttributesParams) -> ExtractorRespons
     return ExtractorResponse(content="execute_tool")
 
 
+def tool_gen_ai_span_kind(params: ToolAttributesParams) -> ExtractorResponse:
+    return ExtractorResponse(content="tool")
+
+
 def tool_gen_ai_tool_message(params: ToolAttributesParams) -> ExtractorResponse:
     tool_input = {
         "role": "tool",
@@ -73,4 +77,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
 }

veadk/tracing/telemetry/exporters/inmemory_exporter.py

@@ -79,14 +79,20 @@ def __init__(self, exporter: _InMemoryExporter) -> None:
     def on_start(self, span, parent_context) -> None:
         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 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")
 
             ctx = set_value("agent_run_span_instance", span, context=parent_context)
             token = attach(ctx)

veadk/tracing/telemetry/telemetry.py

@@ -91,6 +91,9 @@ 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
+        span.set_attribute("gen_ai.input", user_content.model_dump(exclude_none=True))
+
         span.add_event(
             "gen_ai.user.message",
             {
@@ -128,6 +131,9 @@ def _set_agent_input_attribute(
 def _set_agent_output_attribute(span: Span, llm_response: LlmResponse) -> None:
     content = llm_response.content
     if content and content.parts:
+        # set gen_ai.output attribute required by APMPlus
+        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(
@@ -173,10 +179,15 @@ def set_common_attributes_on_model_span(
                 current_step_token_usage + int(prev_total_token_usage)  # type: ignore
             )  # we can ignore this warning, cause we manually set the attribute to int before
             invocation_span.set_attribute(
-                "gen_ai.usage.total_tokens", accumulated_total_token_usage
+                # record input/output 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():