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
2 changes: 1 addition & 1 deletion tests/test_agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -88,7 +88,7 @@ def test_agent_default_values():
assert agent.sub_agents == []
assert agent.knowledgebase is None
assert agent.long_term_memory is None
assert agent.tracers == []
# assert agent.tracers == []


@patch.dict("os.environ", {"MODEL_AGENT_API_KEY": "mock_api_key"})
Expand Down
8 changes: 4 additions & 4 deletions tests/test_tracing.py
Original file line number Diff line number Diff line change
Expand Up @@ -82,7 +82,7 @@ async def test_tracing():
exporters = init_exporters()
tracer = OpentelemetryTracer(exporters=exporters)

assert len(tracer.exporters) == 3
assert len(tracer.exporters) == 4 # with extra 1 built-in exporters

# TODO: Ensure the tracing provider is set correctly after loading SDK

Expand All @@ -98,7 +98,7 @@ async def test_tracing_with_global_provider():
#
tracer = OpentelemetryTracer(exporters=exporters)

assert len(tracer.exporters) == 3 # with extra 1 built-in exporters
assert len(tracer.exporters) == 4 # with extra 1 built-in exporters


@pytest.mark.asyncio
Expand All @@ -112,5 +112,5 @@ async def test_tracing_with_apmplus_global_provider():
# init OpentelemetryTracer
tracer = OpentelemetryTracer(exporters=exporters)

# apmplus exporter won't init again
assert len(tracer.exporters) == 2 # with extra 1 built-in exporters
# apmplus exporter won't init again, so there are cozeloop, tls, in_memory exporter
assert len(tracer.exporters) == 3 # with extra 1 built-in exporters
39 changes: 39 additions & 0 deletions veadk/agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@

from __future__ import annotations

import os
from typing import Optional, Union

from google.adk.agents import LlmAgent, RunConfig
Expand Down Expand Up @@ -132,6 +133,8 @@ def model_post_init(self, __context: Any) -> None:
"You are trying to use your own LiteLLM client, some default request headers may be missing."
)

self._prepare_tracers()

if self.knowledgebase:
from veadk.tools.builtin_tools.load_knowledgebase import (
LoadKnowledgebaseTool,
Expand Down Expand Up @@ -194,6 +197,42 @@ async def event_generator():

return final_output

def _prepare_tracers(self):
enable_apmplus_tracer = os.getenv("ENABLE_APMPLUS", "false").lower() == "true"
enable_cozeloop_tracer = os.getenv("ENABLE_COZELOOP", "false").lower() == "true"
enable_tls_tracer = os.getenv("ENABLE_TLS", "false").lower() == "true"

if not self.tracers:
from veadk.tracing.telemetry.opentelemetry_tracer import OpentelemetryTracer

self.tracers.append(OpentelemetryTracer())

exporters = self.tracers[0].exporters # type: ignore

from veadk.tracing.telemetry.exporters.apmplus_exporter import APMPlusExporter
from veadk.tracing.telemetry.exporters.cozeloop_exporter import CozeloopExporter
from veadk.tracing.telemetry.exporters.tls_exporter import TLSExporter

if enable_apmplus_tracer and not any(
isinstance(e, APMPlusExporter) for e in exporters
):
self.tracers[0].exporters.append(APMPlusExporter()) # type: ignore
logger.info("Enable APMPlus exporter by env.")

if enable_cozeloop_tracer and not any(
isinstance(e, CozeloopExporter) for e in exporters
):
self.tracers[0].exporters.append(CozeloopExporter()) # type: ignore
logger.info("Enable CozeLoop exporter by env.")

if enable_tls_tracer and not any(isinstance(e, TLSExporter) for e in exporters):
self.tracers[0].exporters.append(TLSExporter()) # type: ignore
logger.info("Enable TLS exporter by env.")

logger.debug(
f"Opentelemetry Tracer init {len(self.tracers[0].exporters)} exporters" # type: ignore
)

async def run(
self,
prompt: str | list[str],
Expand Down
9 changes: 6 additions & 3 deletions veadk/tracing/telemetry/opentelemetry_tracer.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@
from opentelemetry import trace as trace_api
from opentelemetry.sdk import trace as trace_sdk
from opentelemetry.sdk.resources import Resource
from opentelemetry.sdk.trace import TracerProvider, SpanLimits
from opentelemetry.sdk.trace import SpanLimits, TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor, SimpleSpanProcessor
from pydantic import BaseModel, ConfigDict, Field, field_validator
from typing_extensions import override
Expand All @@ -31,8 +31,8 @@
from veadk.tracing.telemetry.exporters.base_exporter import BaseExporter
from veadk.tracing.telemetry.exporters.inmemory_exporter import InMemoryExporter
from veadk.utils.logger import get_logger
from veadk.utils.patches import patch_google_adk_telemetry
from veadk.utils.misc import get_temp_dir
from veadk.utils.patches import patch_google_adk_telemetry

logger = get_logger(__name__)

Expand Down Expand Up @@ -131,12 +131,15 @@ def _init_global_tracer_provider(self) -> None:
) + global_tracer_provider._active_span_processor._span_processors

self._processors.append(self._inmemory_exporter.processor)
self.exporters.append(self._inmemory_exporter)
else:
logger.warning(
"InMemoryExporter processor is not initialized, cannot add to OpentelemetryTracer."
)

logger.info(f"Init OpentelemetryTracer with {len(self._processors)} exporters.")
logger.info(
f"Init OpentelemetryTracer with {len(self._processors)} exporter(s)."
)

@property
def trace_file_path(self) -> str:
Expand Down