arize 8.0.0a14__py3-none-any.whl → 8.0.0a16__py3-none-any.whl

arize/experiments/evaluators/rate_limiters.py

@@ -0,0 +1,277 @@
+import asyncio
+import time
+from functools import wraps
+from math import exp
+from typing import Any, Callable, Coroutine, Optional, Tuple, Type, TypeVar
+
+from typing_extensions import ParamSpec
+
+from .exceptions import ArizeException
+from .utils import printif
+
+ParameterSpec = ParamSpec("ParameterSpec")
+GenericType = TypeVar("GenericType")
+AsyncCallable = Callable[ParameterSpec, Coroutine[Any, Any, GenericType]]
+
+
+class UnavailableTokensError(ArizeException):
+    pass
+
+
+class AdaptiveTokenBucket:
+    """
+    An adaptive rate-limiter that adjusts the rate based on the number of rate limit errors.
+
+    This rate limiter does not need to know the exact rate limit. Instead, it starts with a high
+    rate and reduces it whenever a rate limit error occurs. The rate is increased slowly over time
+    if no further errors occur.
+
+    Args:
+    initial_per_second_request_rate (float): The allowed request rate.
+    maximum_per_second_request_rate (float): The maximum allowed request rate.
+    enforcement_window_minutes (float): The time window over which the rate limit is enforced.
+    rate_reduction_factor (float): Multiplier used to reduce the rate limit after an error.
+    rate_increase_factor (float): Exponential factor increasing the rate limit over time.
+    cooldown_seconds (float): The minimum time before allowing the rate limit to decrease again.
+
+    """
+
+    def __init__(
+        self,
+        initial_per_second_request_rate: float,
+        maximum_per_second_request_rate: Optional[float] = None,
+        minimum_per_second_request_rate: float = 0.1,
+        enforcement_window_minutes: float = 1,
+        rate_reduction_factor: float = 0.5,
+        rate_increase_factor: float = 0.01,
+        cooldown_seconds: float = 5,
+    ):
+        self._initial_rate = initial_per_second_request_rate
+        self.rate_reduction_factor = rate_reduction_factor
+        self.enforcement_window = enforcement_window_minutes * 60
+        self.rate_increase_factor = rate_increase_factor
+        self.rate = initial_per_second_request_rate
+        self.minimum_rate = minimum_per_second_request_rate
+
+        if maximum_per_second_request_rate is None:
+            # if unset, do not allow the maximum rate to exceed 3 consecutive rate reductions
+            # assuming the initial rate is the advertised API rate limit
+
+            maximum_rate_multiple = (1 / rate_reduction_factor) ** 3
+            maximum_per_second_request_rate = (
+                initial_per_second_request_rate * maximum_rate_multiple
+            )
+
+        maximum_per_second_request_rate = float(maximum_per_second_request_rate)
+        assert isinstance(maximum_per_second_request_rate, float)
+        self.maximum_rate = maximum_per_second_request_rate
+
+        self.cooldown = cooldown_seconds
+
+        now = time.time()
+        self.last_rate_update = now
+        self.last_checked = now
+        self.last_error = now - self.cooldown
+        self.tokens = 0.0
+
+    def increase_rate(self) -> None:
+        time_since_last_update = time.time() - self.last_rate_update
+        if time_since_last_update > self.enforcement_window:
+            self.rate = self._initial_rate
+        else:
+            self.rate *= exp(self.rate_increase_factor * time_since_last_update)
+            self.rate = min(self.rate, self.maximum_rate)
+        self.last_rate_update = time.time()
+
+    def on_rate_limit_error(
+        self, request_start_time: float, verbose: bool = False
+    ) -> None:
+        now = time.time()
+        if request_start_time < (self.last_error + self.cooldown):
+            # do not reduce the rate for concurrent requests
+            return
+
+        original_rate = self.rate
+
+        self.rate = original_rate * self.rate_reduction_factor
+        printif(
+            verbose,
+            f"Reducing rate from {original_rate} to {self.rate} after rate limit error",
+        )
+
+        self.rate = max(self.rate, self.minimum_rate)
+
+        # reset request tokens on a rate limit error
+        self.tokens = 0
+        self.last_checked = now
+        self.last_rate_update = now
+        self.last_error = now
+        time.sleep(self.cooldown)  # block for a bit to let the rate limit reset
+
+    def max_tokens(self) -> float:
+        return self.rate * self.enforcement_window
+
+    def available_requests(self) -> float:
+        now = time.time()
+        time_since_last_checked = time.time() - self.last_checked
+        self.tokens = min(
+            self.max_tokens(), self.rate * time_since_last_checked + self.tokens
+        )
+        self.last_checked = now
+        return self.tokens
+
+    def make_request_if_ready(self) -> None:
+        if self.available_requests() <= 1:
+            raise UnavailableTokensError
+        self.tokens -= 1
+
+    def wait_until_ready(
+        self,
+        max_wait_time: float = 300,
+    ) -> None:
+        start = time.time()
+        while (time.time() - start) < max_wait_time:
+            try:
+                self.increase_rate()
+                self.make_request_if_ready()
+                break
+            except UnavailableTokensError:
+                time.sleep(0.1 / self.rate)
+                continue
+
+    async def async_wait_until_ready(
+        self,
+        max_wait_time: float = 10,  # defeat the token bucket rate limiter at low rates (<.1 req/s)
+    ) -> None:
+        start = time.time()
+        while (time.time() - start) < max_wait_time:
+            try:
+                self.increase_rate()
+                self.make_request_if_ready()
+                break
+            except UnavailableTokensError:
+                await asyncio.sleep(0.1 / self.rate)
+                continue
+
+
+class RateLimitError(ArizeException): ...
+
+
+class RateLimiter:
+    def __init__(
+        self,
+        rate_limit_error: Optional[Type[BaseException]] = None,
+        max_rate_limit_retries: int = 3,
+        initial_per_second_request_rate: float = 1.0,
+        maximum_per_second_request_rate: Optional[float] = None,
+        enforcement_window_minutes: float = 1,
+        rate_reduction_factor: float = 0.5,
+        rate_increase_factor: float = 0.01,
+        cooldown_seconds: float = 5,
+        verbose: bool = False,
+    ) -> None:
+        self._rate_limit_error: Tuple[Type[BaseException], ...]
+        self._rate_limit_error = (
+            (rate_limit_error,) if rate_limit_error is not None else tuple()
+        )
+
+        self._max_rate_limit_retries = max_rate_limit_retries
+        self._throttler = AdaptiveTokenBucket(
+            initial_per_second_request_rate=initial_per_second_request_rate,
+            maximum_per_second_request_rate=maximum_per_second_request_rate,
+            enforcement_window_minutes=enforcement_window_minutes,
+            rate_reduction_factor=rate_reduction_factor,
+            rate_increase_factor=rate_increase_factor,
+            cooldown_seconds=cooldown_seconds,
+        )
+        self._rate_limit_handling: Optional[asyncio.Event] = None
+        self._rate_limit_handling_lock: Optional[asyncio.Lock] = None
+        self._current_loop: Optional[asyncio.AbstractEventLoop] = None
+        self._verbose = verbose
+
+    def limit(
+        self, fn: Callable[ParameterSpec, GenericType]
+    ) -> Callable[ParameterSpec, GenericType]:
+        @wraps(fn)
+        def wrapper(*args: Any, **kwargs: Any) -> GenericType:
+            try:
+                self._throttler.wait_until_ready()
+                request_start_time = time.time()
+                return fn(*args, **kwargs)
+            except self._rate_limit_error:
+                self._throttler.on_rate_limit_error(
+                    request_start_time, verbose=self._verbose
+                )
+                for _attempt in range(self._max_rate_limit_retries):
+                    try:
+                        request_start_time = time.time()
+                        self._throttler.wait_until_ready()
+                        return fn(*args, **kwargs)
+                    except self._rate_limit_error:
+                        self._throttler.on_rate_limit_error(
+                            request_start_time, verbose=self._verbose
+                        )
+                        continue
+            raise RateLimitError(
+                f"Exceeded max ({self._max_rate_limit_retries}) retries"
+            )
+
+        return wrapper
+
+    def _initialize_async_primitives(self) -> None:
+        """
+        Lazily initialize async primitives to ensure they are created in the correct event loop.
+        """
+        loop = asyncio.get_running_loop()
+        if loop is not self._current_loop:
+            self._current_loop = loop
+            self._rate_limit_handling = asyncio.Event()
+            self._rate_limit_handling.set()
+            self._rate_limit_handling_lock = asyncio.Lock()
+
+    def alimit(
+        self, fn: AsyncCallable[ParameterSpec, GenericType]
+    ) -> AsyncCallable[ParameterSpec, GenericType]:
+        @wraps(fn)
+        async def wrapper(*args: Any, **kwargs: Any) -> GenericType:
+            self._initialize_async_primitives()
+            assert self._rate_limit_handling_lock is not None and isinstance(
+                self._rate_limit_handling_lock, asyncio.Lock
+            )
+            assert self._rate_limit_handling is not None and isinstance(
+                self._rate_limit_handling, asyncio.Event
+            )
+            try:
+                try:
+                    await asyncio.wait_for(
+                        self._rate_limit_handling.wait(), 120
+                    )
+                except asyncio.TimeoutError:
+                    self._rate_limit_handling.set()  # Set the event as a failsafe
+                await self._throttler.async_wait_until_ready()
+                request_start_time = time.time()
+                return await fn(*args, **kwargs)
+            except self._rate_limit_error:
+                async with self._rate_limit_handling_lock:
+                    self._rate_limit_handling.clear()  # prevent new requests from starting
+                    self._throttler.on_rate_limit_error(
+                        request_start_time, verbose=self._verbose
+                    )
+                    try:
+                        for _attempt in range(self._max_rate_limit_retries):
+                            try:
+                                request_start_time = time.time()
+                                await self._throttler.async_wait_until_ready()
+                                return await fn(*args, **kwargs)
+                            except self._rate_limit_error:
+                                self._throttler.on_rate_limit_error(
+                                    request_start_time, verbose=self._verbose
+                                )
+                                continue
+                    finally:
+                        self._rate_limit_handling.set()  # allow new requests to start
+            raise RateLimitError(
+                f"Exceeded max ({self._max_rate_limit_retries}) retries"
+            )
+
+        return wrapper

arize/experiments/evaluators/types.py

@@ -0,0 +1,122 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any, Dict, List, Mapping, Tuple
+
+JSONSerializable = Dict[str, Any] | List[Any] | str | int | float | bool
+
+
+class AnnotatorKind(Enum):
+    CODE = "CODE"
+    LLM = "LLM"
+
+
+EvaluatorKind = str
+EvaluatorName = str
+
+Score = bool | int | float | None
+Label = str | None
+Explanation = str | None
+
+
+@dataclass(frozen=True)
+class EvaluationResult:
+    """
+    Represents the result of an evaluation.
+    Args:
+        score: The score of the evaluation.
+        label: The label of the evaluation.
+        explanation: The explanation of the evaluation.
+        metadata: Additional metadata for the evaluation.
+    """
+
+    score: float | None = None
+    label: str | None = None
+    explanation: str | None = None
+    metadata: Mapping[str, JSONSerializable] = field(default_factory=dict)
+
+    @classmethod
+    def from_dict(
+        cls, obj: Mapping[str, Any] | None
+    ) -> EvaluationResult | None:
+        if not obj:
+            return None
+        return cls(
+            score=obj.get("score"),
+            label=obj.get("label"),
+            explanation=obj.get("explanation"),
+            metadata=obj.get("metadata") or {},
+        )
+
+    def __post_init__(self) -> None:
+        if self.score is None and not self.label:
+            raise ValueError("Must specify score or label, or both")
+        if self.score is None and not self.label:
+            object.__setattr__(self, "score", 0)
+        for k in ("label", "explanation"):
+            v = getattr(self, k, None)
+            if v is not None:
+                object.__setattr__(self, k, str(v) or None)
+
+
+EvaluatorOutput = (
+    EvaluationResult
+    | bool
+    | int
+    | float
+    | str
+    | Tuple[Score, Label, Explanation]
+)
+
+
+@dataclass
+class EvaluationResultFieldNames:
+    """Column names for mapping evaluation results in a DataFrame.
+
+    Args:
+        score: Optional name of column containing evaluation scores
+        label: Optional name of column containing evaluation labels
+        explanation: Optional name of column containing evaluation explanations
+        metadata: Optional mapping of metadata keys to column names. If a column name
+            is None or empty string, the metadata key will be used as the column name.
+
+    Examples:
+        >>> # Basic usage with score and label columns
+        >>> EvaluationResultColumnNames(
+        ...     score="quality.score", label="quality.label"
+        ... )
+
+        >>> # Using metadata with same key and column name
+        >>> EvaluationResultColumnNames(
+        ...     score="quality.score",
+        ...     metadata={
+        ...         "version": None
+        ...     },  # Will look for column named "version"
+        ... )
+
+        >>> # Using metadata with different key and column name
+        >>> EvaluationResultColumnNames(
+        ...     score="quality.score",
+        ...     metadata={
+        ...         # Will look for "column_in_my_df.version" column and ingest as
+        ...         # "eval.{EvaluatorName}.meatadata.model_version"
+        ...         "model_version": "column_in_my_df.version",
+        ...         # Will look for "column_in_my_df.ts" column and ingest as
+        ...         # "eval.{EvaluatorName}.metadata.timestamp"
+        ...         "timestamp": "column_in_my_df.ts",
+        ...     },
+        ... )
+
+    Raises:
+        ValueError: If neither score nor label column names are specified
+    """
+
+    score: str | None = None
+    label: str | None = None
+    explanation: str | None = None
+    metadata: Dict[str, str | None] | None = None
+
+    def __post_init__(self) -> None:
+        if self.score is None and self.label is None:
+            raise ValueError("Must specify score or label column name, or both")

arize/experiments/evaluators/utils.py

@@ -0,0 +1,198 @@
+import functools
+import inspect
+from typing import TYPE_CHECKING, Any, Callable, Optional
+
+from tqdm.auto import tqdm
+
+from arize.experiments.evaluators.types import (
+    EvaluationResult,
+    JSONSerializable,
+)
+
+
+def get_func_name(fn: Callable[..., Any]) -> str:
+    """
+    Makes a best-effort attempt to get the name of the function.
+    """
+    if isinstance(fn, functools.partial):
+        return fn.func.__qualname__
+    if hasattr(fn, "__qualname__") and not fn.__qualname__.endswith("<lambda>"):
+        return fn.__qualname__.split(".<locals>.")[-1]
+    return str(fn)
+
+
+if TYPE_CHECKING:
+    from ..evaluators.base import Evaluator
+
+
+def unwrap_json(obj: JSONSerializable) -> JSONSerializable:
+    if isinstance(obj, dict) and len(obj) == 1:
+        key = next(iter(obj.keys()))
+        output = obj[key]
+        assert isinstance(
+            output, (dict, list, str, int, float, bool, type(None))
+        ), "Output must be JSON serializable"
+        return output
+    return obj
+
+
+def validate_evaluator_signature(sig: inspect.Signature) -> None:
+    # Check that the wrapped function has a valid signature for use as an evaluator
+    # If it does not, raise an error to exit early before running evaluations
+    params = sig.parameters
+    valid_named_params = {
+        "input",
+        "output",
+        "experiment_output",
+        "dataset_output",
+        "metadata",
+        "dataset_row",
+    }
+    if len(params) == 0:
+        raise ValueError(
+            "Evaluation function must have at least one parameter."
+        )
+    if len(params) > 1:
+        for not_found in set(params) - valid_named_params:
+            param = params[not_found]
+            if (
+                param.kind is inspect.Parameter.VAR_KEYWORD
+                or param.default is not inspect.Parameter.empty
+            ):
+                continue
+            raise ValueError(
+                f"Invalid parameter names in evaluation function: {', '.join(not_found)}. "
+                "Parameters names for multi-argument functions must be "
+                f"any of: {', '.join(valid_named_params)}."
+            )
+
+
+def _bind_evaluator_signature(
+    sig: inspect.Signature, **kwargs: Any
+) -> inspect.BoundArguments:
+    parameter_mapping = {
+        "input": kwargs.get("input"),
+        "output": kwargs.get("output"),
+        "experiment_output": kwargs.get("experiment_output"),
+        "dataset_output": kwargs.get("dataset_output"),
+        "metadata": kwargs.get("metadata"),
+        "dataset_row": kwargs.get("dataset_row"),
+    }
+    params = sig.parameters
+    if len(params) == 1:
+        parameter_name = next(iter(params))
+        if parameter_name in parameter_mapping:
+            return sig.bind(parameter_mapping[parameter_name])
+        else:
+            return sig.bind(parameter_mapping["experiment_output"])
+    return sig.bind_partial(
+        **{
+            name: parameter_mapping[name]
+            for name in set(parameter_mapping).intersection(params)
+        }
+    )
+
+
+def create_evaluator(
+    name: Optional[str] = None,
+    scorer: Optional[Callable[[Any], EvaluationResult]] = None,
+) -> Callable[[Callable[..., Any]], "Evaluator"]:
+    if scorer is None:
+        scorer = _default_eval_scorer
+
+    def wrapper(func: Callable[..., Any]) -> "Evaluator":
+        nonlocal name
+        if not name:
+            name = get_func_name(func)
+        assert name is not None
+
+        wrapped_signature = inspect.signature(func)
+        validate_evaluator_signature(wrapped_signature)
+
+        if inspect.iscoroutinefunction(func):
+            return _wrap_coroutine_evaluation_function(
+                name, wrapped_signature, scorer
+            )(func)
+
+        return _wrap_sync_evaluation_function(name, wrapped_signature, scorer)(
+            func
+        )
+
+    return wrapper
+
+
+def _wrap_coroutine_evaluation_function(
+    name: str,
+    sig: inspect.Signature,
+    convert_to_score: Callable[[Any], EvaluationResult],
+) -> Callable[[Callable[..., Any]], "Evaluator"]:
+    from ..evaluators.base import Evaluator
+
+    def wrapper(func: Callable[..., Any]) -> "Evaluator":
+        class AsyncEvaluator(Evaluator):
+            def __init__(self) -> None:
+                self._name = name
+
+            @functools.wraps(func)
+            async def __call__(self, *args: Any, **kwargs: Any) -> Any:
+                return await func(*args, **kwargs)
+
+            async def async_evaluate(self, **kwargs: Any) -> EvaluationResult:
+                bound_signature = _bind_evaluator_signature(sig, **kwargs)
+                result = await func(
+                    *bound_signature.args, **bound_signature.kwargs
+                )
+                return convert_to_score(result)
+
+        return AsyncEvaluator()
+
+    return wrapper
+
+
+def _wrap_sync_evaluation_function(
+    name: str,
+    sig: inspect.Signature,
+    convert_to_score: Callable[[Any], EvaluationResult],
+) -> Callable[[Callable[..., Any]], "Evaluator"]:
+    from ..evaluators.base import Evaluator
+
+    def wrapper(func: Callable[..., Any]) -> "Evaluator":
+        class SyncEvaluator(Evaluator):
+            def __init__(self) -> None:
+                self._name = name
+
+            @functools.wraps(func)
+            def __call__(self, *args: Any, **kwargs: Any) -> Any:
+                return func(*args, **kwargs)
+
+            def evaluate(self, **kwargs: Any) -> EvaluationResult:
+                bound_signature = _bind_evaluator_signature(sig, **kwargs)
+                result = func(*bound_signature.args, **bound_signature.kwargs)
+                return convert_to_score(result)
+
+        return SyncEvaluator()
+
+    return wrapper
+
+
+def _default_eval_scorer(result: Any) -> EvaluationResult:
+    if isinstance(result, EvaluationResult):
+        return result
+    if isinstance(result, bool):
+        return EvaluationResult(score=float(result), label=str(result))
+    if hasattr(result, "__float__"):
+        return EvaluationResult(score=float(result))
+    if isinstance(result, str):
+        return EvaluationResult(label=result)
+    if isinstance(result, (tuple, list)) and len(result) == 2:
+        # If the result is a 2-tuple, the first item will be recorded as the score
+        # and the second item will recorded as the explanation.
+        return EvaluationResult(
+            score=float(result[0]), explanation=str(result[1])
+        )
+    raise ValueError(f"Unsupported evaluation result type: {type(result)}")
+
+
+def printif(condition: bool, *args: Any, **kwargs: Any) -> None:
+    if condition:
+        tqdm.write(*args, **kwargs)