feat(trace): support invoation id in spans (#410)

Author: yaozheng-fang

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

@@ -115,6 +115,23 @@ def common_gen_ai_session_id(**kwargs) -> str:
     return session_id or "<unknown_session_id>"
 
 
+def common_gen_ai_invocation_id(**kwargs) -> str:
+    """Extract the invocation identifier from context.
+
+    Provides invocation-level identification for organizing telemetry data
+    by individual API calls and enabling detailed analysis of each operation.
+
+    Args:
+        **kwargs: Keyword arguments containing context information.
+            Expected to include 'invocation_id' key.
+
+    Returns:
+        str: Invocation identifier or placeholder if not available
+    """
+    invocation_id = kwargs.get("invocation_id")
+    return invocation_id or "<unknown_invocation_id>"
+
+
 def common_cozeloop_report_source(**kwargs) -> str:
     """Extract the CozeLoop report source identifier.
 
@@ -175,6 +192,7 @@ def llm_openinference_instrumentation_veadk(**kwargs) -> str:
     "app.name": common_gen_ai_app_name,  # TLS required
     "user.id": common_gen_ai_user_id,  # CozeLoop / TLS required
     "session.id": common_gen_ai_session_id,  # CozeLoop / TLS required
+    "invocation.id": common_gen_ai_invocation_id,  # CozeLoop required
     "cozeloop.report.source": common_cozeloop_report_source,  # CozeLoop required
     "cozeloop.call_type": common_cozeloop_call_type,  # CozeLoop required
 }

veadk/tracing/telemetry/telemetry.py

@@ -147,7 +147,10 @@ def _set_agent_input_attribute(
             if part.text:
                 span.add_event(
                     "gen_ai.user.message",
-                    {f"parts.{idx}.type": "text", f"parts.{idx}.content": part.text},
+                    {
+                        f"parts.{idx}.type": "text",
+                        f"parts.{idx}.content": part.text,
+                    },
                 )
             if part.inline_data:
                 span.add_event(
@@ -186,7 +189,8 @@ def _set_agent_output_attribute(span: Span, llm_response: LlmResponse) -> None:
     if content and content.parts:
         # set gen_ai.output attribute required by APMPlus
         span.set_attribute(
-            "gen_ai.output", safe_json_serialize(content.model_dump(exclude_none=True))
+            "gen_ai.output",
+            safe_json_serialize(content.model_dump(exclude_none=True)),
         )
 
         for idx, part in enumerate(content.parts):
@@ -374,6 +378,7 @@ def trace_call_llm(
         user_id=invocation_context.user_id,
         app_name=invocation_context.app_name,
         session_id=invocation_context.session.id,
+        invocation_id=invocation_context.invocation_id,
         model_provider=invocation_context.agent.model_provider
         if isinstance(invocation_context.agent, Agent)
         else "",