@arizeai/phoenix-mcp

import time from contextlib import contextmanager from math import exp, isclose from typing import Any, Callable, Iterator, List, Optional from unittest import mock import pytest from phoenix.client.utils.rate_limiters import ( AdaptiveTokenBucket, UnavailableTokensError, ) @contextmanager def freeze_time(frozen_time: Optional[float] = None) -> Iterator[Callable[[], None]]: frozen_time = time.time() if frozen_time is None else frozen_time with mock.patch("time.time") as mock_time: mock_time.return_value = frozen_time yield mock_time @contextmanager def warp_time(start: Optional[float]) -> Iterator[Any]: sleeps: List[float] = [0] current_time = start if start is not None else time.time() def instant_sleep(sleep_time: float) -> None: nonlocal sleeps sleeps.append(sleep_time) def time_warp() -> float: try: nonlocal current_time nonlocal sleeps current_time += sleeps.pop() return current_time except IndexError: return current_time with mock.patch("time.time") as mock_time: with mock.patch("time.sleep") as mock_sleep: mock_sleep.side_effect = instant_sleep mock_time.side_effect = time_warp yield mock_sleep @contextmanager def async_warp_time(start: Optional[float]) -> Iterator[Any]: sleeps: List[float] = [0] current_time = start if start is not None else time.time() def instant_sleep(sleep_time: float) -> None: nonlocal sleeps sleeps.append(sleep_time) def time_warp() -> float: try: nonlocal current_time nonlocal sleeps current_time += sleeps.pop() return current_time except IndexError: return current_time with mock.patch("time.time") as mock_time: with mock.patch("asyncio.sleep") as mock_sleep: mock_sleep.side_effect = instant_sleep mock_time.side_effect = time_warp yield mock_sleep def test_token_bucket_gains_tokens_over_time() -> None: start = time.time() with freeze_time(start): bucket = AdaptiveTokenBucket( initial_per_second_request_rate=1, maximum_per_second_request_rate=1, enforcement_window_minutes=1, rate_reduction_factor=1, rate_increase_factor=0, cooldown_seconds=5, ) with freeze_time(start + 5): assert isclose(bucket.available_requests(), 5) with freeze_time(start + 10): assert isclose(bucket.available_requests(), 10) def test_token_rate_limiter_can_max_out_on_requests() -> None: start = time.time() with freeze_time(start): bucket = AdaptiveTokenBucket( initial_per_second_request_rate=1, maximum_per_second_request_rate=1, enforcement_window_minutes=2, rate_reduction_factor=1, rate_increase_factor=0, cooldown_seconds=5, ) with freeze_time(start + 30): assert bucket.available_requests() == 30 with freeze_time(start + 120): assert bucket.available_requests() == 120 with freeze_time(start + 130): assert bucket.available_requests() == 120 def test_token_rate_limiter_spends_tokens() -> None: start = time.time() with freeze_time(start): bucket = AdaptiveTokenBucket( initial_per_second_request_rate=1, maximum_per_second_request_rate=1, enforcement_window_minutes=1, rate_reduction_factor=1, rate_increase_factor=0, cooldown_seconds=5, ) with freeze_time(start + 3): assert bucket.available_requests() == 3 bucket.make_request_if_ready() assert bucket.available_requests() == 2 def test_token_rate_limiter_cannot_spend_unavailable_tokens() -> None: start = time.time() with freeze_time(start): bucket = AdaptiveTokenBucket( initial_per_second_request_rate=1, maximum_per_second_request_rate=1, enforcement_window_minutes=2, rate_reduction_factor=1, rate_increase_factor=0, cooldown_seconds=5, ) assert bucket.available_requests() == 0 with pytest.raises(UnavailableTokensError): bucket.make_request_if_ready() def test_token_rate_limiter_can_block_until_tokens_are_available() -> None: start = time.time() with freeze_time(start): rate = 0.5 bucket = AdaptiveTokenBucket( initial_per_second_request_rate=rate, maximum_per_second_request_rate=rate * 2, enforcement_window_minutes=2, rate_reduction_factor=1, rate_increase_factor=0, cooldown_seconds=5, ) with warp_time(start) as mock_sleep: assert bucket.available_requests() == 0 bucket.wait_until_ready() sleeps = [call.args[0] for call in mock_sleep.call_args_list] time_cost = 1 / rate assert isclose(sum(sleeps), time_cost, rel_tol=0.2) async def test_token_rate_limiter_async_waits_until_tokens_are_available() -> None: start = time.time() with freeze_time(start): rate = 0.5 bucket = AdaptiveTokenBucket( initial_per_second_request_rate=rate, maximum_per_second_request_rate=rate * 2, enforcement_window_minutes=2, rate_reduction_factor=1, rate_increase_factor=0, cooldown_seconds=5, ) with async_warp_time(start) as mock_sleep: assert bucket.available_requests() == 0 await bucket.async_wait_until_ready() sleeps = [call.args[0] for call in mock_sleep.call_args_list] time_cost = 1 / rate assert isclose(sum(sleeps), time_cost, rel_tol=0.2) def test_token_rate_limiter_can_accumulate_tokens_before_waiting() -> None: start = time.time() with freeze_time(start): rate = 0.1 bucket = AdaptiveTokenBucket( initial_per_second_request_rate=rate, maximum_per_second_request_rate=rate * 2, enforcement_window_minutes=2, rate_reduction_factor=1, rate_increase_factor=0, cooldown_seconds=5, ) with warp_time(start + 5) as mock_sleep: assert bucket.available_requests() == 0.5, "should have accumulated half a request" bucket.wait_until_ready() sleeps = [call.args[0] for call in mock_sleep.call_args_list] time_cost = (1 / rate) - 5 assert isclose(sum(sleeps), time_cost, rel_tol=0.2) async def test_token_rate_limiter_can_async_accumulate_tokens_before_waiting() -> None: start = time.time() with freeze_time(start): rate = 0.1 bucket = AdaptiveTokenBucket( initial_per_second_request_rate=rate, maximum_per_second_request_rate=rate * 2, enforcement_window_minutes=2, rate_reduction_factor=1, rate_increase_factor=0, cooldown_seconds=5, ) with async_warp_time(start + 5) as mock_sleep: assert bucket.available_requests() == 0.5, "should have accumulated half a request" await bucket.async_wait_until_ready() sleeps = [call.args[0] for call in mock_sleep.call_args_list] time_cost = (1 / rate) - 5 assert isclose(sum(sleeps), time_cost, rel_tol=0.2) def test_token_bucket_adaptively_increases_rate_over_time() -> None: start = time.time() with freeze_time(start): rate = 0.1 bucket = AdaptiveTokenBucket( initial_per_second_request_rate=rate, maximum_per_second_request_rate=rate * 2, enforcement_window_minutes=1, rate_reduction_factor=1, rate_increase_factor=0.01, cooldown_seconds=5, ) with warp_time(start + 5) as mock_sleep: assert bucket.available_requests() == 0.5, "should have accumulated half a request" bucket.wait_until_ready() sleeps = [call.args[0] for call in mock_sleep.call_args_list] elapsed_time = sum(sleeps) + 5 assert isclose(bucket.rate, 0.1 * exp(0.01 * elapsed_time)) def test_token_bucket_does_not_increase_rate_past_maximum() -> None: start = time.time() with freeze_time(start): rate = 0.1 bucket = AdaptiveTokenBucket( initial_per_second_request_rate=rate, maximum_per_second_request_rate=rate * 2, enforcement_window_minutes=1, rate_reduction_factor=1, rate_increase_factor=100, cooldown_seconds=5, ) with warp_time(start + 5): assert bucket.available_requests() == 0.5, "should have accumulated half a request" bucket.wait_until_ready() assert isclose(bucket.rate, rate * 2) def test_token_bucket_resets_rate_after_inactivity() -> None: start = time.time() with freeze_time(start): rate = 0.1 bucket = AdaptiveTokenBucket( initial_per_second_request_rate=rate, maximum_per_second_request_rate=rate * 2, enforcement_window_minutes=1, rate_reduction_factor=1, rate_increase_factor=100, cooldown_seconds=5, ) with warp_time(start + 5): assert bucket.available_requests() == 0.5, "should have accumulated half a request" bucket.wait_until_ready() assert isclose(bucket.rate, rate * 2) with warp_time(start + 100): bucket.wait_until_ready() assert isclose(bucket.rate, rate) def test_token_bucket_decreases_rate() -> None: start = time.time() with warp_time(start): rate = 100 bucket = AdaptiveTokenBucket( initial_per_second_request_rate=rate, maximum_per_second_request_rate=rate * 2, enforcement_window_minutes=1, rate_reduction_factor=0.25, rate_increase_factor=0.01, cooldown_seconds=5, ) bucket.on_rate_limit_error(request_start_time=time.time()) assert isclose(bucket.rate, 25) assert bucket.tokens == 0 assert time.time() == start + 5 def test_token_bucket_decreases_rate_once_per_cooldown_period() -> None: start = time.time() with warp_time(start): rate = 100 bucket = AdaptiveTokenBucket( initial_per_second_request_rate=rate, maximum_per_second_request_rate=rate * 2, enforcement_window_minutes=1, rate_reduction_factor=0.25, rate_increase_factor=0.01, cooldown_seconds=5, ) bucket.on_rate_limit_error(request_start_time=time.time()) assert isclose(bucket.rate, 25) with warp_time(start + 3): bucket.on_rate_limit_error(request_start_time=time.time()) assert isclose(bucket.rate, 25), "3 seconds is still within the cooldown period" with warp_time(start - 6): bucket.on_rate_limit_error(request_start_time=time.time()) assert isclose(bucket.rate, 25), "requests before the rate limited request are ignored" with warp_time(start + 6): bucket.on_rate_limit_error(request_start_time=time.time()) assert isclose(bucket.rate, 6.25)