Migrated provider ratelimiter to aiolimiter

providers/nvidia_nim.py

@@ -7,7 +7,7 @@ import uuid
 from typing import Dict, Any, AsyncIterator
 
 import httpx
-from httpx import TimeoutException, ReadTimeout, ConnectTimeout
+from httpx import TimeoutException, ConnectTimeout
 
 from .base import BaseProvider, ProviderConfig
 from .utils import (
@@ -107,21 +107,19 @@ class NvidiaNimProvider(BaseProvider):
         self, request: Any, input_tokens: int = 0
     ) -> AsyncIterator[str]:
         """Stream response in Anthropic SSE format."""
-        # Wait if globally rate limited
-        blocked = await self._global_rate_limiter.wait_if_blocked()
-        if blocked:
-            # Yield error event for rate limit blocking
+        # Wait if globally rate limited (proactive throttle + reactive block)
+        waited_reactively = await self._global_rate_limiter.wait_if_blocked()
+
+        if waited_reactively:
+            # Yield error event for reactive rate limit blocking (user feedback)
             message_id = f"msg_{uuid.uuid4()}"
             sse = SSEBuilder(message_id, request.model, input_tokens)
-            error_msg = "⏱️ Rate limit exceeded. Please try again in a minute."
-            logger.warning(f"NIM_STREAM: Rate limit blocked, yielding error event")
+            error_msg = "⏱️ Global rate limit active. Resuming now..."
+            logger.info(f"NIM_STREAM: Reactive block detected, notified user")
             yield sse.message_start()
             for event in sse.emit_error(error_msg):
                 yield event
-            yield sse.message_delta("stop", 0)
-            yield sse.message_stop()
-            yield sse.done()
-            return
+            # After notification, we continue to the actual request
 
         body = self._build_request_body(request, stream=True)
         # Log compact request summary
@@ -228,13 +226,13 @@ class NvidiaNimProvider(BaseProvider):
             logger.error(f"NIM_ERROR: {type(e).__name__}: {e}")
             error_occurred = True
             error_message = str(e)
-            logger.info(f"NIM_STREAM: Emitting SSE error event for exception")
+            logger.info("NIM_STREAM: Emitting SSE error event for exception")
 
         # Handle errors
         if error_occurred:
             for event in sse.emit_error(error_message):
                 yield event
-            logger.info(f"NIM_STREAM: Error event yielded, total events emitted")
+            logger.info("NIM_STREAM: Error event yielded, total events emitted")
 
         # Flush remaining content from parsers
         remaining = think_parser.flush()
@@ -406,13 +404,8 @@ class NvidiaNimProvider(BaseProvider):
 
     async def complete(self, request: Any) -> dict:
         """Make a non-streaming completion request."""
-        # Wait if globally rate limited
-        wait_time = await self._global_rate_limiter.wait_if_blocked()
-        if wait_time > 0:
-            # Raise error for rate limit blocking in non-streaming mode
-            error_msg = "⏱️ Rate limit exceeded. Please try again in a minute."
-            logger.warning(f"NIM_COMPLETE: Rate limit blocked for {wait_time:.1f}s, raising error")
-            raise APIError(error_msg, status_code=429)
+        # Wait if globally rate limited (proactive throttle + reactive block)
+        await self._global_rate_limiter.wait_if_blocked()
 
         body = self._build_request_body(request, stream=False)
         # Log compact request summary

providers/rate_limit.py

@@ -3,7 +3,9 @@
 import asyncio
 import time
 import logging
+import os
 from typing import Optional
+from aiolimiter import AsyncLimiter
 
 logger = logging.getLogger(__name__)
 
@@ -11,17 +13,31 @@ logger = logging.getLogger(__name__)
 class GlobalRateLimiter:
     """
     Global singleton rate limiter that blocks all requests
-    when a rate limit error is encountered.
+    when a rate limit error is encountered (reactive) and
+    throttles requests (proactive) using aiolimiter.
 
-    No proactive limits - only reactive when 429 is hit.
-    No retry logic - just pauses all requests until cooldown expires.
+    Proactive limits - throttles requests to stay within API limits.
+    Reactive limits - pauses all requests when a 429 is hit.
     """
 
     _instance: Optional["GlobalRateLimiter"] = None
 
     def __init__(self):
+        # Prevent double initialization in singleton
+        if hasattr(self, "_initialized"):
+            return
+
+        rate_limit = int(os.getenv("NVIDIA_NIM_RATE_LIMIT", "40"))
+        rate_window = float(os.getenv("NVIDIA_NIM_RATE_WINDOW", "60.0"))
+
+        self.limiter = AsyncLimiter(rate_limit, rate_window)
         self._blocked_until: float = 0
         self._lock = asyncio.Lock()
+        self._initialized = True
+
+        logger.info(
+            f"GlobalRateLimiter (Provider) initialized ({rate_limit} req / {rate_window}s)"
+        )
 
     @classmethod
     def get_instance(cls) -> "GlobalRateLimiter":
@@ -37,33 +53,40 @@ class GlobalRateLimiter:
 
     async def wait_if_blocked(self) -> bool:
         """
-        Wait if currently rate limited.
+        Wait if currently rate limited or throttle to meet quota.
 
         Returns:
-            True if was blocked and waited, False if not blocked
+            True if was reactively blocked and waited, False otherwise.
         """
+        # 1. Reactive check: Wait if someone hit a 429
+        waited_reactively = False
         now = time.time()
         if now < self._blocked_until:
             wait_time = self._blocked_until - now
-            logger.warning(f"Global rate limit active, waiting {wait_time:.1f}s...")
+            logger.warning(
+                f"Global provider rate limit active (reactive), waiting {wait_time:.1f}s..."
+            )
             await asyncio.sleep(wait_time)
-            return True
-        return False
+            waited_reactively = True
+
+        # 2. Proactive check: Acquire slot from aiolimiter
+        async with self.limiter:
+            return waited_reactively
 
     def set_blocked(self, seconds: float = 60) -> None:
         """
-        Set global block for specified seconds.
+        Set global block for specified seconds (reactive).
 
         Args:
             seconds: How long to block (default 60s)
         """
         self._blocked_until = time.time() + seconds
-        logger.warning(f"Global rate limit set for {seconds:.1f}s")
+        logger.warning(f"Global provider rate limit set for {seconds:.1f}s (reactive)")
 
     def is_blocked(self) -> bool:
-        """Check if currently blocked."""
+        """Check if currently reactively blocked."""
         return time.time() < self._blocked_until
 
     def remaining_wait(self) -> float:
-        """Get remaining wait time in seconds."""
+        """Get remaining reactive wait time in seconds."""
         return max(0, self._blocked_until - time.time())

providers/utils/__init__.py

@@ -1,6 +1,5 @@
 """Utility modules for providers."""
 
-from .rate_limiter import SlidingWindowRateLimiter
 from .sse_builder import SSEBuilder, ContentBlockManager, map_stop_reason
 from .think_parser import (
     ThinkTagParser,
@@ -17,7 +16,6 @@ from .message_converter import (
 )
 
 __all__ = [
-    "SlidingWindowRateLimiter",
     "SSEBuilder",
     "ContentBlockManager",
     "map_stop_reason",

providers/utils/rate_limiter.py

@@ -1,72 +0,0 @@
-"""Reusable async rate limiter with sliding window."""
-
-import asyncio
-import time
-import threading
-from collections import deque
-from typing import Deque
-
-
-class SlidingWindowRateLimiter:
-    """
-    Async rate limiter using sliding window algorithm.
-
-    Thread-safe for use across multiple async contexts.
-    """
-
-    def __init__(
-        self,
-        rate_limit: int = 40,
-        window_seconds: int = 60,
-        max_retries: int = 120,
-    ):
-        self.rate_limit = rate_limit
-        self.window_seconds = window_seconds
-        self.max_retries = max_retries
-        self._timestamps: Deque[float] = deque()
-        self._lock = threading.Condition()
-
-    async def acquire(self) -> None:
-        """
-        Acquire a rate limit slot, waiting if necessary.
-
-        Uses exponential backoff up to max_retries attempts.
-        """
-        for _ in range(self.max_retries):
-            now = time.time()
-
-            with self._lock:
-                # Remove expired timestamps
-                while self._timestamps and now - self._timestamps[0] > self.window_seconds:
-                    self._timestamps.popleft()
-
-                # Check if we can proceed
-                if len(self._timestamps) < self.rate_limit:
-                    self._timestamps.append(now)
-                    return
-
-                # Calculate wait time
-                wait_time = self._timestamps[0] + self.window_seconds - now
-
-            if wait_time <= 0:
-                continue
-
-            await asyncio.sleep(min(wait_time, 1.0))
-
-        # Fallback: allow request after max retries
-        with self._lock:
-            self._timestamps.append(time.time())
-
-    def reset(self) -> None:
-        """Clear all timestamps."""
-        with self._lock:
-            self._timestamps.clear()
-
-    @property
-    def current_count(self) -> int:
-        """Current number of requests in the window."""
-        now = time.time()
-        with self._lock:
-            while self._timestamps and now - self._timestamps[0] > self.window_seconds:
-                self._timestamps.popleft()
-            return len(self._timestamps)

tests/verify_provider_limiter.py

@@ -0,0 +1,87 @@
+import asyncio
+import os
+import time
+import logging
+import sys
+
+# Add current directory to path
+sys.path.append(os.getcwd())
+
+from providers.rate_limit import GlobalRateLimiter
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+async def test_proactive_throttling():
+    print("\n--- Testing Proactive Throttling (3 req / 1 sec) ---")
+    os.environ["NVIDIA_NIM_RATE_LIMIT"] = "3"
+    os.environ["NVIDIA_NIM_RATE_WINDOW"] = "1.0"
+
+    GlobalRateLimiter.reset_instance()
+    limiter = GlobalRateLimiter.get_instance()
+
+    start_time = time.time()
+
+    async def call_limiter(i):
+        await limiter.wait_if_blocked()
+        print(f"[{time.time() - start_time:.2f}s] Request {i} passed")
+
+    print("Sending 5 requests sequentially...")
+    # 3 should pass immediately, 4th and 5th should wait
+    for i in range(5):
+        await call_limiter(i)
+
+    total_time = time.time() - start_time
+    print(f"All requests completed in {total_time:.2f}s")
+
+    # 5 requests with limit 3 per 1s:
+    # R0, R1, R2 at ~0s
+    # R3 waits until R0 is 1s old -> ~1s
+    # R4 waits until R1 is 1s old -> ~1s
+    if total_time >= 0.9:  # Allow some jitter
+        print("[SUCCESS] Proactive throttling working!")
+    else:
+        print(f"[FAILURE] Proactive throttling failed! Took only {total_time:.2f}s")
+
+
+async def test_reactive_blocking():
+    print("\n--- Testing Reactive Blocking ---")
+    GlobalRateLimiter.reset_instance()
+    limiter = GlobalRateLimiter.get_instance()
+
+    start_time = time.time()
+
+    print("Setting manual block for 2s...")
+    limiter.set_blocked(2)
+
+    async def call_limiter(i):
+        waited = await limiter.wait_if_blocked()
+        print(
+            f"[{time.time() - start_time:.2f}s] Request {i} passed (waited reactively: {waited})"
+        )
+        return waited
+
+    results = await asyncio.gather(*[call_limiter(i) for i in range(2)])
+
+    total_time = time.time() - start_time
+    print(f"Requests completed in {total_time:.2f}s")
+
+    if total_time >= 1.9 and any(results):
+        print("[SUCCESS] Reactive blocking working!")
+    else:
+        print(f"[FAILURE] Reactive blocking failed! Took only {total_time:.2f}s")
+
+
+async def main():
+    try:
+        await test_proactive_throttling()
+        await test_reactive_blocking()
+    except Exception as e:
+        print(f"Error during verification: {e}")
+    finally:
+        print("\nVerification complete.")
+
+
+if __name__ == "__main__":
+    asyncio.run(main())