ragbits-core 0.16.0__py3-none-any.whl → 1.4.0.dev202512021005__py3-none-any.whl

ragbits/core/__init__.py

@@ -1,9 +1,28 @@
 import os
+from concurrent.futures import ThreadPoolExecutor
 
 import typer
 
-from ragbits.core import audit
+from ragbits.core.audit.traces import set_trace_handlers
+
+_config_executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="config-import")
+_config_future = None
+
+
+def _import_and_run_config() -> None:
+    from ragbits.core.config import import_modules_from_config
+
+    import_modules_from_config()
+
+
+def ensure_config_loaded() -> None:
+    """Wait for config import to complete if it hasn't already."""
+    if _config_future:
+        _config_future.result()
+
 
 if os.getenv("RAGBITS_VERBOSE", "0") == "1":
     typer.echo('Verbose mode is enabled with environment variable "RAGBITS_VERBOSE".')
-    audit.set_trace_handlers("cli")
+    set_trace_handlers("cli")
+
+_config_future = _config_executor.submit(_import_and_run_config)

ragbits/core/audit/__init__.py

@@ -1,157 +1,15 @@
-import asyncio
-import inspect
-from collections.abc import Callable, Iterator
-from contextlib import ExitStack, contextmanager
-from functools import wraps
-from types import SimpleNamespace
-from typing import Any, ParamSpec, TypeVar
-
-from ragbits.core.audit.base import TraceHandler
-
-__all__ = ["TraceHandler", "set_trace_handlers", "trace", "traceable"]
-
-_trace_handlers: list[TraceHandler] = []
-
-Handler = str | TraceHandler
-
-P = ParamSpec("P")
-R = TypeVar("R")
-
-
-def set_trace_handlers(handlers: Handler | list[Handler]) -> None:
-    """
-    Setup trace handlers.
-
-    Args:
-        handlers: List of trace handlers to be used.
-
-    Raises:
-        ValueError: If handler is not found.
-        TypeError: If handler type is invalid.
-    """
-    global _trace_handlers  # noqa: PLW0602
-
-    if isinstance(handlers, Handler):
-        handlers = [handlers]
-
-    for handler in handlers:  # type: ignore
-        if isinstance(handler, TraceHandler):
-            _trace_handlers.append(handler)
-        elif isinstance(handler, str):
-            if handler == "otel":
-                from ragbits.core.audit.otel import OtelTraceHandler
-
-                if not any(isinstance(item, OtelTraceHandler) for item in _trace_handlers):
-                    _trace_handlers.append(OtelTraceHandler())
-            elif handler == "cli":
-                from ragbits.core.audit.cli import CLITraceHandler
-
-                if not any(isinstance(item, CLITraceHandler) for item in _trace_handlers):
-                    _trace_handlers.append(CLITraceHandler())
-            else:
-                raise ValueError(f"Handler {handler} not found.")
-        else:
-            raise TypeError(f"Invalid handler type: {type(handler)}")
-
-
-def clear_event_handlers() -> None:
-    """
-    Clear all trace handlers.
-    """
-    global _trace_handlers  # noqa: PLW0602
-
-    _trace_handlers.clear()
-
-
-@contextmanager
-def trace(name: str | None = None, **inputs: Any) -> Iterator[SimpleNamespace]:  # noqa: ANN401
-    """
-    Context manager for processing a trace.
-
-    Args:
-        name: The name of the trace.
-        inputs: The input data.
-
-    Yields:
-        The output data.
-    """
-    # We need to go up 2 frames (trace() and __enter__()) to get the parent function.
-    parent_frame = inspect.stack()[2].frame
-    name = (
-        (
-            f"{cls.__class__.__qualname__}.{parent_frame.f_code.co_name}"
-            if (cls := parent_frame.f_locals.get("self"))
-            else parent_frame.f_code.co_name
-        )
-        if name is None
-        else name
-    )
-
-    with ExitStack() as stack:
-        outputs = [stack.enter_context(handler.trace(name, **inputs)) for handler in _trace_handlers]
-        yield (out := SimpleNamespace())
-        for output in outputs:
-            output.__dict__.update(vars(out))
-
-
-def traceable(func: Callable[P, R]) -> Callable[P, R]:
-    """
-    Decorator for making a function traceable.
-
-    Args:
-        func: The function to be decorated.
-
-    Returns:
-        The decorated function.
-    """
-
-    @wraps(func)
-    def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
-        inputs = _get_function_inputs(func, args, kwargs)
-        with trace(name=func.__qualname__, **inputs) as outputs:
-            returned = func(*args, **kwargs)
-            if returned is not None:
-                outputs.returned = returned
-        return returned
-
-    @wraps(func)
-    async def wrapper_async(*args: P.args, **kwargs: P.kwargs) -> R:
-        inputs = _get_function_inputs(func, args, kwargs)
-        with trace(name=func.__qualname__, **inputs) as outputs:
-            returned = await func(*args, **kwargs)  # type: ignore
-            if returned is not None:
-                outputs.returned = returned
-        return returned
-
-    return wrapper_async if asyncio.iscoroutinefunction(func) else wrapper  # type: ignore
-
-
-def _get_function_inputs(func: Callable, args: tuple, kwargs: dict) -> dict:
-    """
-    Get the dictionary of inputs for a function based on positional and keyword arguments.
-
-    Args:
-        func: The function to get inputs for.
-        args: The positional arguments.
-        kwargs: The keyword arguments.
-
-    Returns:
-        The dictionary of inputs.
-    """
-    sig_params = inspect.signature(func).parameters
-    merged = {}
-    pos_args_used = 0
-
-    for param_name, param in sig_params.items():
-        if param_name in kwargs:
-            merged[param_name] = kwargs[param_name]
-        elif pos_args_used < len(args):
-            if param_name not in ("self", "cls", "args", "kwargs"):
-                merged[param_name] = args[pos_args_used]
-            pos_args_used += 1
-        elif param.default is not param.empty:
-            merged[param_name] = param.default
-
-    merged.update({k: v for k, v in kwargs.items() if k not in merged})
-
-    return merged
+from ragbits.core.audit.metrics import clear_metric_handlers, set_metric_handlers
+from ragbits.core.audit.metrics.base import MetricHandler
+from ragbits.core.audit.traces import clear_trace_handlers, set_trace_handlers, trace, traceable
+from ragbits.core.audit.traces.base import TraceHandler
+
+__all__ = [
+    "MetricHandler",
+    "TraceHandler",
+    "clear_metric_handlers",
+    "clear_trace_handlers",
+    "set_metric_handlers",
+    "set_trace_handlers",
+    "trace",
+    "traceable",
+]

ragbits/core/audit/metrics/__init__.py

@@ -0,0 +1,83 @@
+from enum import Enum
+from typing import Any
+
+from ragbits.core.audit.metrics.base import MetricHandler, MetricType, register_metric
+
+__all__ = [
+    "MetricHandler",
+    "MetricType",
+    "clear_metric_handlers",
+    "record_metric",
+    "register_metric",
+    "set_metric_handlers",
+]
+
+Handler = str | MetricHandler
+
+_metric_handlers: list[MetricHandler] = []
+
+
+def set_metric_handlers(handlers: Handler | list[Handler]) -> None:
+    """
+    Set the global metric handlers.
+
+    Args:
+        handlers: List of metric handlers to be used.
+
+    Raises:
+        ValueError: If handler is not found.
+        TypeError: If handler type is invalid.
+    """
+    global _metric_handlers  # noqa: PLW0602
+
+    if isinstance(handlers, Handler):
+        handlers = [handlers]
+
+    for handler in handlers:
+        if isinstance(handler, MetricHandler):
+            _metric_handlers.append(handler)
+        elif isinstance(handler, str):
+            match handler.lower():
+                case "otel":
+                    from ragbits.core.audit.metrics.otel import OtelMetricHandler
+
+                    if not any(isinstance(item, OtelMetricHandler) for item in _metric_handlers):
+                        _metric_handlers.append(OtelMetricHandler())
+
+                case "logfire":
+                    from ragbits.core.audit.metrics.logfire import LogfireMetricHandler
+
+                    if not any(isinstance(item, LogfireMetricHandler) for item in _metric_handlers):
+                        _metric_handlers.append(LogfireMetricHandler())
+
+                case _:
+                    raise ValueError(f"Not found handler: {handler}")
+        else:
+            raise TypeError(f"Invalid handler type: {type(handler)}")
+
+
+def clear_metric_handlers() -> None:
+    """
+    Clear all metric handlers.
+    """
+    global _metric_handlers  # noqa: PLW0602
+    _metric_handlers.clear()
+
+
+def record_metric(
+    metric: str | Enum,
+    value: int | float,
+    metric_type: MetricType,
+    **attributes: Any,  # noqa: ANN401
+) -> None:
+    """
+    Record a metric of any type using the global metric handlers.
+
+    Args:
+        metric: The metric key (name or enum value) to record
+        value: The value to record
+        metric_type: The type of metric (histogram, counter, gauge)
+        **attributes: Additional metadata for the metric
+    """
+    for handler in _metric_handlers:
+        handler.record_metric(metric_key=metric, value=value, attributes=attributes, metric_type=metric_type)

ragbits/core/audit/metrics/base.py

@@ -0,0 +1,198 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from enum import Enum, auto
+from typing import Any
+
+
+class MetricType(Enum):
+    """Supported metric types."""
+
+    HISTOGRAM = "histogram"
+    COUNTER = "counter"
+    GAUGE = "gauge"
+
+
+@dataclass
+class Metric:
+    """
+    Represents the metric configuration data.
+    """
+
+    name: str
+    description: str
+    unit: str
+    type: MetricType
+
+
+class LLMMetric(Enum):
+    """
+    LLM-related metrics that can be recorded.
+    Each metric has a predefined type and is registered in the global registry.
+    """
+
+    # Histogram metrics
+    PROMPT_THROUGHPUT = auto()
+    TOKEN_THROUGHPUT = auto()
+    INPUT_TOKENS = auto()
+    TIME_TO_FIRST_TOKEN = auto()
+
+
+# Global registry for all metrics by type
+METRICS_REGISTRY: dict[MetricType, dict[Any, Metric]] = {
+    MetricType.HISTOGRAM: {},
+    MetricType.COUNTER: {},
+    MetricType.GAUGE: {},
+}
+
+
+# Register default LLM metrics
+METRICS_REGISTRY[MetricType.HISTOGRAM][LLMMetric.PROMPT_THROUGHPUT] = Metric(
+    name="prompt_throughput",
+    description="Tracks the response time of LLM calls in seconds",
+    unit="s",
+    type=MetricType.HISTOGRAM,
+)
+METRICS_REGISTRY[MetricType.HISTOGRAM][LLMMetric.TOKEN_THROUGHPUT] = Metric(
+    name="token_throughput",
+    description="Tracks tokens generated per second",
+    unit="tokens/s",
+    type=MetricType.HISTOGRAM,
+)
+METRICS_REGISTRY[MetricType.HISTOGRAM][LLMMetric.INPUT_TOKENS] = Metric(
+    name="input_tokens",
+    description="Tracks the number of input tokens per request",
+    unit="tokens",
+    type=MetricType.HISTOGRAM,
+)
+METRICS_REGISTRY[MetricType.HISTOGRAM][LLMMetric.TIME_TO_FIRST_TOKEN] = Metric(
+    name="time_to_first_token",
+    description="Tracks the time to first token in seconds",
+    unit="s",
+    type=MetricType.HISTOGRAM,
+)
+
+
+def register_metric(key: str | Enum, metric: Metric) -> None:
+    """
+    Register a new metric in the global registry by type.
+
+    Args:
+        key: The metric key (enum value or string)
+        metric: The metric configuration
+    """
+    METRICS_REGISTRY[metric.type][key] = metric
+
+
+def get_metric(key: str | Enum, metric_type: MetricType) -> Metric | None:
+    """
+    Get a metric from the registry by key and type.
+
+    Args:
+        key: The metric key (enum value or string)
+        metric_type: The type of metric to retrieve
+
+    Returns:
+        The metric configuration if found, None otherwise
+    """
+    return METRICS_REGISTRY[metric_type].get(key)
+
+
+class MetricHandler(ABC):
+    """
+    Base class for all metric handlers.
+    """
+
+    def __init__(self, metric_prefix: str = "ragbits") -> None:
+        """
+        Initialize the MetricHandler instance.
+
+        Args:
+            metric_prefix: Prefix for all metric names.
+        """
+        super().__init__()
+        self._metric_prefix = metric_prefix
+        self._metrics: dict[str, Any] = {}
+
+    @abstractmethod
+    def create_metric(
+        self, name: str, unit: str = "", description: str = "", metric_type: MetricType = MetricType.HISTOGRAM
+    ) -> Any:  # noqa: ANN401
+        """
+        Create a metric of the given type.
+
+        Args:
+            name: The metric name.
+            unit: The metric unit.
+            description: The metric description.
+            metric_type: The type of the metric (histogram, counter, gauge).
+
+        Returns:
+            The initialized metric.
+        """
+
+    @abstractmethod
+    def _record(self, metric: Any, value: int | float, attributes: dict | None = None) -> None:  # noqa: ANN401
+        """
+        Low-level method to record a value for a specified metric.
+        This method should not be called directly, use record_metric instead.
+
+        Args:
+            metric: The metric to record.
+            value: The value to record for the metric.
+            attributes: Additional metadata for the metric.
+        """
+
+    def register_metric_instance(
+        self, name: str, unit: str = "", description: str = "", metric_type: MetricType = MetricType.HISTOGRAM
+    ) -> None:
+        """
+        Register a metric instance.
+
+        Args:
+            name: The metric name.
+            unit: The metric unit.
+            description: The metric description.
+            metric_type: The type of the metric (histogram, counter, gauge).
+        """
+        self._metrics[name] = self.create_metric(
+            name=f"{self._metric_prefix}_{name}",
+            unit=unit,
+            description=description,
+            metric_type=metric_type,
+        )
+
+    def record_metric(
+        self,
+        metric_key: str | Enum,
+        value: int | float,
+        attributes: dict | None = None,
+        metric_type: MetricType = MetricType.HISTOGRAM,
+    ) -> None:
+        """
+        Record the value for a specified metric.
+
+        Args:
+            metric_key: The metric key (name or enum value) to record.
+            value: The value to record for the metric.
+            attributes: Additional metadata for the metric.
+            metric_type: The type of the metric (histogram, counter, gauge).
+        """
+        metric_cfg = get_metric(metric_key, metric_type)
+        if metric_cfg:
+            metric_name = metric_cfg.name
+            if metric_name not in self._metrics:
+                self.register_metric_instance(
+                    name=metric_name,
+                    unit=metric_cfg.unit,
+                    description=metric_cfg.description,
+                    metric_type=metric_type,
+                )
+        else:
+            metric_name = str(metric_key)
+            if metric_name not in self._metrics:
+                self.register_metric_instance(metric_name, metric_type=metric_type)
+        self._record(
+            metric=self._metrics[metric_name],
+            value=value,
+            attributes=attributes,
+        )

ragbits/core/audit/metrics/logfire.py

@@ -0,0 +1,19 @@
+from typing import Any
+
+import logfire
+
+from ragbits.core.audit.metrics.otel import OtelMetricHandler
+
+
+class LogfireMetricHandler(OtelMetricHandler):
+    """
+    Logfire metric handler.
+    """
+
+    def __init__(self, metric_prefix: str = "ragbits", *args: Any, **kwargs: Any) -> None:  # noqa: ANN401
+        """
+        Initialize the LogfireMetricHandler instance.
+        """
+        logfire.configure(*args, **kwargs)
+        logfire.instrument_system_metrics()
+        super().__init__(metric_prefix=metric_prefix)

ragbits/core/audit/metrics/otel.py

@@ -0,0 +1,65 @@
+from typing import Any
+
+from opentelemetry.metrics import MeterProvider, get_meter
+
+from ragbits.core.audit.metrics.base import MetricHandler, MetricType
+
+
+class OtelMetricHandler(MetricHandler):
+    """
+    OpenTelemetry metric handler.
+    """
+
+    def __init__(self, provider: MeterProvider | None = None, metric_prefix: str = "ragbits") -> None:
+        """
+        Initialize the OtelMetricHandler instance.
+
+        Args:
+            provider: The meter provider to use.
+            metric_prefix: Prefix for all metric names.
+        """
+        super().__init__(metric_prefix=metric_prefix)
+        self._meter = get_meter(name=__name__, meter_provider=provider)
+
+    def create_metric(
+        self, name: str, unit: str = "", description: str = "", metric_type: MetricType = MetricType.HISTOGRAM
+    ) -> Any:  # noqa: ANN401
+        """
+        Create a metric of the specified type.
+
+        Args:
+            name: The metric name.
+            unit: The metric unit.
+            description: The metric description.
+            metric_type: The type of the metric (histogram, counter, gauge).
+
+        Returns:
+            The initialized metric.
+        """
+        if metric_type == MetricType.HISTOGRAM:
+            return self._meter.create_histogram(name=name, unit=unit, description=description)
+        elif metric_type == MetricType.COUNTER:
+            return self._meter.create_counter(name=name, unit=unit, description=description)
+        elif metric_type == MetricType.GAUGE:
+            return self._meter.create_gauge(name=name, unit=unit, description=description)
+        else:
+            raise ValueError(f"Unsupported metric type: {metric_type}")
+
+    def _record(self, metric: Any, value: int | float, attributes: dict | None = None) -> None:  # noqa
+        """
+        Record the value for a specified metric.
+
+        Args:
+            metric: The metric to record (histogram, counter, or gauge).
+            value: The value to record for the metric.
+            attributes: Additional metadata for the metric.
+        """
+        # Determine metric type by instance
+        if hasattr(metric, "record"):
+            # Histogram or Gauge (OpenTelemetry Python API)
+            metric.record(value, attributes=attributes)
+        elif hasattr(metric, "add"):
+            # Counter
+            metric.add(value, attributes=attributes)
+        else:
+            raise TypeError("Unsupported metric instance for recording")