From 60eba8a2a547e37c586f836e60a048cc83b5a325 Mon Sep 17 00:00:00 2001 From: "hanzhi.421" Date: Wed, 13 Aug 2025 10:46:40 +0800 Subject: [PATCH] feat: tracer callbacks update --- tests/test_agent.py | 7 ++- veadk/agent.py | 9 +++- veadk/cli/services/vefaas/template/src/app.py | 12 +++-- veadk/runner.py | 4 ++ veadk/tracing/base_tracer.py | 47 +++++++++++++++++-- .../tracing/telemetry/opentelemetry_tracer.py | 5 ++ 6 files changed, 71 insertions(+), 13 deletions(-) diff --git a/tests/test_agent.py b/tests/test_agent.py index f7634f1e..622e5caa 100644 --- a/tests/test_agent.py +++ b/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 diff --git a/veadk/agent.py b/veadk/agent.py index 16472133..ffbf8722 100644 --- a/veadk/agent.py +++ b/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 diff --git a/veadk/cli/services/vefaas/template/src/app.py b/veadk/cli/services/vefaas/template/src/app.py index 0a2d0c09..23703a40 100644 --- a/veadk/cli/services/vefaas/template/src/app.py +++ b/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 ================================================================ # ============================================================================== diff --git a/veadk/runner.py b/veadk/runner.py index f0d6db50..8d6c297a 100644 --- a/veadk/runner.py +++ b/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)])] diff --git a/veadk/tracing/base_tracer.py b/veadk/tracing/base_tracer.py index 04638dd5..780c58dc 100644 --- a/veadk/tracing/base_tracer.py +++ b/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 diff --git a/veadk/tracing/telemetry/opentelemetry_tracer.py b/veadk/tracing/telemetry/opentelemetry_tracer.py index b82dc7c7..a1b1da1c 100644 --- a/veadk/tracing/telemetry/opentelemetry_tracer.py +++ b/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