fleet-memory/hindsight-api/hindsight_api/metrics.py

"""
OpenTelemetry metrics instrumentation for Hindsight API.

This module provides metrics for:
- Operation latency (retain, recall, reflect) with percentiles
- Token usage (input/output) per operation
- Per-bank granularity via labels
"""
import logging
from typing import Dict, Any, Optional
from contextlib import contextmanager
import time

from opentelemetry import metrics
from opentelemetry.sdk.metrics import MeterProvider
from opentelemetry.sdk.resources import Resource
from opentelemetry.exporter.prometheus import PrometheusMetricReader
from prometheus_client import REGISTRY

logger = logging.getLogger(__name__)

# Global meter instance
_meter = None


def initialize_metrics(service_name: str = "hindsight-api", service_version: str = "1.0.0"):
    """
    Initialize OpenTelemetry metrics with Prometheus exporter.

    This should be called once during application startup.

    Args:
        service_name: Name of the service for resource attributes
        service_version: Version of the service

    Returns:
        PrometheusMetricReader instance (for accessing metrics endpoint)
    """
    global _meter

    # Create resource with service information
    resource = Resource.create({
        "service.name": service_name,
        "service.version": service_version,
    })

    # Create Prometheus metric reader
    prometheus_reader = PrometheusMetricReader()

    # Create meter provider with Prometheus exporter
    provider = MeterProvider(
        resource=resource,
        metric_readers=[prometheus_reader]
    )

    # Set the global meter provider
    metrics.set_meter_provider(provider)

    # Get meter for this application
    _meter = metrics.get_meter(__name__)

    return prometheus_reader


def get_meter():
    """Get the global meter instance."""
    if _meter is None:
        raise RuntimeError("Metrics not initialized. Call initialize_metrics() first.")
    return _meter


class MetricsCollectorBase:
    """Base class for metrics collectors."""

    @contextmanager
    def record_operation(self, operation: str, bank_id: str, budget: Optional[str] = None, max_tokens: Optional[int] = None):
        """Context manager to record operation duration and status."""
        raise NotImplementedError

    def record_tokens(self, operation: str, bank_id: str, input_tokens: int = 0, output_tokens: int = 0, budget: Optional[str] = None, max_tokens: Optional[int] = None):
        """Record token usage for an operation."""
        raise NotImplementedError


class NoOpMetricsCollector(MetricsCollectorBase):
    """No-op metrics collector that does nothing. Used when metrics are disabled."""

    @contextmanager
    def record_operation(self, operation: str, bank_id: str, budget: Optional[str] = None, max_tokens: Optional[int] = None):
        """No-op context manager."""
        yield

    def record_tokens(self, operation: str, bank_id: str, input_tokens: int = 0, output_tokens: int = 0, budget: Optional[str] = None, max_tokens: Optional[int] = None):
        """No-op token recording."""
        pass


class MetricsCollector(MetricsCollectorBase):
    """
    Collector for Hindsight API metrics.

    Provides methods to record latency and token usage for operations.
    """

    def __init__(self):
        self.meter = get_meter()

        # Operation latency histogram (in seconds)
        # Records duration of retain, recall, reflect operations
        self.operation_duration = self.meter.create_histogram(
            name="hindsight.operation.duration",
            description="Duration of Hindsight operations in seconds",
            unit="s"
        )

        # Token usage counters
        self.tokens_input = self.meter.create_counter(
            name="hindsight.tokens.input",
            description="Number of input tokens consumed",
            unit="tokens"
        )

        self.tokens_output = self.meter.create_counter(
            name="hindsight.tokens.output",
            description="Number of output tokens generated",
            unit="tokens"
        )

        # Operation counter (success/failure)
        self.operation_total = self.meter.create_counter(
            name="hindsight.operation.total",
            description="Total number of operations executed",
            unit="operations"
        )

    @contextmanager
    def record_operation(self, operation: str, bank_id: str, budget: Optional[str] = None, max_tokens: Optional[int] = None):
        """
        Context manager to record operation duration and status.

        Usage:
            with metrics.record_operation("recall", bank_id="user123", budget="mid", max_tokens=4096):
                # ... perform operation
                pass

        Args:
            operation: Operation name (retain, recall, reflect)
            bank_id: Memory bank ID
            budget: Optional budget level (low, mid, high)
            max_tokens: Optional max tokens for the operation
        """
        start_time = time.time()
        attributes = {
            "operation": operation,
            "bank_id": bank_id,
        }
        if budget:
            attributes["budget"] = budget
        if max_tokens:
            attributes["max_tokens"] = str(max_tokens)

        success = True
        try:
            yield
        except Exception:
            success = False
            raise
        finally:
            duration = time.time() - start_time
            attributes["success"] = str(success).lower()

            # Record duration
            self.operation_duration.record(duration, attributes)

            # Record operation count
            self.operation_total.add(1, attributes)

    def record_tokens(self, operation: str, bank_id: str, input_tokens: int = 0, output_tokens: int = 0, budget: Optional[str] = None, max_tokens: Optional[int] = None):
        """
        Record token usage for an operation.

        Args:
            operation: Operation name (retain, recall, reflect)
            bank_id: Memory bank ID
            input_tokens: Number of input tokens
            output_tokens: Number of output tokens
            budget: Optional budget level
            max_tokens: Optional max tokens for the operation
        """
        attributes = {
            "operation": operation,
            "bank_id": bank_id,
        }
        if budget:
            attributes["budget"] = budget
        if max_tokens:
            attributes["max_tokens"] = str(max_tokens)

        if input_tokens > 0:
            self.tokens_input.add(input_tokens, attributes)

        if output_tokens > 0:
            self.tokens_output.add(output_tokens, attributes)


# Global metrics collector instance (defaults to no-op)
_metrics_collector: MetricsCollectorBase = NoOpMetricsCollector()


def get_metrics_collector() -> MetricsCollectorBase:
    """
    Get the global metrics collector instance.

    Returns a no-op collector if metrics are not initialized.
    """
    return _metrics_collector


def create_metrics_collector() -> MetricsCollector:
    """
    Create and set the global metrics collector.

    Should be called after initialize_metrics().
    """
    global _metrics_collector
    _metrics_collector = MetricsCollector()
    return _metrics_collector