feat(tracing): reconstruct tracing module (#93)

* feat(tracing): reconstruct tracing module

* update attributions

* fix exporter bugs

* realize invocation and agent run attributes

* fix: fix the bugs in tracing so that it can pass the tests.

* update tracing

* fix function response

* update attributes

* fix bugs

* update attributes

* fix: fix apmplus metrics

* fix: remove do hooks of tracer

* fix: add license header

* fix: fix function_name of metrics

* fix: rewrite init meter uploader

* update attributes

* refine meter uploader

* fix attr

* add attributes

* remove useless class

---------

Co-authored-by: hanzhi.421 <[email protected]>

Author: yaozheng-fang

tests/test_agent.py

@@ -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

tests/test_tracing.py

@@ -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
 
@@ -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
@@ -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

veadk/agent.py

@@ -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'})}"
@@ -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

veadk/agents/loop_agent.py

@@ -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
@@ -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.")

veadk/agents/parallel_agent.py

@@ -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
@@ -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.")

veadk/agents/sequential_agent.py

@@ -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
@@ -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.")

veadk/integrations/ve_faas/template/{{cookiecutter.local_dir_name}}/src/app.py

@@ -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:

veadk/runner.py

@@ -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
@@ -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,
@@ -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.")
@@ -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}")