feat: tracer callbacks update (#48)

Author: zakahan

tests/test_agent.py

@@ -49,7 +49,6 @@ def test_agent():
     assert agent.long_term_memory.backend == "local"
     assert load_memory in agent.tools
 
-    assert tracer.llm_metrics_hook in agent.before_model_callback
-    assert tracer.token_metrics_hook in agent.after_model_callback
-    assert tracer.token_metrics_hook in agent.after_model_callback
-    assert tracer.token_metrics_hook in agent.after_model_callback
+    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

veadk/agent.py

@@ -118,9 +118,11 @@ def model_post_init(self, __context: Any) -> None:
         if self.tracers:
             self.before_model_callback = []
             self.after_model_callback = []
+            self.after_tool_callback = []
             for tracer in self.tracers:
-                self.before_model_callback.append(tracer.llm_metrics_hook)
-                self.after_model_callback.append(tracer.token_metrics_hook)
+                self.before_model_callback.append(tracer.tracer_hook_before_model)
+                self.after_model_callback.append(tracer.tracer_hook_after_model)
+                self.after_tool_callback.append(tracer.tracer_hook_after_tool)
 
         logger.info(f"Agent `{self.name}` init done.")
         logger.debug(
@@ -221,6 +223,9 @@ 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/cli/services/vefaas/template/src/app.py

@@ -49,11 +49,15 @@
     agent.before_model_callback = []
 if not getattr(agent, "after_model_callback", None):
     agent.after_model_callback = []
+if not getattr(agent, "after_tool_callback", None):
+    agent.after_tool_callback = []
 for tracer in TRACERS:
-    if tracer.llm_metrics_hook not in agent.before_model_callback:
-        agent.before_model_callback.append(tracer.llm_metrics_hook)
-    if tracer.token_metrics_hook not in agent.after_model_callback:
-        agent.after_model_callback.append(tracer.token_metrics_hook)
+    if tracer.tracer_hook_before_model not in agent.before_model_callback:
+        agent.before_model_callback.append(tracer.tracer_hook_before_model)
+    if tracer.tracer_hook_after_model not in agent.after_model_callback:
+        agent.after_model_callback.append(tracer.tracer_hook_after_model)
+    if tracer.tracer_hook_after_tool not in agent.after_tool_callback:
+        agent.after_tool_callback.append(tracer.tracer_hook_after_tool)
 
 # Tracer Config ================================================================
 # ==============================================================================

veadk/runner.py

@@ -70,6 +70,10 @@ def __init__(
             memory_service=self.long_term_memory,
         )
 
+        if getattr(self.agent, "tracers", None):
+            for tracers in self.agent.tracers:
+                tracers.set_app_name(self.app_name)
+
     def _convert_messages(self, messages) -> list:
         if isinstance(messages, str):
             messages = [types.Content(role="user", parts=[types.Part(text=messages)])]

veadk/tracing/base_tracer.py

@@ -14,10 +14,11 @@
 
 import json
 from abc import ABC, abstractmethod
-from typing import Optional
+from typing import Optional, Any
 
 from google.adk.agents.callback_context import CallbackContext
 from google.adk.models import LlmRequest, LlmResponse
+from google.adk.tools import BaseTool, ToolContext
 from opentelemetry import trace
 
 from veadk.utils.logger import get_logger
@@ -27,20 +28,24 @@
 
 class BaseTracer(ABC):
     def __init__(self, name: str):
+        self.app_name = "veadk_app_name"
         pass
 
     @abstractmethod
     def dump(self) -> str: ...
 
-    def llm_metrics_hook(
+    def tracer_hook_before_model(
         self, callback_context: CallbackContext, llm_request: LlmRequest
     ) -> Optional[LlmResponse]:
+        """agent run stage"""
         trace.get_tracer("gcp.vertex.agent")
         span = trace.get_current_span()
         # logger.debug(f"llm_request: {llm_request}")
 
         req = llm_request.model_dump()
 
+        app_name = getattr(self, "app_name", "veadk_app")
+        agent_name = callback_context.agent_name
         model_name = req.get("model", "unknown")
         max_tokens = (
             None
@@ -59,6 +64,8 @@ def llm_metrics_hook(
         )
 
         attributes = {}
+        attributes["agent.name"] = agent_name
+        attributes["app.name"] = app_name
         attributes["gen_ai.system"] = "veadk"
         if model_name:
             attributes["gen_ai.request.model"] = model_name
@@ -84,9 +91,10 @@ def llm_metrics_hook(
             for k, v in attributes.items():
                 span.set_attribute(k, v)
 
-    def token_metrics_hook(
+    def tracer_hook_after_model(
         self, callback_context: CallbackContext, llm_response: LlmResponse
     ) -> Optional[LlmResponse]:
+        """call llm stage"""
         trace.get_tracer("gcp.vertex.agent")
         span = trace.get_current_span()
         # logger.debug(f"llm_response: {llm_response}")
@@ -95,6 +103,10 @@ def token_metrics_hook(
         # Refined: collect all attributes, use set_attributes, print for debugging
         attributes = {}
 
+        app_name = getattr(self, "app_name", "veadk_app")
+        agent_name = callback_context.agent_name
+        attributes["agent.name"] = agent_name
+        attributes["app.name"] = app_name
         # prompt
         user_content = callback_context.user_content
         if getattr(user_content, "role", None):
@@ -170,3 +182,32 @@ def token_metrics_hook(
             # Fallback for OpenTelemetry versions without set_attributes
             for k, v in attributes.items():
                 span.set_attribute(k, v)
+
+    def tracer_hook_after_tool(
+        self,
+        tool: BaseTool,
+        args: dict[str, Any],
+        tool_context: ToolContext,
+        tool_response: dict,
+    ):
+        trace.get_tracer("gcp.vertex.agent")
+        span = trace.get_current_span()
+        agent_name = tool_context.agent_name
+        tool_name = tool.name
+        app_name = getattr(self, "app_name", "veadk_app")
+        attributes = {
+            "agent.name": agent_name,
+            "app.name": app_name,
+            "tool.name": tool_name,
+        }
+
+        # Set all attributes at once if possible, else fallback to individual
+        if hasattr(span, "set_attributes"):
+            span.set_attributes(attributes)
+        else:
+            # Fallback for OpenTelemetry versions without set_attributes
+            for k, v in attributes.items():
+                span.set_attribute(k, v)
+
+    def set_app_name(self, app_name):
+        self.app_name = app_name

veadk/tracing/telemetry/opentelemetry_tracer.py

@@ -51,6 +51,11 @@ class OpentelemetryTracer(BaseModel, BaseTracer):
         DEFAULT_VEADK_TRACER_NAME, description="The identifier of tracer."
     )
 
+    app_name: str = Field(
+        "veadk_app",
+        description="The identifier of app.",
+    )
+
     def model_post_init(self, context: Any, /) -> None:
         self._processors = []
         self._inmemory_exporter: InMemoryExporter = None