refine meter uploader

Author: yaozheng-fang

veadk/tracing/telemetry/exporters/apmplus_exporter.py

@@ -14,9 +14,13 @@
 
 from typing import Any
 
+from google.adk.models.llm_request import LlmRequest
+from google.adk.models.llm_response import LlmResponse
+from opentelemetry import metrics
 from opentelemetry import metrics as metrics_api
 from opentelemetry.exporter.otlp.proto.grpc.metric_exporter import OTLPMetricExporter
 from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
+from opentelemetry.metrics._internal import Meter
 from opentelemetry.sdk import metrics as metrics_sdk
 from opentelemetry.sdk.metrics.export import PeriodicExportingMetricReader
 from opentelemetry.sdk.resources import Resource
@@ -26,12 +30,76 @@
 
 from veadk.config import getenv
 from veadk.tracing.telemetry.exporters.base_exporter import BaseExporter
-from veadk.tracing.telemetry.metrics.opentelemetry_metrics import MeterUploader
 from veadk.utils.logger import get_logger
 
 logger = get_logger(__name__)
 
 
+class MeterUploader:
+    def __init__(
+        self, name: str, endpoint: str, headers: dict, resource_attributes: dict
+    ) -> None:
+        # global_metrics_provider -> global_tracer_provider
+        # exporter -> exporter
+        # metric_reader -> processor
+        global_metrics_provider = metrics_api.get_meter_provider()
+
+        # 1. init resource
+        if hasattr(global_metrics_provider, "_sdk_config"):
+            global_resource = global_metrics_provider._sdk_config.resource  # type: ignore
+        else:
+            global_resource = Resource.create()
+
+        resource = global_resource.merge(Resource.create(resource_attributes))
+
+        # 2. init exporter and reader
+        exporter = OTLPMetricExporter(endpoint=endpoint, headers=headers)
+        metric_reader = PeriodicExportingMetricReader(exporter)
+
+        metrics_api.set_meter_provider(
+            metrics_sdk.MeterProvider(metric_readers=[metric_reader], resource=resource)
+        )
+
+        # 3. init meter
+        self.meter: Meter = metrics.get_meter(name=name)
+
+        # create meter attributes
+        self.llm_invoke_counter = self.meter.create_counter(
+            name="gen_ai.chat.count",
+            description="Number of LLM invocations",
+            unit="count",
+        )
+        self.token_usage = self.meter.create_histogram(
+            name="gen_ai.client.token.usage",
+            description="Token consumption of LLM invocations",
+            unit="count",
+        )
+
+    def record(self, llm_request: LlmRequest, llm_response: LlmResponse) -> None:
+        attributes = {
+            "gen_ai_system": "volcengine",
+            "gen_ai_response_model": llm_request.model,
+            "gen_ai_operation_name": "chat_completions",
+            "stream": "false",
+            "server_address": "api.volcengine.com",
+        }  # required by Volcengine APMPlus
+
+        if llm_response.usage_metadata:
+            # llm invocation number += 1
+            self.llm_invoke_counter.add(1, attributes)
+
+            # upload token usage
+            input_token = llm_response.usage_metadata.prompt_token_count
+            output_token = llm_response.usage_metadata.candidates_token_count
+
+            if input_token:
+                token_attributes = {**attributes, "gen_ai_token_type": "input"}
+                self.token_usage.record(input_token, attributes=token_attributes)
+            if output_token:
+                token_attributes = {**attributes, "gen_ai_token_type": "output"}
+                self.token_usage.record(output_token, attributes=token_attributes)
+
+
 class APMPlusExporterConfig(BaseModel):
     endpoint: str = Field(
         default_factory=lambda: getenv(
@@ -69,25 +137,13 @@ def model_post_init(self, context: Any) -> None:
             endpoint=self.config.endpoint, insecure=True, headers=self.headers
         )
         self.processor = BatchSpanProcessor(self._exporter)
-        self.meter_uploader = self._init_meter_uploader(exporter_id="apmplus")
-
-    def _init_meter_uploader(self, exporter_id: str) -> MeterUploader:
-        global_metrics_provider = metrics_api.get_meter_provider()
-        if hasattr(global_metrics_provider, "_sdk_config"):
-            global_resource = global_metrics_provider._sdk_config.resource
-        else:
-            global_resource = Resource.create()
-        resource = global_resource.merge(Resource.create(self.resource_attributes))
-
-        exporter = OTLPMetricExporter(
-            endpoint=self.config.endpoint, headers=self.headers
-        )
-        metric_reader = PeriodicExportingMetricReader(exporter)
 
-        metrics_api.set_meter_provider(
-            metrics_sdk.MeterProvider(metric_readers=[metric_reader], resource=resource)
+        self.meter_uploader = MeterUploader(
+            name="apmplus_meter",
+            endpoint=self.config.endpoint,
+            headers=self.headers,
+            resource_attributes=self.resource_attributes,
         )
-        return MeterUploader(exporter_id=exporter_id)
 
     @override
     def export(self) -> None:

veadk/tracing/telemetry/metrics/__init__.py

@@ -34,7 +34,9 @@
 
 
 def upload_metrics(
-    invocation_context: InvocationContext, llm_response: LlmResponse
+    invocation_context: InvocationContext,
+    llm_request: LlmRequest,
+    llm_response: LlmResponse,
 ) -> None:
     from veadk.agent import Agent
 
@@ -43,7 +45,7 @@ def upload_metrics(
         for tracer in tracers:
             for exporter in getattr(tracer, "exporters", []):
                 if getattr(exporter, "meter_uploader", None):
-                    exporter.meter_uploader.record(llm_response)
+                    exporter.meter_uploader.record(llm_request, llm_response)
 
 
 def trace_send_data(): ...
@@ -144,4 +146,4 @@ def trace_call_llm(
         response: ExtractorResponse = attr_extractor(params)
         ExtractorResponse.update_span(span, attr_name, response)
 
-    upload_metrics(invocation_context, llm_response)
+    upload_metrics(invocation_context, llm_request, llm_response)