Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions veadk/tracing/telemetry/exporters/inmemory_exporter.py
Original file line number Diff line number Diff line change
Expand Up @@ -106,3 +106,6 @@ def __init__(self, name: str = "inmemory_exporter") -> None:

self._exporter = _InMemoryExporter()
self.processor = _InMemorySpanProcessor(self._exporter)


_INMEMORY_EXPORTER_INSTANCE = InMemoryExporter()
5 changes: 4 additions & 1 deletion veadk/tracing/telemetry/opentelemetry_tracer.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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
Expand Down
53 changes: 23 additions & 30 deletions veadk/tracing/telemetry/telemetry.py
Original file line number Diff line number Diff line change
Expand Up @@ -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__)
Expand All @@ -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:
Expand All @@ -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(
Expand Down