support tool call record metrics

Author: yaozheng-fang

veadk/tracing/telemetry/exporters/apmplus_exporter.py

@@ -17,8 +17,10 @@
 from typing import Any
 
 from google.adk.agents.invocation_context import InvocationContext
+from google.adk.events import Event
 from google.adk.models.llm_request import LlmRequest
 from google.adk.models.llm_response import LlmResponse
+from google.adk.tools import BaseTool
 from opentelemetry import metrics, trace
 from opentelemetry import metrics as metrics_api
 from opentelemetry.exporter.otlp.proto.grpc.metric_exporter import OTLPMetricExporter
@@ -195,7 +197,7 @@ def __init__(
             explicit_bucket_boundaries_advisory=_GEN_AI_SERVER_TIME_PER_OUTPUT_TOKEN_BUCKETS,
         )
 
-    def record(
+    def record_call_llm(
         self,
         invocation_context: InvocationContext,
         event_id: str,
@@ -267,6 +269,14 @@ def record(
             #         time_per_output_token, attributes=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(

veadk/tracing/telemetry/telemetry.py

@@ -33,8 +33,10 @@
 
 logger = get_logger(__name__)
 
+meter_uploader = None
 
-def _upload_metrics(
+
+def _upload_call_llm_metrics(
     invocation_context: InvocationContext,
     event_id: str,
     llm_request: LlmRequest,
@@ -47,11 +49,27 @@ def _upload_metrics(
         for tracer in tracers:
             for exporter in getattr(tracer, "exporters", []):
                 if getattr(exporter, "meter_uploader", None):
-                    exporter.meter_uploader.record(
+                    global meter_uploader
+                    meter_uploader = exporter.meter_uploader
+                    exporter.meter_uploader.record_call_llm(
                         invocation_context, event_id, llm_request, llm_response
                     )
 
 
+def _upload_tool_call_metrics(
+    tool: BaseTool,
+    args: dict[str, Any],
+    function_response_event: Event,
+):
+    global meter_uploader
+    if meter_uploader:
+        meter_uploader.record_tool_call(tool, args, function_response_event)
+    else:
+        logger.warning(
+            "Meter uploader is not initialized yet. Skip recording tool call metrics."
+        )
+
+
 def _set_agent_input_attribute(
     span: Span, invocation_context: InvocationContext
 ) -> None:
@@ -203,6 +221,8 @@ def trace_tool_call(
         response: ExtractorResponse = attr_extractor(params)
         ExtractorResponse.update_span(span, attr_name, response)
 
+    _upload_tool_call_metrics(tool, args, function_response_event)
+
 
 def trace_call_llm(
     invocation_context: InvocationContext,
@@ -252,7 +272,7 @@ def trace_call_llm(
         response: ExtractorResponse = attr_extractor(params)
         ExtractorResponse.update_span(span, attr_name, response)
 
-    _upload_metrics(invocation_context, event_id, llm_request, llm_response)
+    _upload_call_llm_metrics(invocation_context, event_id, llm_request, llm_response)
 
 
 # Do not modify this function