arize 8.0.0a21__py3-none-any.whl → 8.0.0a23__py3-none-any.whl

arize/experiments/evaluators/rate_limiters.py

@@ -1,8 +1,11 @@
+"""Rate limiting utilities for evaluator execution."""
+
 import asyncio
 import time
+from collections.abc import Callable, Coroutine
 from functools import wraps
 from math import exp
-from typing import Any, Callable, Coroutine, Optional, Tuple, Type, TypeVar
+from typing import Any, TypeVar
 
 from typing_extensions import ParamSpec
 
@@ -15,12 +18,11 @@ AsyncCallable = Callable[ParameterSpec, Coroutine[Any, Any, GenericType]]
 
 
 class UnavailableTokensError(ArizeException):
-    pass
+    """Raised when insufficient tokens are available for rate limiting."""
 
 
 class AdaptiveTokenBucket:
-    """
-    An adaptive rate-limiter that adjusts the rate based on the number of rate limit errors.
+    """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
@@ -39,13 +41,24 @@ class AdaptiveTokenBucket:
     def __init__(
         self,
         initial_per_second_request_rate: float,
-        maximum_per_second_request_rate: Optional[float] = None,
+        maximum_per_second_request_rate: float | None = 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,
-    ):
+    ) -> None:
+        """Initialize the adaptive rate limit state.
+
+        Args:
+            initial_per_second_request_rate: Starting request rate per second.
+            maximum_per_second_request_rate: Maximum allowed rate limit.
+            minimum_per_second_request_rate: Minimum allowed rate limit.
+            enforcement_window_minutes: Time window for rate enforcement.
+            rate_reduction_factor: Factor to reduce rate on errors.
+            rate_increase_factor: Factor to gradually increase rate.
+            cooldown_seconds: Cooldown period before rate adjustments.
+        """
         self._initial_rate = initial_per_second_request_rate
         self.rate_reduction_factor = rate_reduction_factor
         self.enforcement_window = enforcement_window_minutes * 60
@@ -63,7 +76,6 @@ class AdaptiveTokenBucket:
             )
 
         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
@@ -75,6 +87,7 @@ class AdaptiveTokenBucket:
         self.tokens = 0.0
 
     def increase_rate(self) -> None:
+        """Increase the rate limit based on time elapsed since last update."""
         time_since_last_update = time.time() - self.last_rate_update
         if time_since_last_update > self.enforcement_window:
             self.rate = self._initial_rate
@@ -86,6 +99,7 @@ class AdaptiveTokenBucket:
     def on_rate_limit_error(
         self, request_start_time: float, verbose: bool = False
     ) -> None:
+        """Handle rate limit error by reducing the rate and adding cooldown period."""
         now = time.time()
         if request_start_time < (self.last_error + self.cooldown):
             # do not reduce the rate for concurrent requests
@@ -109,9 +123,11 @@ class AdaptiveTokenBucket:
         time.sleep(self.cooldown)  # block for a bit to let the rate limit reset
 
     def max_tokens(self) -> float:
+        """Return the maximum number of tokens allowed in the enforcement window."""
         return self.rate * self.enforcement_window
 
     def available_requests(self) -> float:
+        """Return the current number of available request tokens."""
         now = time.time()
         time_since_last_checked = time.time() - self.last_checked
         self.tokens = min(
@@ -121,6 +137,7 @@ class AdaptiveTokenBucket:
         return self.tokens
 
     def make_request_if_ready(self) -> None:
+        """Make a request if tokens are available, otherwise raise error."""
         if self.available_requests() <= 1:
             raise UnavailableTokensError
         self.tokens -= 1
@@ -129,6 +146,7 @@ class AdaptiveTokenBucket:
         self,
         max_wait_time: float = 300,
     ) -> None:
+        """Wait until tokens are available for making a request."""
         start = time.time()
         while (time.time() - start) < max_wait_time:
             try:
@@ -143,6 +161,7 @@ class AdaptiveTokenBucket:
         self,
         max_wait_time: float = 10,  # defeat the token bucket rate limiter at low rates (<.1 req/s)
     ) -> None:
+        """Asynchronously wait until tokens are available for making a request."""
         start = time.time()
         while (time.time() - start) < max_wait_time:
             try:
@@ -154,25 +173,41 @@ class AdaptiveTokenBucket:
                 continue
 
 
-class RateLimitError(ArizeException): ...
+class RateLimitError(ArizeException):
+    """Raised when a rate limit is exceeded."""
 
 
 class RateLimiter:
+    """Rate limiter for controlling request frequency with adaptive token bucket algorithm."""
+
     def __init__(
         self,
-        rate_limit_error: Optional[Type[BaseException]] = None,
+        rate_limit_error: type[BaseException] | None = None,
         max_rate_limit_retries: int = 3,
         initial_per_second_request_rate: float = 1.0,
-        maximum_per_second_request_rate: Optional[float] = None,
+        maximum_per_second_request_rate: float | None = 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], ...]
+        """Initialize the rate limiter with adaptive token bucket algorithm.
+
+        Args:
+            rate_limit_error: Exception type to catch for rate limiting.
+            max_rate_limit_retries: Maximum retries for rate limit errors.
+            initial_per_second_request_rate: Initial request rate per second.
+            maximum_per_second_request_rate: Maximum allowed rate limit.
+            enforcement_window_minutes: Time window for rate enforcement.
+            rate_reduction_factor: Factor to reduce rate on errors.
+            rate_increase_factor: Factor to gradually increase rate.
+            cooldown_seconds: Cooldown period before rate adjustments.
+            verbose: Whether to print rate limit adjustments.
+        """
+        self._rate_limit_error: tuple[type[BaseException], ...]
         self._rate_limit_error = (
-            (rate_limit_error,) if rate_limit_error is not None else tuple()
+            (rate_limit_error,) if rate_limit_error is not None else ()
         )
 
         self._max_rate_limit_retries = max_rate_limit_retries
@@ -184,44 +219,50 @@ class RateLimiter:
             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._rate_limit_handling: asyncio.Event | None = None
+        self._rate_limit_handling_lock: asyncio.Lock | None = None
+        self._current_loop: asyncio.AbstractEventLoop | None = None
         self._verbose = verbose
 
+    def _retry_with_rate_limit_sync(
+        self,
+        fn: Callable[..., GenericType],
+        remaining_attempts: int,
+        *args: object,
+        **kwargs: object,
+    ) -> GenericType:
+        """Recursively retry a function call with rate limiting."""
+        try:
+            request_start_time = time.time()
+            self._throttler.wait_until_ready()
+            return fn(*args, **kwargs)
+        except self._rate_limit_error as e:
+            self._throttler.on_rate_limit_error(
+                request_start_time, verbose=self._verbose
+            )
+            if remaining_attempts <= 1:
+                raise RateLimitError(
+                    f"Exceeded max ({self._max_rate_limit_retries}) retries"
+                ) from e
+            return self._retry_with_rate_limit_sync(
+                fn, remaining_attempts - 1, *args, **kwargs
+            )
+
     def limit(
         self, fn: Callable[ParameterSpec, GenericType]
     ) -> Callable[ParameterSpec, GenericType]:
+        """Apply rate limiting to a synchronous function."""
+
         @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"
+        def wrapper(*args: object, **kwargs: object) -> GenericType:
+            return self._retry_with_rate_limit_sync(
+                fn, self._max_rate_limit_retries, *args, **kwargs
             )
 
         return wrapper
 
     def _initialize_async_primitives(self) -> None:
-        """
-        Lazily initialize async primitives to ensure they are created in the correct event loop.
-        """
+        """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
@@ -232,15 +273,19 @@ class RateLimiter:
     def alimit(
         self, fn: AsyncCallable[ParameterSpec, GenericType]
     ) -> AsyncCallable[ParameterSpec, GenericType]:
+        """Apply rate limiting to an asynchronous function."""
+
         @wraps(fn)
-        async def wrapper(*args: Any, **kwargs: Any) -> GenericType:
+        async def wrapper(*args: object, **kwargs: object) -> GenericType:
             self._initialize_async_primitives()
-            assert self._rate_limit_handling_lock is not None and isinstance(
+            if self._rate_limit_handling_lock is None or not isinstance(
                 self._rate_limit_handling_lock, asyncio.Lock
-            )
-            assert self._rate_limit_handling is not None and isinstance(
+            ):
+                raise RuntimeError("Rate limit lock not properly initialized")
+            if self._rate_limit_handling is None or not isinstance(
                 self._rate_limit_handling, asyncio.Event
-            )
+            ):
+                raise RuntimeError("Rate limit event not properly initialized")
             try:
                 try:
                     await asyncio.wait_for(
@@ -257,21 +302,25 @@ class RateLimiter:
                     self._throttler.on_rate_limit_error(
                         request_start_time, verbose=self._verbose
                     )
+
+                    async def _retry_async(remaining: int) -> GenericType:
+                        try:
+                            request_start_time = time.time()
+                            await self._throttler.async_wait_until_ready()
+                            return await fn(*args, **kwargs)
+                        except self._rate_limit_error as e:
+                            self._throttler.on_rate_limit_error(
+                                request_start_time, verbose=self._verbose
+                            )
+                            if remaining <= 1:
+                                raise RateLimitError(
+                                    f"Exceeded max ({self._max_rate_limit_retries}) retries"
+                                ) from e
+                            return await _retry_async(remaining - 1)
+
                     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
+                        return await _retry_async(self._max_rate_limit_retries)
                     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

@@ -1,13 +1,29 @@
+"""Type definitions for evaluators and evaluation results."""
+
 from __future__ import annotations
 
 from dataclasses import dataclass, field
 from enum import Enum
-from typing import Any, Dict, List, Mapping, Tuple
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from collections.abc import Mapping
 
-JSONSerializable = Dict[str, Any] | List[Any] | str | int | float | bool
+# Recursive type alias for JSON-serializable values
+JSONSerializable = (
+    dict[str, "JSONSerializable"]
+    | list["JSONSerializable"]
+    | str
+    | int
+    | float
+    | bool
+    | None
+)
 
 
 class AnnotatorKind(Enum):
+    """Enum representing the type of annotator used for evaluation."""
+
     CODE = "CODE"
     LLM = "LLM"
 
@@ -22,8 +38,8 @@ Explanation = str | None
 
 @dataclass(frozen=True)
 class EvaluationResult:
-    """
-    Represents the result of an evaluation.
+    """Represents the result of an evaluation.
+
     Args:
         score: The score of the evaluation.
         label: The label of the evaluation.
@@ -38,8 +54,9 @@ class EvaluationResult:
 
     @classmethod
     def from_dict(
-        cls, obj: Mapping[str, Any] | None
+        cls, obj: Mapping[str, object] | None
     ) -> EvaluationResult | None:
+        """Create an EvaluationResult instance from a dictionary."""
         if not obj:
             return None
         return cls(
@@ -50,6 +67,11 @@ class EvaluationResult:
         )
 
     def __post_init__(self) -> None:
+        """Validate and normalize evaluation result fields.
+
+        Raises:
+            ValueError: If neither score nor label is specified.
+        """
         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:
@@ -66,7 +88,7 @@ EvaluatorOutput = (
     | int
     | float
     | str
-    | Tuple[Score, Label, Explanation]
+    | tuple[Score, Label, Explanation]
 )
 
 
@@ -115,8 +137,13 @@ class EvaluationResultFieldNames:
     score: str | None = None
     label: str | None = None
     explanation: str | None = None
-    metadata: Dict[str, str | None] | None = None
+    metadata: dict[str, str | None] | None = None
 
     def __post_init__(self) -> None:
+        """Validate that at least one output column is specified.
+
+        Raises:
+            ValueError: If neither score nor label column name is specified.
+        """
         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

@@ -1,6 +1,9 @@
+"""Utility functions for evaluator operations."""
+
 import functools
 import inspect
-from typing import TYPE_CHECKING, Any, Callable, Optional
+from collections.abc import Callable
+from typing import TYPE_CHECKING
 
 from tqdm.auto import tqdm
 
@@ -10,10 +13,8 @@ from arize.experiments.evaluators.types import (
 )
 
 
-def get_func_name(fn: Callable[..., Any]) -> str:
-    """
-    Makes a best-effort attempt to get the name of the function.
-    """
+def get_func_name(fn: Callable[..., object]) -> 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>"):
@@ -26,17 +27,36 @@ if TYPE_CHECKING:
 
 
 def unwrap_json(obj: JSONSerializable) -> JSONSerializable:
+    """Unwrap a single-key JSON object to extract its value.
+
+    Args:
+        obj: A JSON-serializable object to unwrap.
+
+    Returns:
+        The unwrapped value if obj is a single-key dict, otherwise the original obj.
+    """
     if isinstance(obj, dict) and len(obj) == 1:
         key = next(iter(obj.keys()))
         output = obj[key]
-        assert isinstance(
+        if not isinstance(
             output, (dict, list, str, int, float, bool, type(None))
-        ), "Output must be JSON serializable"
+        ):
+            raise TypeError(
+                f"Evaluator output must be JSON serializable, got {type(output).__name__}"
+            )
         return output
     return obj
 
 
 def validate_evaluator_signature(sig: inspect.Signature) -> None:
+    """Validate that a function signature is compatible for use as an evaluator.
+
+    Args:
+        sig: The function signature to validate.
+
+    Raises:
+        ValueError: If the signature is invalid for use as an evaluator.
+    """
     # 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
@@ -68,7 +88,7 @@ def validate_evaluator_signature(sig: inspect.Signature) -> None:
 
 
 def _bind_evaluator_signature(
-    sig: inspect.Signature, **kwargs: Any
+    sig: inspect.Signature, **kwargs: object
 ) -> inspect.BoundArguments:
     parameter_mapping = {
         "input": kwargs.get("input"),
@@ -83,8 +103,7 @@ def _bind_evaluator_signature(
         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(parameter_mapping["experiment_output"])
     return sig.bind_partial(
         **{
             name: parameter_mapping[name]
@@ -94,17 +113,27 @@ def _bind_evaluator_signature(
 
 
 def create_evaluator(
-    name: Optional[str] = None,
-    scorer: Optional[Callable[[Any], EvaluationResult]] = None,
-) -> Callable[[Callable[..., Any]], "Evaluator"]:
+    name: str | None = None,
+    scorer: Callable[[object], EvaluationResult] | None = None,
+) -> Callable[[Callable[..., object]], "Evaluator"]:
+    """Create an evaluator decorator for wrapping evaluation functions.
+
+    Args:
+        name: Optional name for the evaluator. Defaults to None (uses function name).
+        scorer: Optional custom scoring function. Defaults to None (uses default scorer).
+
+    Returns:
+        A decorator that wraps a function as an Evaluator instance.
+    """
     if scorer is None:
         scorer = _default_eval_scorer
 
-    def wrapper(func: Callable[..., Any]) -> "Evaluator":
+    def wrapper(func: Callable[..., object]) -> "Evaluator":
         nonlocal name
         if not name:
             name = get_func_name(func)
-        assert name is not None
+        if name is None:
+            raise ValueError("Evaluator name cannot be None")
 
         wrapped_signature = inspect.signature(func)
         validate_evaluator_signature(wrapped_signature)
@@ -124,20 +153,22 @@ def create_evaluator(
 def _wrap_coroutine_evaluation_function(
     name: str,
     sig: inspect.Signature,
-    convert_to_score: Callable[[Any], EvaluationResult],
-) -> Callable[[Callable[..., Any]], "Evaluator"]:
+    convert_to_score: Callable[[object], EvaluationResult],
+) -> Callable[[Callable[..., object]], "Evaluator"]:
     from ..evaluators.base import Evaluator
 
-    def wrapper(func: Callable[..., Any]) -> "Evaluator":
+    def wrapper(func: Callable[..., object]) -> "Evaluator":
         class AsyncEvaluator(Evaluator):
             def __init__(self) -> None:
                 self._name = name
 
             @functools.wraps(func)
-            async def __call__(self, *args: Any, **kwargs: Any) -> Any:
+            async def __call__(self, *args: object, **kwargs: object) -> object:
                 return await func(*args, **kwargs)
 
-            async def async_evaluate(self, **kwargs: Any) -> EvaluationResult:
+            async def async_evaluate(
+                self, **kwargs: object
+            ) -> EvaluationResult:
                 bound_signature = _bind_evaluator_signature(sig, **kwargs)
                 result = await func(
                     *bound_signature.args, **bound_signature.kwargs
@@ -152,20 +183,20 @@ def _wrap_coroutine_evaluation_function(
 def _wrap_sync_evaluation_function(
     name: str,
     sig: inspect.Signature,
-    convert_to_score: Callable[[Any], EvaluationResult],
-) -> Callable[[Callable[..., Any]], "Evaluator"]:
+    convert_to_score: Callable[[object], EvaluationResult],
+) -> Callable[[Callable[..., object]], "Evaluator"]:
     from ..evaluators.base import Evaluator
 
-    def wrapper(func: Callable[..., Any]) -> "Evaluator":
+    def wrapper(func: Callable[..., object]) -> "Evaluator":
         class SyncEvaluator(Evaluator):
             def __init__(self) -> None:
                 self._name = name
 
             @functools.wraps(func)
-            def __call__(self, *args: Any, **kwargs: Any) -> Any:
+            def __call__(self, *args: object, **kwargs: object) -> object:
                 return func(*args, **kwargs)
 
-            def evaluate(self, **kwargs: Any) -> EvaluationResult:
+            def evaluate(self, **kwargs: object) -> EvaluationResult:
                 bound_signature = _bind_evaluator_signature(sig, **kwargs)
                 result = func(*bound_signature.args, **bound_signature.kwargs)
                 return convert_to_score(result)
@@ -175,7 +206,7 @@ def _wrap_sync_evaluation_function(
     return wrapper
 
 
-def _default_eval_scorer(result: Any) -> EvaluationResult:
+def _default_eval_scorer(result: object) -> EvaluationResult:
     if isinstance(result, EvaluationResult):
         return result
     if isinstance(result, bool):
@@ -193,6 +224,13 @@ def _default_eval_scorer(result: Any) -> EvaluationResult:
     raise ValueError(f"Unsupported evaluation result type: {type(result)}")
 
 
-def printif(condition: bool, *args: Any, **kwargs: Any) -> None:
+def printif(condition: bool, *args: object, **kwargs: object) -> None:
+    """Print to tqdm output if the condition is true.
+
+    Args:
+        condition: Whether to print the message.
+        *args: Positional arguments to pass to tqdm.write.
+        **kwargs: Keyword arguments to pass to tqdm.write.
+    """
     if condition:
         tqdm.write(*args, **kwargs)