feat: add tool metrics

Author: feng-95

agent.py

@@ -140,7 +140,9 @@ def model_post_init(self, __context: Any) -> None:
 
         if self.long_term_memory is not None:
             from google.adk.tools import load_memory
-
+            if not load_memory.custom_metadata:
+                load_memory.custom_metadata = {}
+            load_memory.custom_metadata["backend"] = self.long_term_memory.backend
             self.tools.append(load_memory)
 
         logger.info(f"VeADK version: {VERSION}")

deep_eval_test.py

@@ -42,7 +42,7 @@
 }
 
 DEFAULT_APMPLUS_OTEL_EXPORTER_ENDPOINT = "http://apmplus-cn-beijing.volces.com:4317"
-DEFAULT_APMPLUS_OTEL_EXPORTER_SERVICE_NAME = "veadk_tracing_demo"
+DEFAULT_APMPLUS_OTEL_EXPORTER_SERVICE_NAME = "veadk_tracing"
 
 DEFAULT_COZELOOP_OTEL_EXPORTER_ENDPOINT = (
     "https://api.coze.cn/v1/loop/opentelemetry/v1/traces"

uv.lock

@@ -25,6 +25,9 @@
 
 from veadk.knowledgebase import KnowledgeBase
 from veadk.knowledgebase.entry import KnowledgebaseEntry
+from veadk.utils.logger import get_logger
+
+logger = get_logger(__name__)
 
 if TYPE_CHECKING:
     from google.adk.models.llm_request import LlmRequest
@@ -96,6 +99,15 @@ class LoadKnowledgebaseTool(FunctionTool):
 
     def __init__(self):
         super().__init__(load_knowledgebase)
+        global knowledgebase
+        if knowledgebase is None:
+            logger.info(
+                "Get global knowledgebase instance failed, failed to set knowledgebase tool backend."
+            )
+        else:
+            if not self.custom_metadata:
+                self.custom_metadata = {}
+            self.custom_metadata["backend"] = knowledgebase.backend
 
     @override
     def _get_declaration(self) -> types.FunctionDeclaration | None:

veadk/agent.py

@@ -127,6 +127,11 @@ class Meters:
         "gen_ai.chat_completions.streaming_time_per_output_token"
     )
 
+    # apmplus metrics
+    # span duration
+    APMPLUS_SPAN_LATENCY = "apmplus_span_latency"
+    # tool token usage
+    APMPLUS_TOOL_TOKEN_USAGE = "apmplus_tool_token_usage"
 
 class MeterUploader:
     def __init__(
@@ -289,21 +294,6 @@ def record_call_llm(
             span = trace.get_current_span()
             if not span:
                 return
-            # operation = "tool"
-            # operation_type = "load_knowledgebase"
-            # # 能否获取到 backend 信息？
-            # operation_backend = "knowledgebase"
-            # if random.randint(0, 1) == 0:
-            #     operation_type = "load_memory"
-            #     operation_backend = "mem0"
-            # span_name_attributes = {**attributes,
-            # #                        "name": span_name,
-            #                         "gen_ai_kind": "tool",
-            #                         "gen_ai_operation_name": operation,
-            #                         "gen_ai_operation_type": operation_type,
-            #                         "gen_ai_operation_backend": operation_backend,
-
-            #                        }
 
             # record span latency
             if hasattr(span, "start_time") and self.apmplus_span_latency:
@@ -313,62 +303,54 @@ def record_call_llm(
                     duration, attributes=attributes
                 )
 
-            # # TODO: trigger with tool record
-            # if self.apmplus_tool_token_usage:
-            #     tool_token_usage = 122 # tool token 数量，使用文本长度/4
-            #     # TODO: 设置 token_type: input, output
-            #     tool_token_attributes = {**span_name_attributes, "token_type": "input"}
-            #
-            #     self.apmplus_tool_token_usage.record(
-            #         tool_token_usage, attributes=tool_token_attributes
-            #     )
-    def record_tool(self):
-        # TODO: trigger with tool record
-        attributes = {
-            "gen_ai_system": "volcengine",
-        }  # required by Volcengine APMPlus
-        # add span name attribute
+    def record_tool_call(
+            self,
+            tool: BaseTool,
+            args: dict[str, Any],
+            function_response_event: Event,
+    ):
+        logger.debug(f"Record tool call work in progress. Tool: {tool.name}")
         span = trace.get_current_span()
         if not span:
             return
         operation = "tool"
-        operation_type = "load_knowledgebase"
-        # 能否获取到 backend 信息？
-        operation_backend = "knowledgebase"
-        if random.randint(0, 1) == 0:
-            operation_type = "load_memory"
-            operation_backend = "mem0"
-        span_name_attributes = {**attributes,
-                                #                        "name": span_name,
-                                "gen_ai_kind": "tool",
-                                "gen_ai_operation_name": operation,
-                                "gen_ai_operation_type": operation_type,
-                                "gen_ai_operation_backend": operation_backend,
-
-                                }
+        operation_type = tool.name
+        operation_backend = ""
+        if tool.custom_metadata:
+            operation_backend = tool.custom_metadata.get("backend", "")
+        # if random.randint(0, 1) == 0:
+        #     operation_type = "load_memory"
+        #     operation_backend = "mem0"
+        attributes = {
+            "gen_ai_operation_name": operation,
+            "gen_ai_operation_type": operation_type,
+            "gen_ai_operation_backend": operation_backend,
+        }
         if hasattr(span, "start_time") and self.apmplus_span_latency:
             # span 耗时
             duration = (time.time_ns() - span.start_time) / 1e9  # type: ignore
             self.apmplus_span_latency.record(
-                duration, attributes=span_name_attributes
+                duration, attributes=attributes
             )
-        if self.apmplus_tool_token_usage:
-            tool_token_usage = 122  # tool token 数量，使用文本长度/4
-            # TODO: 设置 token_type: input, output
-            tool_token_attributes = {**span_name_attributes, "token_type": "input"}
 
+        if self.apmplus_tool_token_usage and hasattr(span, "attributes"):
+            tool_input = span.attributes["gen_ai.tool.input"]
+            tool_token_usage = len(tool_input) / 4  # tool token 数量，使用文本长度/4 # tool token 数量，使用文本长度/4
+            input_tool_token_attributes = {**attributes, "token_type": "input"}
             self.apmplus_tool_token_usage.record(
-                tool_token_usage, attributes=tool_token_attributes
+                tool_token_usage, attributes=input_tool_token_attributes
             )
 
+            tool_output = span.attributes["gen_ai.tool.output"]
+            tool_token_usage = len(tool_output) / 4  # tool token 数量，使用文本长度/4 # tool token 数量，使用文本长度/4
+            output_tool_token_attributes = {**attributes, "token_type": "output"}
+            self.apmplus_tool_token_usage.record(
+                tool_token_usage, attributes=output_tool_token_attributes
+            )
+
+
+
 
-    def record_tool_call(
-        self,
-        tool: BaseTool,
-        args: dict[str, Any],
-        function_response_event: Event,
-    ):
-        logger.debug(f"Record tool call work in progress. Tool: {tool.name}")
 
 class APMPlusExporterConfig(BaseModel):
     endpoint: str = Field(