feat: add gen_ai input

Author: feng-95

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

@@ -324,6 +324,8 @@ def llm_gen_ai_is_streaming(params: LLMAttributesParams) -> ExtractorResponse:
 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 = {
@@ -559,6 +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
     # -> 1.5. inputs
     "gen_ai.prompt": llm_gen_ai_prompt,
     # -> 1.6. outputs

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

@@ -22,6 +22,9 @@
 def tool_gen_ai_operation_name(params: ToolAttributesParams) -> ExtractorResponse:
     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 = {
@@ -73,4 +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
 }

veadk/tracing/telemetry/exporters/inmemory_exporter.py

@@ -77,16 +77,20 @@ def __init__(self, exporter: _InMemoryExporter) -> None:
         super().__init__(exporter)
 
     def on_start(self, span, parent_context) -> None:
-        if span.name.startswith("invocation"):
+        if span.name.startswith("invocation") or span.name.startswith("invoke") :
             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)
+
             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"):
             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

@@ -20,9 +20,10 @@
 from google.adk.models.llm_response import LlmResponse
 from google.adk.tools import BaseTool
 from opentelemetry import trace
-from opentelemetry.context import get_value
+from opentelemetry.context import get_value, set_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,
@@ -91,6 +92,10 @@ 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.add_event(
             "gen_ai.user.message",
             {
@@ -127,7 +132,15 @@ 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)
         for idx, part in enumerate(content.parts):
             if part.text:
                 span.add_event(
@@ -150,7 +163,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"):
+        if invocation_span and (invocation_span.name.startswith("invocation") or invocation_span.name.startswith("invoke")):
             _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():
@@ -173,6 +186,7 @@ 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(
+                # record input/output token usage?
                 "gen_ai.usage.total_tokens", accumulated_total_token_usage
             )