fix: fix apmplus metrics

Author: zakahan

veadk/tracing/telemetry/exporters/apmplus_exporter.py

@@ -14,10 +14,10 @@
 
 from typing import Any
 
-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.sdk.metrics import MeterProvider
+from opentelemetry.sdk import metrics as metrics_sdk
 from opentelemetry.sdk.metrics.export import PeriodicExportingMetricReader
 from opentelemetry.sdk.resources import Resource
 from opentelemetry.sdk.trace.export import BatchSpanProcessor
@@ -26,6 +26,9 @@
 
 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.tracing.telemetry.metrics.opentelemetry_metrics import meter_uploader_manager
 from veadk.utils.logger import get_logger
 
 logger = get_logger(__name__)
@@ -68,15 +71,32 @@ 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:
         # init meter
-        resource = Resource.create()
-        exporter = OTLPMetricExporter(endpoint=self.config.endpoint, headers=headers)
+        exporter = OTLPMetricExporter(
+            endpoint=self.config.endpoint, headers=self.headers
+        )
         metric_reader = PeriodicExportingMetricReader(exporter)
-        provider = MeterProvider(metric_readers=[metric_reader], resource=resource)
-        metrics.set_meter_provider(provider)
 
-        # metrics.get_meter("veadk.apmplus.meter")
+        global_metrics_provider = metrics_api.get_meter_provider()
+
+        if getattr(global_metrics_provider, "_sdk_config", None):
+            global_resource = getattr(global_metrics_provider, "_sdk_config").resource
+        else:
+            global_resource = Resource.create()
+
+        new_resource = Resource.create(self.resource_attributes)
+        merged_resource = global_resource.merge(new_resource)
+
+        provider = metrics_sdk.MeterProvider(
+            metric_readers=[metric_reader], resource=merged_resource
+        )
+        metrics_api.set_meter_provider(provider)
+
+        meter_uploader = MeterUploader(exporter_id=exporter_id)
+        return meter_uploader
 
     @override
     def export(self) -> None:

veadk/tracing/telemetry/metrics/opentelemetry_metrics.py

@@ -11,33 +11,20 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
-import time
-
+from google.adk.models import LlmResponse
+from opentelemetry import metrics
 from opentelemetry.metrics._internal import Meter
-from opentelemetry.sdk.metrics import MeterProvider
-from opentelemetry.sdk.metrics.export import PeriodicExportingMetricReader
 
 from veadk.config import getenv
 
-
-class MeterContext:
-    def __init__(
-        self,
-        meter: Meter,
-        provider: MeterProvider,
-        reader: PeriodicExportingMetricReader,
-    ):
-        self.meter = meter
-        self.provider = provider
-        self.reader = reader
+METER_NAME_TEMPLATE = "veadk.{exporter_id}.meter"
 
 
 class MeterUploader:
-    def __init__(self, meter_context: MeterContext):
-        self.meter = meter_context.meter
-        self.provider = meter_context.provider
-        self.reader = meter_context.reader
+    def __init__(self, exporter_id: str):
+        self.meter: Meter = metrics.get_meter(
+            METER_NAME_TEMPLATE.format(exporter_id=exporter_id)
+        )
 
         self.base_attributes = {
             "gen_ai_system": "volcengine",
@@ -57,17 +44,11 @@ def __init__(self, meter_context: MeterContext):
             unit="count",
         )
 
-    def record(self, prompt_tokens: list[int], completion_tokens: list[int]):
-        self.llm_invoke_counter.add(len(completion_tokens), self.base_attributes)
-
-        for prompt_token in prompt_tokens:
-            token_attributes = {**self.base_attributes, "gen_ai_token_type": "input"}
-            self.token_usage.record(prompt_token, attributes=token_attributes)
-        for completion_token in completion_tokens:
-            token_attributes = {**self.base_attributes, "gen_ai_token_type": "output"}
-            self.token_usage.record(completion_token, attributes=token_attributes)
-
-    def close(self):
-        time.sleep(0.05)
-        self.reader.force_flush()
-        self.provider.shutdown()
+    def record(self, llm_response: LlmResponse):
+        input_token = llm_response.usage_metadata.prompt_token_count
+        output_token = llm_response.usage_metadata.candidates_token_count
+        self.llm_invoke_counter.add(1, self.base_attributes)
+        token_attributes = {**self.base_attributes, "gen_ai_token_type": "input"}
+        self.token_usage.record(input_token, attributes=token_attributes)
+        token_attributes = {**self.base_attributes, "gen_ai_token_type": "output"}
+        self.token_usage.record(output_token, attributes=token_attributes)

veadk/tracing/telemetry/telemetry.py

@@ -19,7 +19,12 @@
 logger = get_logger(__name__)
 
 
-def trace_send_data(): ...
+def trace_send_data(invocation_context: InvocationContext, llm_response: LlmResponse):
+    tracers = invocation_context.agent.tracers
+    for tracer in tracers:
+        for exporter in getattr(tracer, "exporters", []):
+            if getattr(exporter, "meter_uploader", None):
+                exporter.meter_uploader.record(llm_response)
 
 
 def set_common_attributes(
@@ -116,3 +121,6 @@ def trace_call_llm(
     for attr_name, attr_extractor in llm_attributes_mapping.items():
         response: ExtractorResponse = attr_extractor(params)
         ExtractorResponse.update_span(span, attr_name, response)
+
+    # Report meter
+    trace_send_data(invocation_context, llm_response)