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
4 changes: 0 additions & 4 deletions tests/test_agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -48,7 +48,3 @@ def test_agent():

assert agent.long_term_memory.backend == "local"
assert load_memory in agent.tools

assert tracer.tracer_hook_before_model in agent.before_model_callback
assert tracer.tracer_hook_after_model in agent.after_model_callback
assert tracer.tracer_hook_after_tool in agent.after_tool_callback
6 changes: 3 additions & 3 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) == 4 # with extra 1 built-in exporters
assert len(tracer.exporters) == 3

# 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) == 4 # with extra 1 built-in exporters
assert len(tracer.exporters) == 3 # with extra 1 built-in exporters


@pytest.mark.asyncio
Expand All @@ -113,4 +113,4 @@ async def test_tracing_with_apmplus_global_provider():
tracer = OpentelemetryTracer(exporters=exporters)

# apmplus exporter won't init again
assert len(tracer.exporters) == 3 # with extra 1 built-in exporters
assert len(tracer.exporters) == 2 # with extra 1 built-in exporters
7 changes: 0 additions & 7 deletions veadk/agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -112,10 +112,6 @@ def model_post_init(self, __context: Any) -> None:

self.tools.append(load_memory)

if self.tracers:
for tracer in self.tracers:
tracer.do_hooks(self)

logger.info(f"{self.__class__.__name__} `{self.name}` init done.")
logger.debug(
f"Agent: {self.model_dump(include={'name', 'model_name', 'model_api_base', 'tools', 'serve_url'})}"
Expand Down Expand Up @@ -215,9 +211,6 @@ async def run(
session_service=session_service,
memory_service=self.long_term_memory,
)
if getattr(self, "tracers", None):
for tracer in self.tracers:
tracer.set_app_name(app_name)

logger.info(f"Begin to process prompt {prompt}")
# run
Expand Down
15 changes: 0 additions & 15 deletions veadk/agents/loop_agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,6 @@
from pydantic import ConfigDict, Field
from typing_extensions import Any

from veadk.agent import Agent
from veadk.prompts.agent_default_prompt import DEFAULT_DESCRIPTION, DEFAULT_INSTRUCTION
from veadk.tracing.base_tracer import BaseTracer
from veadk.utils.logger import get_logger
Expand Down Expand Up @@ -50,21 +49,7 @@ class LoopAgent(GoogleADKLoopAgent):
tracers: list[BaseTracer] = []
"""The tracers provided to agent."""

def set_sub_agents_tracer(self, tracer) -> None:
from veadk.agents.parallel_agent import ParallelAgent
from veadk.agents.sequential_agent import SequentialAgent

for sub_agent in self.sub_agents:
if isinstance(sub_agent, Agent):
tracer.do_hooks(sub_agent)
elif isinstance(sub_agent, (SequentialAgent, LoopAgent, ParallelAgent)):
sub_agent.set_sub_agents_tracer(tracer)

def model_post_init(self, __context: Any) -> None:
super().model_post_init(None) # for sub_agents init

if self.tracers:
for tracer in self.tracers:
self.set_sub_agents_tracer(tracer)

logger.info(f"{self.__class__.__name__} `{self.name}` init done.")
13 changes: 0 additions & 13 deletions veadk/agents/parallel_agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,6 @@
from pydantic import ConfigDict, Field
from typing_extensions import Any

from veadk.agent import Agent
from veadk.prompts.agent_default_prompt import DEFAULT_DESCRIPTION, DEFAULT_INSTRUCTION
from veadk.tracing.base_tracer import BaseTracer
from veadk.utils.logger import get_logger
Expand Down Expand Up @@ -50,24 +49,12 @@ class ParallelAgent(GoogleADKParallelAgent):
tracers: list[BaseTracer] = []
"""The tracers provided to agent."""

def set_sub_agents_tracer(self, tracer) -> None:
from veadk.agents.loop_agent import LoopAgent
from veadk.agents.sequential_agent import SequentialAgent

for sub_agent in self.sub_agents:
if isinstance(sub_agent, Agent):
tracer.do_hooks(sub_agent)
elif isinstance(sub_agent, (SequentialAgent, LoopAgent, ParallelAgent)):
sub_agent.set_sub_agents_tracer(tracer)

def model_post_init(self, __context: Any) -> None:
super().model_post_init(None) # for sub_agents init

if self.tracers:
logger.warning(
"Enable tracing in ParallelAgent may cause OpenTelemetry context error. Issue see https://github.com/google/adk-python/issues/1670"
)
for tracer in self.tracers:
self.set_sub_agents_tracer(tracer)

logger.info(f"{self.__class__.__name__} `{self.name}` init done.")
15 changes: 0 additions & 15 deletions veadk/agents/sequential_agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,6 @@
from pydantic import ConfigDict, Field
from typing_extensions import Any

from veadk.agent import Agent
from veadk.prompts.agent_default_prompt import DEFAULT_DESCRIPTION, DEFAULT_INSTRUCTION
from veadk.tracing.base_tracer import BaseTracer
from veadk.utils.logger import get_logger
Expand Down Expand Up @@ -50,21 +49,7 @@ class SequentialAgent(GoogleADKSequentialAgent):
tracers: list[BaseTracer] = []
"""The tracers provided to agent."""

def set_sub_agents_tracer(self, tracer) -> None:
from veadk.agents.loop_agent import LoopAgent
from veadk.agents.parallel_agent import ParallelAgent

for sub_agent in self.sub_agents:
if isinstance(sub_agent, Agent):
tracer.do_hooks(sub_agent)
elif isinstance(sub_agent, (SequentialAgent, LoopAgent, ParallelAgent)):
sub_agent.set_sub_agents_tracer(tracer)

def model_post_init(self, __context: Any) -> None:
super().model_post_init(None) # for sub_agents init

if self.tracers:
for tracer in self.tracers:
self.set_sub_agents_tracer(tracer)

logger.info(f"{self.__class__.__name__} `{self.name}` init done.")
Original file line number Diff line number Diff line change
Expand Up @@ -70,11 +70,8 @@ def load_tracer() -> None:
else:
exporters.append(exporter_cls())

tracer = OpentelemetryTracer(
name="veadk_tracer", app_name=agent_run_config.app_name, exporters=exporters
)
tracer = OpentelemetryTracer(name="veadk_tracer", exporters=exporters)
agent_run_config.agent.tracers.extend([tracer])
tracer.do_hooks(agent=agent_run_config.agent)


def build_mcp_run_agent_func() -> Callable:
Expand Down
38 changes: 22 additions & 16 deletions veadk/runner.py
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,6 @@
from veadk.agents.sequential_agent import SequentialAgent
from veadk.evaluation import EvalSetRecorder
from veadk.memory.short_term_memory import ShortTermMemory
from veadk.tracing.base_tracer import UserMessagePlugin
from veadk.types import MediaMessage
from veadk.utils.logger import get_logger
from veadk.utils.misc import read_png_to_bytes
Expand Down Expand Up @@ -68,22 +67,9 @@ def __init__(
# prevent VeRemoteAgent has no long-term memory attr
if isinstance(self.agent, Agent):
self.long_term_memory = self.agent.long_term_memory
for tracer in self.agent.tracers:
tracer.set_app_name(self.app_name)
else:
self.long_term_memory = None

# process plugins
try:
# try to detect tracer
_ = self.agent.tracers[0]
if not plugins:
plugins = [UserMessagePlugin(name="user_message_plugin")]
else:
plugins.append(UserMessagePlugin(name="user_message_plugin"))
except Exception:
logger.debug("Agent has no tracers, telemetry plugin not added.")

self.runner = ADKRunner(
app_name=self.app_name,
agent=self.agent,
Expand Down Expand Up @@ -187,7 +173,27 @@ async def run(

return final_output

def _print_trace_id(self):
def get_trace_id(self) -> str:
if not isinstance(self.agent, Agent):
logger.warning(
("The agent is not an instance of VeADK Agent, no trace id provided.")
)
return "<unknown_trace_id>"

if not self.agent.tracers:
logger.warning(
"No tracer is configured in the agent, no trace id provided."
)
return "<unknown_trace_id>"

try:
trace_id = self.agent.tracers[0].trace_id # type: ignore
return trace_id
except Exception as e:
logger.warning(f"Get tracer id failed as {e}")
return "<unknown_trace_id>"

def _print_trace_id(self) -> None:
if not isinstance(self.agent, Agent):
logger.warning(
("The agent is not an instance of VeADK Agent, no trace id provided.")
Expand All @@ -201,7 +207,7 @@ def _print_trace_id(self):
return

try:
trace_id = self.agent.tracers[0].get_trace_id() # type: ignore
trace_id = self.agent.tracers[0].trace_id # type: ignore
logger.info(f"Trace id: {trace_id}")
except Exception as e:
logger.warning(f"Get tracer id failed as {e}")
Expand Down
Loading