realize invocation and agent run attributes

Author: yaozheng-fang

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

@@ -14,6 +14,11 @@ def common_gen_ai_app_name(**kwargs) -> str:
     return app_name or "<unknown_app_name>"
 
 
+def common_gen_ai_agent_name(**kwargs) -> str:
+    agent_name = kwargs.get("agent_name")
+    return agent_name or "<unknown_agent_name>"
+
+
 def common_gen_ai_user_id(**kwargs) -> str:
     user_id = kwargs.get("user_id")
     return user_id or "<unknown_user_id>"
@@ -28,6 +33,7 @@ def common_gen_ai_session_id(**kwargs) -> str:
     "gen_ai.system": common_gen_ai_system,
     "gen_ai.system.version": common_gen_ai_system_version,
     "gen_ai.app.name": common_gen_ai_app_name,
+    "gen_ai.agent.name": common_gen_ai_agent_name,
     "gen_ai.user.id": common_gen_ai_user_id,
     "gen_ai.session.id": common_gen_ai_session_id,
 }

veadk/tracing/telemetry/exporters/inmemory_exporter.py

@@ -14,6 +14,12 @@
 
 from typing import Sequence
 
+from opentelemetry.context import (
+    _SUPPRESS_INSTRUMENTATION_KEY,
+    attach,
+    detach,
+    set_value,
+)
 from opentelemetry.sdk.trace import ReadableSpan, export
 from typing_extensions import override
 
@@ -23,7 +29,7 @@
 logger = get_logger(__name__)
 
 
-in_memory_exporter_instance = None
+inmemory_span_processor = None
 
 
 # ======== Adapted from Google ADK ========
@@ -64,6 +70,30 @@ def clear(self):
         self._spans.clear()
 
 
+class _InMemorySpanProcessor(export.SimpleSpanProcessor):
+    def __init__(self, exporter: _InMemoryExporter) -> None:
+        super().__init__(exporter)
+        self.spans = []
+
+    def on_start(self, span, parent_context) -> None:
+        if span.context:
+            self.spans.append(span)
+
+    def on_end(self, span: ReadableSpan) -> None:
+        if span.context:
+            if not span.context.trace_flags.sampled:
+                return
+            token = attach(set_value(_SUPPRESS_INSTRUMENTATION_KEY, True))
+            try:
+                self.span_exporter.export((span,))
+            # pylint: disable=broad-exception-caught
+            except Exception:
+                logger.exception("Exception while exporting Span.")
+            detach(token)
+            if span in self.spans:
+                self.spans.remove(span)
+
+
 class InMemoryExporter(BaseExporter):
     """InMemory Exporter mainly for store spans in memory for debugging / observability purposes."""
 
@@ -73,4 +103,7 @@ def __init__(self, name: str = "inmemory_exporter") -> None:
         self.name = name
 
         self._exporter = _InMemoryExporter()
-        self.processor = export.SimpleSpanProcessor(self._exporter)
+        self.processor = _InMemorySpanProcessor(self._exporter)
+
+        global inmemory_span_processor
+        inmemory_span_processor = self.processor

veadk/tracing/telemetry/opentelemetry_tracer.py

@@ -25,7 +25,7 @@
 from opentelemetry.sdk.resources import Resource
 from opentelemetry.sdk.trace import TracerProvider
 from opentelemetry.sdk.trace.export import BatchSpanProcessor, SimpleSpanProcessor
-from pydantic import BaseModel, ConfigDict, Field, field_validator
+from pydantic import BaseModel, ConfigDict, Field, PrivateAttr, field_validator
 from typing_extensions import override
 
 from veadk.tracing.base_tracer import BaseTracer
@@ -37,8 +37,6 @@
 
 logger = get_logger(__name__)
 
-DEFAULT_VEADK_TRACER_NAME = "veadk_global_tracer"
-
 
 def update_resource_attributions(provider: TracerProvider, resource_attributes: dict):
     provider._resource = provider._resource.merge(Resource.create(resource_attributes))
@@ -47,20 +45,37 @@ def update_resource_attributions(provider: TracerProvider, resource_attributes:
 class OpentelemetryTracer(BaseModel, BaseTracer):
     model_config = ConfigDict(arbitrary_types_allowed=True)
 
-    name: str = Field(
-        default=DEFAULT_VEADK_TRACER_NAME, description="The identifier of tracer."
-    )
-
-    app_name: str = Field(
-        default="veadk_app",
-        description="The identifier of app.",
-    )
+    name: str = Field(default="veadk_tracer", description="The identifier of tracer.")
 
     exporters: list[BaseExporter] = Field(
         default_factory=list,
         description="The exporters to export spans.",
     )
 
+    _app_name: str = PrivateAttr(default="<unknown_app_name>")
+
+    _agent_name: str = PrivateAttr(default="<unknown_agent_name>")
+
+    @property
+    def app_name(self) -> str:
+        return self._app_name
+
+    @app_name.setter
+    def app_name(self, value: str) -> None:
+        self._app_name = value
+        # update_common_attributes(self._tracer_provider, {"app_name": self._app_name})
+
+    @property
+    def agent_name(self) -> str:
+        return self._agent_name
+
+    @agent_name.setter
+    def agent_name(self, value: str) -> None:
+        self._agent_name = value
+        # update_common_attributes(
+        #     self._tracer_provider, {"agent_name": self._agent_name}
+        # )
+
     @field_validator("exporters")
     @classmethod
     def forbid_inmemory_exporter(cls, v: list[BaseExporter]) -> list[BaseExporter]:
@@ -73,8 +88,6 @@ def model_post_init(self, context: Any) -> None:
         patch_google_adk_telemetry()
 
         self._processors = []
-        self._inmemory_exporter = InMemoryExporter()
-        self.exporters.append(self._inmemory_exporter)
 
         # VeADK operates on global OpenTelemetry provider, return nothing
         self._init_global_tracer_provider()
@@ -121,9 +134,22 @@ def _init_global_tracer_provider(self) -> None:
                     f"Add span processor for exporter `{exporter.__class__.__name__}` to OpentelemetryTracer failed."
                 )
 
-        logger.debug(
-            f"Init OpentelemetryTracer with {len(self._processors)} exporters."
-        )
+        self._inmemory_exporter = InMemoryExporter()
+        self._inmemory_exporter_processor = self._inmemory_exporter.processor
+
+        # make sure the in memory exporter processor is added at index 0
+        if self._inmemory_exporter_processor:
+            global_tracer_provider._active_span_processor._span_processors = (
+                self._inmemory_exporter_processor,
+            ) + global_tracer_provider._active_span_processor._span_processors
+
+            self._processors.append(self._inmemory_exporter_processor)
+        else:
+            logger.warning(
+                "InMemoryExporter processor is not initialized, cannot add to OpentelemetryTracer."
+            )
+
+        logger.info(f"Init OpentelemetryTracer with {len(self._processors)} exporters.")
 
     @property
     def trace_file_path(self) -> str:

veadk/tracing/telemetry/telemetry.py

@@ -6,7 +6,7 @@
 from google.adk.models.llm_response import LlmResponse
 from google.adk.tools import BaseTool
 from opentelemetry import trace
-from opentelemetry.trace.span import Span
+from opentelemetry.sdk.trace import _Span
 
 from veadk.tracing.telemetry.attributes.attributes import ATTRIBUTES
 from veadk.tracing.telemetry.attributes.extractors.llm_attributes_extractors import (
@@ -15,23 +15,32 @@
 from veadk.tracing.telemetry.attributes.extractors.tool_attributes_extractors import (
     ToolAttributesParams,
 )
+from veadk.tracing.telemetry.exporters import inmemory_exporter
 from veadk.utils.logger import get_logger
 
 logger = get_logger(__name__)
 
 
-def _set_common_attributes(
-    span: Span, app_name: str, user_id: str, session_id: str
-) -> None:
-    common_attributes = ATTRIBUTES.get("common", {})
-    for attr_name, attr_extractor in common_attributes.items():
-        value = attr_extractor(
-            app_name=app_name, user_id=user_id, session_id=session_id
-        )
-        span.set_attribute(attr_name, value)
+def trace_send_data(): ...
 
 
-def trace_send_data(): ...
+def set_common_attributes(current_span: _Span, **kwargs) -> None:
+    if current_span.context:
+        current_span_id = current_span.context.trace_id
+
+    spans = inmemory_exporter.inmemory_span_processor.spans
+
+    spans_in_current_trace = [
+        span
+        for span in spans
+        if span.context and span.context.trace_id == current_span_id
+    ]
+
+    common_attributes = ATTRIBUTES.get("common", {})
+    for span in spans_in_current_trace:
+        for attr_name, attr_extractor in common_attributes.items():
+            value = attr_extractor(**kwargs)
+            span.set_attribute(attr_name, value)
 
 
 def trace_tool_call(
@@ -63,10 +72,11 @@ def trace_call_llm(
 ) -> None:
     span = trace.get_current_span()
 
-    _set_common_attributes(
-        span=span,
-        app_name=invocation_context.session.app_name,
-        user_id=invocation_context.session.user_id,
+    set_common_attributes(
+        current_span=span,  # type: ignore
+        agent_name=invocation_context.agent.name,
+        app_name=invocation_context.app_name,
+        user_id=invocation_context.user_id,
         session_id=invocation_context.session.id,
     )