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
7 changes: 3 additions & 4 deletions tests/test_agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
9 changes: 7 additions & 2 deletions veadk/agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down Expand Up @@ -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
Expand Down
12 changes: 8 additions & 4 deletions veadk/cli/services/vefaas/template/src/app.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 ================================================================
# ==============================================================================
Expand Down
4 changes: 4 additions & 0 deletions veadk/runner.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)])]
Expand Down
47 changes: 44 additions & 3 deletions veadk/tracing/base_tracer.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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}")
Expand All @@ -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):
Expand Down Expand Up @@ -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
5 changes: 5 additions & 0 deletions veadk/tracing/telemetry/opentelemetry_tracer.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down