feat(tracing): reconstruct tracing module

Author: yaozheng-fang

veadk/agent.py

@@ -112,9 +112,9 @@ 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)
+        # if self.tracers:
+        #     for tracer in self.tracers:
+        #         tracer.do_hooks(self)
 
         logger.info(f"{self.__class__.__name__} `{self.name}` init done.")
         logger.debug(

veadk/runner.py

@@ -68,8 +68,8 @@ 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)
+            # for tracer in self.agent.tracers:
+            #     tracer.set_app_name(self.app_name)
         else:
             self.long_term_memory = None
 

veadk/tracing/base_tracer.py

@@ -12,16 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import json
 from abc import ABC, abstractmethod
-from typing import Any, Optional
+from typing import 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
 
@@ -102,201 +97,9 @@ def replace_bytes_with_empty(data):
 
 class BaseTracer(ABC):
     def __init__(self, name: str):
-        self.app_name = "veadk_app_name"
-        pass
+        self.name = name
+        self._trace_id = "<unknown_trace_id>"
+        self._trace_file_path = "<unknown_trace_file_path>"
 
     @abstractmethod
     def dump(self, user_id: str, session_id: str, path: str = "/tmp") -> str: ...
-
-    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
-            if not req.get("live_connect_config")
-            else req["live_connect_config"].get("max_output_tokens", None)
-        )
-        temperature = (
-            None
-            if not req.get("live_connect_config")
-            else req["live_connect_config"].get("temperature", None)
-        )
-        top_p = (
-            None
-            if not req.get("live_connect_config")
-            else req["live_connect_config"].get("top_p", None)
-        )
-
-        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
-            attributes["gen_ai.response.model"] = (
-                model_name  # The req model and the resp model should be consistent.
-            )
-        attributes["gen_ai.request.type"] = "completion"
-        if max_tokens:
-            attributes["gen_ai.request.max_tokens"] = max_tokens
-        if temperature:
-            attributes["gen_ai.request.temperature"] = temperature
-        if top_p:
-            attributes["gen_ai.request.top_p"] = top_p
-
-        # Print attributes for debugging
-        # print("Tracing attributes:", attributes)
-
-        # 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 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}")
-        # logger.debug(f"callback_context: {callback_context}")
-
-        # 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
-        attributes["gen_ai.system"] = "veadk"
-        # prompt
-        user_content = callback_context.user_content
-        role = None
-        content = None
-        if getattr(user_content, "role", None):
-            role = getattr(user_content, "role", None)
-
-        if user_content and getattr(user_content, "parts", None):
-            # content = user_content.model_dump_json(exclude_none=True)
-            content = user_content.model_dump(exclude_none=True).get("parts", None)
-            if content:
-                content = replace_bytes_with_empty(content)
-            content = json.dumps(content, ensure_ascii=False) if content else None
-
-        if role and content:
-            attributes["gen_ai.prompt.0.role"] = role
-            attributes["gen_ai.prompt.0.content"] = content
-
-        # completion
-        completion_content = getattr(llm_response, "content").model_dump(
-            exclude_none=True
-        )
-        if completion_content:
-            content = json.dumps(
-                getattr(llm_response, "content").model_dump(exclude_none=True)["parts"]
-            )
-            role = getattr(llm_response, "content").model_dump(exclude_none=True)[
-                "role"
-            ]
-            if role and content:
-                attributes["gen_ai.completion.0.role"] = role
-                attributes["gen_ai.completion.0.content"] = content
-
-        if not llm_response.usage_metadata:
-            return
-
-        # tokens
-        metadata = llm_response.usage_metadata.model_dump()
-        if metadata:
-            prompt_tokens = metadata.get("prompt_token_count", None)
-            completion_tokens = metadata.get("candidates_token_count", None)
-            total_tokens = metadata.get("total_token_count", None)
-            cache_read_input_tokens = (
-                metadata.get("cache_read_input_tokens") or 0
-            )  # Might change, once openai introduces their equivalent.
-            cache_create_input_tokens = (
-                metadata.get("cache_create_input_tokens") or 0
-            )  # Might change, once openai introduces their equivalent.
-            if prompt_tokens:
-                attributes["gen_ai.usage.prompt_tokens"] = prompt_tokens
-            if completion_tokens:
-                attributes["gen_ai.usage.completion_tokens"] = completion_tokens
-            if total_tokens:
-                attributes["gen_ai.usage.total_tokens"] = total_tokens
-            if cache_read_input_tokens is not None:
-                attributes["gen_ai.usage.cache_read_input_tokens"] = (
-                    cache_read_input_tokens
-                )
-            if cache_create_input_tokens is not None:
-                attributes["gen_ai.usage.cache_create_input_tokens"] = (
-                    cache_create_input_tokens
-                )
-
-        # Print attributes for debugging
-        # print("Tracing attributes:", attributes)
-
-        # 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 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,
-            "gen_ai.system": "veadk",
-        }
-
-        # 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
-
-    def do_hooks(self, agent) -> None:
-        if not getattr(agent, "before_model_callback", None):
-            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 = []
-
-        if self.tracer_hook_before_model not in agent.before_model_callback:
-            agent.before_model_callback.append(self.tracer_hook_before_model)
-        if self.tracer_hook_after_model not in agent.after_model_callback:
-            agent.after_model_callback.append(self.tracer_hook_after_model)
-        if self.tracer_hook_after_tool not in agent.after_tool_callback:
-            agent.after_tool_callback.append(self.tracer_hook_after_tool)

veadk/tracing/telemetry/attributes/attributes.py

@@ -0,0 +1,32 @@
+from veadk.tracing.telemetry.attributes.extractors.common_attributes_extractors import (
+    common_gen_ai_app_name,
+    common_gen_ai_session_id,
+    common_gen_ai_system,
+    common_gen_ai_system_version,
+    common_gen_ai_user_id,
+)
+from veadk.tracing.telemetry.attributes.extractors.llm_attributes_extrators import (
+    llm_gen_ai_completion,
+    llm_gen_ai_prompt,
+    llm_gen_ai_request_model,
+    llm_gen_ai_request_type,
+    llm_gen_ai_response_model,
+)
+
+ATTRIBUTES = {
+    "common": {
+        "gen_ai.system": common_gen_ai_system,
+        "gen_ai.system_version": common_gen_ai_system_version,
+        "gen_ai.app.name": common_gen_ai_app_name,
+        "gen_ai.user.id": common_gen_ai_user_id,
+        "gen_ai.session.id": common_gen_ai_session_id,
+    },
+    "llm": {
+        "gen_ai.request.model": llm_gen_ai_request_model,
+        "gen_ai.request.type": llm_gen_ai_request_type,
+        "gen_ai.response.model": llm_gen_ai_response_model,
+        "gen_ai.prompt": llm_gen_ai_prompt,
+        "gen_ai.completion": llm_gen_ai_completion,
+    },
+    "tool": ...,
+}

veadk/tracing/telemetry/attributes/extractors/common_attributes_extractors.py

@@ -0,0 +1,21 @@
+from veadk.version import VERSION
+
+
+def common_gen_ai_system() -> str:
+    return "veadk"
+
+
+def common_gen_ai_system_version() -> str:
+    return VERSION
+
+
+def common_gen_ai_app_name(**kwargs) -> str:
+    return kwargs.get("app_name", "<unknown_app_name>")
+
+
+def common_gen_ai_user_id(**kwargs) -> str:
+    return kwargs.get("user_id", "")
+
+
+def common_gen_ai_session_id(**kwargs) -> str:
+    return kwargs.get("session_id", "<unknown_session_id>")

veadk/tracing/telemetry/attributes/extractors/llm_attributes_extrators.py

@@ -0,0 +1,65 @@
+from attr import dataclass
+from google.adk.agents.invocation_context import InvocationContext
+from google.adk.models.llm_request import LlmRequest
+from google.adk.models.llm_response import LlmResponse
+
+
+@dataclass
+class LLMAttributesParams:
+    invocation_context: InvocationContext
+    event_id: str
+    llm_request: LlmRequest
+    llm_response: LlmResponse
+
+
+def llm_gen_ai_request_model(params: LLMAttributesParams) -> str:
+    return params.llm_request.model or "<unknown_model_name>"
+
+
+def llm_gen_ai_request_type(params: LLMAttributesParams) -> str | list[str]:
+    type = "completion"
+    return type or "<unknown_type>"
+
+
+def llm_gen_ai_response_model(params: LLMAttributesParams) -> str:
+    return params.llm_request.model or "<unknown_model_name>"
+
+
+def llm_gen_ai_request_max_tokens(params: LLMAttributesParams) -> int | None:
+    return params.llm_request.config.max_output_tokens
+
+
+def llm_gen_ai_request_temperature(params: LLMAttributesParams) -> float | None:
+    return params.llm_request.config.temperature
+
+
+def llm_gen_ai_request_top_p(params: LLMAttributesParams) -> float | None:
+    return params.llm_request.config.top_p
+
+
+def llm_gen_ai_prompt(params: LLMAttributesParams) -> list[dict]:
+    ret = []
+    for idx, content in enumerate(params.llm_request.contents):
+        if content.parts:
+            role = content.role
+            parts = [part for part in content.parts if not part.inline_data]
+        ret.append({f".{idx}.role": role, f".{idx}.content": str(parts)})
+    return ret
+
+
+def llm_gen_ai_completion(params: LLMAttributesParams) -> list[dict] | None:
+    ret = []
+
+    content = params.llm_response.content
+    if content and content.parts:
+        parts = [part for part in content.parts if not part.inline_data]
+        ret.append({f".{0}.role": content.role, f".{0}.content": str(parts)})
+        return ret
+
+
+def llm_gen_ai_response_stop_reason(params: LLMAttributesParams) -> str | None:
+    return params.llm_response.stop_reason
+
+
+def llm_gen_ai_response_finish_reason(params: LLMAttributesParams) -> str | None:
+    return params.llm_response.finish_reason