orchestrata_ultimate.providers.rate_limiter

Reference › orchestrata_ultimate › providers › rate_limiter

rate_limiter tracks how much of a provider's request and token allowance is used (specification 011, user story US-004). RateLimiter is created per provider by provider_manager with a requests-per-minute limit (default 60) and a tokens-per-minute limit (default 100000), and read by the router. record_request adds a request and its token count to a 60-second window, which _maybe_reset_window clears when it has elapsed. headroom_pct gives the fraction of allowance left, is_available says whether more requests fit, has_low_headroom is true under 10% so traffic can be steered away early, seconds_until_reset gives the wait for the window, and get_unavailability_reason and get_status give readable detail. With a Redis client from set_redis, tracking is distributed.

At a glance

Class diagram

Dependency diagram

API

Classes

RateLimiter

Decorators: traces(spec='011', story='US-004', reqs=['REQ-011-041', 'REQ-011-042', 'REQ-011-043', 'REQ-011-044']) · Line 25

Per-provider rate limit tracker.

Source

Permalink: https://gitlab.com/MSamy/orchestrata-ultimate-architecture/-/blob/a4cbc2bdf/src/orchestrata_ultimate/providers/rate_limiter.py

"""Per-provider rate limit tracking.

SPEC-011 FR-010: Tracks RPM and TPM with sliding window.
Routes away from rate-limited providers.
"""

import logging
import threading
import time

import redis.asyncio as redis

from orchestrata_ultimate.core.base.traces import traces

logger = logging.getLogger(__name__)

_RL_PREFIX = "orchestrata:provider:rl"


@traces(
    spec="011",
    story="US-004",
    reqs=["REQ-011-041", "REQ-011-042", "REQ-011-043", "REQ-011-044"],
)
class RateLimiter:
    """Per-provider rate limit tracker.

    Tracks requests per minute (RPM) and tokens per minute (TPM)
    using a sliding window approach. Supports Redis persistence
    for distributed tracking.
    """

    @traces(spec="011", story="US-011-001", reqs=["REQ-011-041"])
    def __init__(
        self,
        provider_name: str,
        rpm_limit: int = 60,
        tpm_limit: int = 100_000,
    ):
        self.provider_name = provider_name
        self.rpm_limit = rpm_limit
        self.tpm_limit = tpm_limit
        self._window_start: float = time.monotonic()
        self._request_count: int = 0
        self._token_count: int = 0
        self._redis: redis.Redis | None = None
        # Guards the sliding-window counters so concurrent callers (free-threaded
        # 3.14, no GIL) cannot lose increments or window resets to a data race.
        self._lock = threading.Lock()

    @property
    @traces(spec="011", story="US-004", reqs=["REQ-011-041"])
    def headroom_pct(self) -> float:
        """Percentage of rate limit remaining (0.0 to 1.0)."""
        with self._lock:
            self._maybe_reset_window_locked()
            if self.rpm_limit <= 0:
                return 1.0  # unlimited
            rpm_headroom = 1.0 - (self._request_count / self.rpm_limit)
            if self.tpm_limit <= 0:
                return max(0.0, rpm_headroom)
            tpm_headroom = 1.0 - (self._token_count / self.tpm_limit)
            return max(0.0, min(rpm_headroom, tpm_headroom))

    @property
    @traces(spec="011", story="US-004", reqs=["REQ-011-041"])
    def is_available(self) -> bool:
        """Whether the provider has headroom for more requests."""
        return self.headroom_pct > 0.0

    @property
    @traces(spec="011", story="US-004", reqs=["REQ-011-041"])
    def seconds_until_reset(self) -> float:
        """Seconds remaining until the 60-second sliding window resets."""
        with self._lock:
            self._maybe_reset_window_locked()
            elapsed = time.monotonic() - self._window_start
            return max(0.0, 60.0 - elapsed)

    @traces(spec="011", story="US-004", reqs=["REQ-011-041"])
    def get_unavailability_reason(self) -> str:
        """Human-readable reason why the provider is currently unavailable.

        Returns:
            One of ``RPM exhausted``, ``TPM exhausted``, or
            ``rate limit exhausted`` when no headroom remains. Caller should
            guard with ``is_available``; this method assumes the limiter is
            exhausted.
        """
        with self._lock:
            self._maybe_reset_window_locked()
            rpm_exhausted = self.rpm_limit > 0 and self._request_count >= self.rpm_limit
            tpm_exhausted = self.tpm_limit > 0 and self._token_count >= self.tpm_limit
            if rpm_exhausted and tpm_exhausted:
                return "RPM and TPM exhausted"
            if rpm_exhausted:
                return "RPM exhausted"
            if tpm_exhausted:
                return "TPM exhausted"
            return "rate limit exhausted"

    @property
    @traces(spec="011", story="US-004", reqs=["REQ-011-043"])
    def has_low_headroom(self) -> bool:
        """Whether headroom is below 10% (should route away)."""
        return self.headroom_pct < 0.10

    @traces(spec="011", story="US-004", reqs=["REQ-011-041"])
    async def set_redis(self, client: redis.Redis) -> None:
        """Set Redis client for distributed tracking."""
        self._redis = client

    @traces(spec="011", story="US-011-001", reqs=["REQ-011-041"])
    async def record_request(self, token_count: int = 0) -> None:
        """Record a request and its token usage."""
        with self._lock:
            self._maybe_reset_window_locked()
            self._request_count += 1
            self._token_count += token_count
        if self._redis:
            key = f"{_RL_PREFIX}:{self.provider_name}"
            pipe = self._redis.pipeline()
            pipe.hincrby(key, "rpm", 1)
            pipe.hincrby(key, "tpm", token_count)
            pipe.expire(key, 120)
            await pipe.execute()

    def _maybe_reset_window_locked(self) -> None:
        """Reset counters if the 60-second window has elapsed.

        Caller MUST hold ``self._lock``.
        """
        now = time.monotonic()
        if now - self._window_start >= 60.0:
            self._window_start = now
            self._request_count = 0
            self._token_count = 0

    @traces(spec="011", story="US-011-001", reqs=["REQ-011-041"])
    def _maybe_reset_window(self) -> None:
        """Reset counters if the 60-second window has elapsed (acquires lock)."""
        with self._lock:
            self._maybe_reset_window_locked()

    @traces(spec="011", story="US-011-001", reqs=["REQ-011-041"])
    def get_status(self) -> dict:
        """Return current rate limit status."""
        with self._lock:
            self._maybe_reset_window_locked()
            request_count = self._request_count
            token_count = self._token_count
        return {
            "provider": self.provider_name,
            "rpm_used": request_count,
            "rpm_limit": self.rpm_limit,
            "tpm_used": token_count,
            "tpm_limit": self.tpm_limit,
            "headroom_pct": round(self.headroom_pct * 100, 1),
            "is_available": self.is_available,
        }