diff --git a/veadk/tracing/telemetry/exporters/inmemory_exporter.py b/veadk/tracing/telemetry/exporters/inmemory_exporter.py index 37b77acf..cc236e6b 100644 --- a/veadk/tracing/telemetry/exporters/inmemory_exporter.py +++ b/veadk/tracing/telemetry/exporters/inmemory_exporter.py @@ -106,3 +106,6 @@ def __init__(self, name: str = "inmemory_exporter") -> None: self._exporter = _InMemoryExporter() self.processor = _InMemorySpanProcessor(self._exporter) + + +_INMEMORY_EXPORTER_INSTANCE = InMemoryExporter() diff --git a/veadk/tracing/telemetry/opentelemetry_tracer.py b/veadk/tracing/telemetry/opentelemetry_tracer.py index 9f1019b9..34e6e436 100644 --- a/veadk/tracing/telemetry/opentelemetry_tracer.py +++ b/veadk/tracing/telemetry/opentelemetry_tracer.py @@ -31,6 +31,9 @@ from veadk.tracing.telemetry.exporters.apmplus_exporter import APMPlusExporter from veadk.tracing.telemetry.exporters.base_exporter import BaseExporter from veadk.tracing.telemetry.exporters.inmemory_exporter import InMemoryExporter +from veadk.tracing.telemetry.exporters.inmemory_exporter import ( + _INMEMORY_EXPORTER_INSTANCE, +) from veadk.utils.logger import get_logger from veadk.utils.patches import patch_google_adk_telemetry @@ -110,7 +113,7 @@ def _init_global_tracer_provider(self) -> None: f"Add span processor for exporter `{exporter.__class__.__name__}` to OpentelemetryTracer failed." ) - self._inmemory_exporter = InMemoryExporter() + self._inmemory_exporter = _INMEMORY_EXPORTER_INSTANCE if self._inmemory_exporter.processor: # make sure the in memory exporter processor is added at index 0 # because we use this to record all spans diff --git a/veadk/tracing/telemetry/telemetry.py b/veadk/tracing/telemetry/telemetry.py index fe0218d5..8f4d1c1f 100644 --- a/veadk/tracing/telemetry/telemetry.py +++ b/veadk/tracing/telemetry/telemetry.py @@ -28,6 +28,9 @@ LLMAttributesParams, ToolAttributesParams, ) +from veadk.tracing.telemetry.exporters.inmemory_exporter import ( + _INMEMORY_EXPORTER_INSTANCE, +) from veadk.utils.logger import get_logger logger = get_logger(__name__) @@ -54,8 +57,6 @@ def trace_send_data(): ... def set_common_attributes( invocation_context: InvocationContext, current_span: _Span, **kwargs ) -> None: - from veadk.agent import Agent - if current_span.context: current_span_id = current_span.context.trace_id else: @@ -64,34 +65,26 @@ def set_common_attributes( ) return - if isinstance(invocation_context.agent, Agent) and invocation_context.agent.tracers: - try: - from veadk.tracing.telemetry.opentelemetry_tracer import OpentelemetryTracer - - tracer: OpentelemetryTracer = invocation_context.agent.tracers[0] # type: ignore - spans = tracer._inmemory_exporter.processor.spans # # type: ignore - - 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: - if span.name.startswith("invocation"): - span.set_attribute("gen_ai.operation.name", "chain") - elif span.name.startswith("agent_run"): - span.set_attribute("gen_ai.operation.name", "agent") - for attr_name, attr_extractor in common_attributes.items(): - value = attr_extractor(**kwargs) - span.set_attribute(attr_name, value) - except Exception as e: - logger.error(f"Failed to set common attributes for spans: {e}") - else: - logger.warning( - "Failed to set common attributes for spans as your agent is not VeADK Agent. Skip this." - ) + try: + spans = _INMEMORY_EXPORTER_INSTANCE.processor.spans # # type: ignore + + 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: + if span.name.startswith("invocation"): + span.set_attribute("gen_ai.operation.name", "chain") + elif span.name.startswith("agent_run"): + span.set_attribute("gen_ai.operation.name", "agent") + for attr_name, attr_extractor in common_attributes.items(): + value = attr_extractor(**kwargs) + span.set_attribute(attr_name, value) + except Exception as e: + logger.error(f"Failed to set common attributes for spans: {e}") def trace_tool_call(