feat: support invocation attrs

Author: yaozheng-fang

veadk/agent.py

@@ -114,13 +114,8 @@ def model_post_init(self, __context: Any) -> None:
             self.tools.append(load_memory)
 
         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.tracer_hook_before_model)
-                self.after_model_callback.append(tracer.tracer_hook_after_model)
-                self.after_tool_callback.append(tracer.tracer_hook_after_tool)
+                tracer.do_hooks(self)
 
         logger.info(f"Agent `{self.name}` init done.")
         logger.debug(

veadk/runner.py

@@ -15,6 +15,7 @@
 
 from google.adk.agents import RunConfig
 from google.adk.agents.run_config import StreamingMode
+from google.adk.plugins.base_plugin import BasePlugin
 from google.adk.runners import Runner as ADKRunner
 from google.genai import types
 from google.genai.types import Blob
@@ -23,6 +24,7 @@
 from veadk.agent import Agent
 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
@@ -44,6 +46,7 @@ def __init__(
         self,
         agent: Agent | RemoteVeAgent,
         short_term_memory: ShortTermMemory,
+        plugins: list[BasePlugin] = [],
         app_name: str = "veadk_default_app",
         user_id: str = "veadk_default_user",
     ):
@@ -65,11 +68,20 @@ def __init__(
         else:
             self.long_term_memory = None
 
+        # process plugins
+        try:
+            # try to detect tracer
+            _ = self.agent.tracers[0]
+            plugins.extend([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,
             session_service=self.session_service,
             memory_service=self.long_term_memory,
+            plugins=plugins,
         )
 
     def _convert_messages(self, messages) -> list:

veadk/tracing/base_tracer.py

@@ -17,16 +17,63 @@
 from typing import Any, Optional
 
 from google.adk.agents.callback_context import CallbackContext
+from google.adk.agents.invocation_context import InvocationContext
 from google.adk.models.llm_request import LlmRequest
 from google.adk.models.llm_response import LlmResponse
+from google.adk.plugins.base_plugin import BasePlugin
 from google.adk.tools import BaseTool, ToolContext
+from google.genai import types
 from opentelemetry import trace
 
 from veadk.utils.logger import get_logger
 
 logger = get_logger(__name__)
 
 
+class UserMessagePlugin(BasePlugin):
+    def __init__(self, name: str):
+        super().__init__(name)
+
+    async def on_user_message_callback(
+        self,
+        *,
+        invocation_context: InvocationContext,
+        user_message: types.Content,
+    ) -> Optional[types.Content]:
+        """Callback executed when a user message is received before an invocation starts.
+
+        This callback helps logging and modifying the user message before the
+        runner starts the invocation.
+
+        Args:
+        invocation_context: The context for the entire invocation.
+        user_message: The message content input by user.
+
+        Returns:
+        An optional `types.Content` to be returned to the ADK. Returning a
+        value to replace the user message. Returning `None` to proceed
+        normally.
+        """
+        trace.get_tracer("gcp.vertex.agent")
+        span = trace.get_current_span()
+
+        logger.debug(f"User message plugin works, catch {span}")
+        span_name = getattr(span, "name", None)
+        if span_name and span_name.startswith("invocation"):
+            agent_name = invocation_context.agent.name
+            invoke_branch = (
+                invocation_context.branch if invocation_context.branch else agent_name
+            )
+            span.set_attribute("agent_name", agent_name)
+            span.set_attribute("invoke_branch", invoke_branch)
+
+            logger.debug(
+                f"Add attributes to {span_name}: agent_name={agent_name}, invoke_branch={invoke_branch}"
+            )
+
+        return None
+
+
 class BaseTracer(ABC):
     def __init__(self, name: str):
         self.app_name = "veadk_app_name"

veadk/tracing/telemetry/opentelemetry_tracer.py

@@ -51,11 +51,11 @@ class OpentelemetryTracer(BaseModel, BaseTracer):
         description="The exporters to export spans.",
     )
     name: str = Field(
-        DEFAULT_VEADK_TRACER_NAME, description="The identifier of tracer."
+        default=DEFAULT_VEADK_TRACER_NAME, description="The identifier of tracer."
     )
 
     app_name: str = Field(
-        "veadk_app",
+        default="veadk_app",
         description="The identifier of app.",
     )