ai-agent-book 精选快照（<2MB 代码与文档，来自 github.com/bojieli/ai-agent-book）

chapter9/gaia-experience/AWorld/aworld/metrics/context_manager.py

@@ -0,0 +1,186 @@
+import time
+import asyncio
+from typing import Callable
+from functools import wraps
+from aworld.metrics.metric import get_metric_provider, MetricType, BaseMetric
+from aworld.metrics.template import MetricTemplate, MetricTemplates
+
+_GLOBAL_METIRCS = {}
+
+
+class MetricContext:
+
+    _initialized = False
+
+    @classmethod
+    def configure(cls,
+                  provider: str,
+                  backend: str,
+                  base_url: str = None,
+                  write_token: str = None,
+                  **kwargs):
+        """
+        Configure the metric provider.
+        Args:
+            provider: The provider of the metric provider.
+            backend: The backend of the metric provider.
+            base_url: The base url of the metric provider.
+            write_token: The write token of the metric provider.
+            export_console: Whether to export the metrics to console.
+            **kwargs: The other parameters of the metric provider.
+        """
+        if cls._initialized:
+            cls.shutdown()
+        if provider == "prometheus":
+            from aworld.metrics.prometheus.prometheus_adapter import configure_prometheus_provider
+            configure_prometheus_provider(
+                backend, base_url, write_token, **kwargs)
+        elif provider == "otlp":
+            from aworld.metrics.opentelemetry.opentelemetry_adapter import configure_otlp_provider
+            configure_otlp_provider(backend, base_url, write_token, **kwargs)
+        cls._initialized = True
+
+    @classmethod
+    def metric_initialized(cls):
+        return cls._initialized
+
+    @staticmethod
+    def get_or_create_metric(template: MetricTemplate):
+        if template.name in _GLOBAL_METIRCS:
+            return _GLOBAL_METIRCS[template.name]
+
+        metric = None
+        if template.type == MetricType.COUNTER:
+            metric = get_metric_provider().create_counter(template.name, template.description, template.unit,
+                                                          template.labels)
+        elif template.type == MetricType.UPDOWNCOUNTER:
+            metric = get_metric_provider().create_un_down_counter(template.name, template.description, template.unit,
+                                                                  template.labels)
+        elif template.type == MetricType.GAUGE:
+            metric = get_metric_provider().create_gauge(template.name, template.description, template.unit,
+                                                        template.labels)
+        elif template.type == MetricType.HISTOGRAM:
+            metric = get_metric_provider().create_histogram(template.name, template.description, template.unit,
+                                                            template.buckets, template.labels)
+
+        _GLOBAL_METIRCS[template.name] = metric
+        return metric
+
+    @classmethod
+    def _validate_type(cls, metric: BaseMetric, type: str):
+        if type != metric._type:
+            raise ValueError(f"metric type {metric._type} is not {type}")
+
+    @classmethod
+    def count(cls, template: MetricTemplate, value: int, labels: dict = None):
+        """
+        Increment a counter metric.
+        """
+        metric = cls.get_or_create_metric(template)
+        cls._validate_type(metric, MetricType.COUNTER)
+        metric.add(value, labels)
+
+    @classmethod
+    def inc(cls, template: MetricTemplate, value: int, labels: dict = None):
+        """
+        Increment a updowncounter metric.
+        """
+        metric = cls.get_or_create_metric(template)
+        cls._validate_type(metric, MetricType.UPDOWNCOUNTER)
+        metric.inc(value, labels)
+
+    @classmethod
+    def dec(cls, template: MetricTemplate, value: int, labels: dict = None):
+        """
+        Decrement a updowncounter metric.
+        """
+        metric = cls.get_or_create_metric(template)
+        cls._validate_type(metric, MetricType.UPDOWNCOUNTER)
+        metric.dec(value, labels)
+
+    @classmethod
+    def gauge_set(cls, template: MetricTemplate, value: int, labels: dict = None):
+        """
+        Set a value to a gauge metric.
+        """
+        metric = cls.get_or_create_metric(template)
+        cls._validate_type(metric, MetricType.GAUGE)
+        metric.set(value, labels)
+
+    @classmethod
+    def histogram_record(cls, template: MetricTemplate, value: int, labels: dict = None):
+        """
+        Set a value to a histogram metric.
+        """
+        metric = cls.get_or_create_metric(template)
+        cls._validate_type(metric, MetricType.HISTOGRAM)
+        metric.record(value, labels)
+
+    @classmethod
+    def shutdown(cls):
+        """
+        Shutdown the metric provider.
+        """
+        provider = get_metric_provider()
+        if provider:
+            provider.shutdown()
+        cls._initialized = False
+
+
+class ApiMetricTracker:
+    """
+    Decorator to track API metrics.
+    """
+
+    def __init__(self, api_name: str = None, func: Callable = None):
+        self.start_time = None
+        self.status = "success"
+        self.func = func
+        self.api_name = api_name
+        if self.api_name is None and self.func is not None:
+            self.api_name = self.func.__name__
+
+    def _new_tracker(self, func: Callable):
+        return self.__class__(func=func)
+
+    def __enter__(self):
+        self.start_time = time.time() * 1000
+
+    def __exit__(self, exc_type, value, traceback):
+        if exc_type is None:
+            self.status = "success"
+        else:
+            self.status = "failure"
+        self._record_metrics(self.api_name, self.start_time, self.status)
+
+    def __call__(self, func: Callable = None) -> Callable:
+        if func is None:
+            return self
+        return self.decorator(func)
+
+    def _record_metrics(self, api_name: str, start_time: float, status: str) -> None:
+        """
+        Record metrics for the API.
+        """
+        elapsed_time = time.time() * 1000 - start_time
+        MetricContext.count(MetricTemplates.REQUEST_COUNT, 1,
+                            labels={"method": api_name, "status": status})
+        MetricContext.histogram_record(MetricTemplates.REQUEST_LATENCY, elapsed_time,
+                                       labels={"method": api_name, "status": status})
+
+    def decorator(self, func):
+        """
+        Decorator to track API metrics.
+        """
+
+        @wraps(func)
+        async def async_wrapper(*args, **kwargs):
+            with self._new_tracker(func):
+                return await func(*args, **kwargs)
+
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            with self._new_tracker(func):
+                return func(*args, **kwargs)
+
+        return async_wrapper if asyncio.iscoroutinefunction(func) else wrapper