228 lines
7.7 KiB
Python
228 lines
7.7 KiB
Python
import types
|
|
import unittest
|
|
from unittest import mock
|
|
|
|
from openai import APIStatusError, RateLimitError
|
|
|
|
import util.translation as T
|
|
from util.translation import AdaptiveLimiter, _header_float, _header_int
|
|
|
|
|
|
class FakeClock:
|
|
"""Deterministic stand-in for the ``time`` module.
|
|
|
|
``sleep`` advances the clock instead of waiting, so pacing/backoff paths
|
|
resolve instantly and predictably.
|
|
"""
|
|
|
|
def __init__(self, start=1000.0):
|
|
self.t = start
|
|
|
|
def monotonic(self):
|
|
return self.t
|
|
|
|
def time(self):
|
|
return self.t
|
|
|
|
def sleep(self, seconds):
|
|
self.t += seconds
|
|
|
|
|
|
def make_exc(cls, **attrs):
|
|
"""Build an openai exception without running its real constructor."""
|
|
err = cls.__new__(cls)
|
|
for key, value in attrs.items():
|
|
setattr(err, key, value)
|
|
return err
|
|
|
|
|
|
class FakeRaw:
|
|
def __init__(self, headers, parsed):
|
|
self._headers = headers
|
|
self._parsed = parsed
|
|
|
|
@property
|
|
def headers(self):
|
|
return self._headers
|
|
|
|
def parse(self):
|
|
return self._parsed
|
|
|
|
|
|
class FakeCreate:
|
|
"""Callable that replays a scripted sequence of returns/raises."""
|
|
|
|
def __init__(self, behaviors):
|
|
self.behaviors = behaviors
|
|
self.calls = []
|
|
|
|
def __call__(self, **params):
|
|
self.calls.append(params)
|
|
behavior = self.behaviors[len(self.calls) - 1]
|
|
if isinstance(behavior, Exception):
|
|
raise behavior
|
|
return behavior
|
|
|
|
|
|
def fake_openai(create):
|
|
completions = types.SimpleNamespace(
|
|
with_raw_response=types.SimpleNamespace(create=create)
|
|
)
|
|
return types.SimpleNamespace(chat=types.SimpleNamespace(completions=completions))
|
|
|
|
|
|
class HeaderHelperTests(unittest.TestCase):
|
|
def test_header_int_first_present_and_floored(self):
|
|
headers = {"b": "125.9"}
|
|
self.assertEqual(_header_int(headers, "a", "b"), 125)
|
|
|
|
def test_header_int_missing_blank_and_garbage(self):
|
|
self.assertIsNone(_header_int({}, "a"))
|
|
self.assertIsNone(_header_int({"a": " "}, "a"))
|
|
self.assertIsNone(_header_int({"a": "abc"}, "a"))
|
|
|
|
def test_header_float_keeps_fraction(self):
|
|
self.assertAlmostEqual(_header_float({"a": "0.42"}, "a"), 0.42)
|
|
|
|
|
|
class AdaptiveLimiterUpdateTests(unittest.TestCase):
|
|
def _limiter(self):
|
|
# Seeded at 0.5 RPS / 50k TPM, matching the production default.
|
|
return AdaptiveLimiter(0.5, 50000, 4000)
|
|
|
|
def test_canonical_headers_correct_pace_and_budget(self):
|
|
lim = self._limiter()
|
|
lim.update({
|
|
"x-ratelimit-limit-requests": "125",
|
|
"x-ratelimit-limit-tokens": "500000",
|
|
"x-ratelimit-remaining-tokens": "498000",
|
|
})
|
|
self.assertAlmostEqual(lim.req_per_sec, 125 / 60.0)
|
|
self.assertAlmostEqual(lim.min_interval, 60.0 / 125)
|
|
self.assertEqual(lim.tok_per_min, 500000)
|
|
self.assertEqual(lim.tokens_remaining, 498000)
|
|
|
|
def test_legacy_minute_spelling_still_works(self):
|
|
lim = self._limiter()
|
|
lim.update({
|
|
"x-ratelimit-limit-req-minute": "25",
|
|
"x-ratelimit-limit-tokens-minute": "60000",
|
|
"x-ratelimit-remaining-tokens-minute": "59000",
|
|
})
|
|
self.assertAlmostEqual(lim.req_per_sec, 25 / 60.0)
|
|
self.assertEqual(lim.tok_per_min, 60000)
|
|
self.assertEqual(lim.tokens_remaining, 59000)
|
|
|
|
def test_true_per_second_header_used_undivided(self):
|
|
lim = self._limiter()
|
|
lim.update({"x-ratelimit-limit-req-second": "2"})
|
|
self.assertAlmostEqual(lim.req_per_sec, 2.0)
|
|
|
|
def test_empty_or_missing_headers_leave_seed_untouched(self):
|
|
for headers in (None, {}):
|
|
lim = self._limiter()
|
|
lim.update(headers)
|
|
self.assertAlmostEqual(lim.req_per_sec, 0.5)
|
|
self.assertEqual(lim.tok_per_min, 50000)
|
|
|
|
def test_garbage_values_ignored(self):
|
|
lim = self._limiter()
|
|
lim.update({"x-ratelimit-limit-requests": "abc"})
|
|
self.assertAlmostEqual(lim.req_per_sec, 0.5)
|
|
|
|
def test_exhausted_request_budget_holds_off(self):
|
|
clock = FakeClock()
|
|
with mock.patch.object(T, "time", clock):
|
|
lim = self._limiter()
|
|
lim.update({"x-ratelimit-remaining-requests": "0"})
|
|
self.assertGreaterEqual(lim.next_request_at, clock.monotonic() + 60.0)
|
|
|
|
|
|
class AdaptiveLimiterAcquireTests(unittest.TestCase):
|
|
def test_requests_spaced_by_min_interval(self):
|
|
clock = FakeClock()
|
|
with mock.patch.object(T, "time", clock):
|
|
lim = AdaptiveLimiter(2.0, 10_000_000, 4000) # min_interval = 0.5s
|
|
lim.acquire(10) # first clears immediately
|
|
self.assertEqual(clock.t, 1000.0)
|
|
lim.acquire(10) # second waits one interval
|
|
self.assertAlmostEqual(clock.t, 1000.5)
|
|
self.assertEqual(lim.tokens_remaining, 10_000_000 - 20)
|
|
|
|
def test_token_gate_waits_for_window_reset(self):
|
|
clock = FakeClock()
|
|
with mock.patch.object(T, "time", clock):
|
|
lim = AdaptiveLimiter(1000.0, 10000, 2000)
|
|
lim.tokens_remaining = 2500 # near depletion: 2500 - 1000 <= 2000
|
|
original_reset = lim.window_reset
|
|
lim.acquire(1000)
|
|
# It blocked until the minute window rolled over.
|
|
self.assertGreaterEqual(clock.t, original_reset)
|
|
# After the reset the budget refills, then this call is charged.
|
|
self.assertEqual(lim.tokens_remaining, 10000 - 1000)
|
|
|
|
|
|
class CallMistralTests(unittest.TestCase):
|
|
def setUp(self):
|
|
self._saved_limiter = T._mistral_limiter
|
|
T._mistral_limiter = None # force a fresh limiter from the seeds
|
|
self.clock = FakeClock()
|
|
self._patches = [
|
|
mock.patch.object(T, "time", self.clock),
|
|
mock.patch.dict(T.os.environ, {}, clear=False),
|
|
]
|
|
for patch in self._patches:
|
|
patch.start()
|
|
|
|
def tearDown(self):
|
|
for patch in self._patches:
|
|
patch.stop()
|
|
T._mistral_limiter = self._saved_limiter
|
|
|
|
def _params(self, **extra):
|
|
base = {"messages": [{"content": "テスト"}], "max_tokens": 100}
|
|
base.update(extra)
|
|
return base
|
|
|
|
def test_success_returns_response_and_syncs_limiter(self):
|
|
sentinel = object()
|
|
create = FakeCreate([
|
|
FakeRaw({"x-ratelimit-limit-requests": "125"}, sentinel),
|
|
])
|
|
with mock.patch.object(T, "openai", fake_openai(create)):
|
|
result = T.callMistral(self._params())
|
|
self.assertIs(result, sentinel)
|
|
self.assertEqual(len(create.calls), 1)
|
|
self.assertAlmostEqual(T._get_mistral_limiter().req_per_sec, 125 / 60.0)
|
|
|
|
def test_rate_limit_then_success_retries(self):
|
|
sentinel = object()
|
|
rate_limited = make_exc(
|
|
RateLimitError,
|
|
response=types.SimpleNamespace(headers={"retry-after": "0.01"}),
|
|
)
|
|
create = FakeCreate([
|
|
rate_limited,
|
|
FakeRaw({"x-ratelimit-limit-requests": "60"}, sentinel),
|
|
])
|
|
with mock.patch.object(T, "openai", fake_openai(create)):
|
|
result = T.callMistral(self._params())
|
|
self.assertIs(result, sentinel)
|
|
self.assertEqual(len(create.calls), 2)
|
|
|
|
def test_json_schema_downgraded_to_json_object(self):
|
|
sentinel = object()
|
|
bad_format = make_exc(APIStatusError, status_code=400)
|
|
create = FakeCreate([
|
|
bad_format,
|
|
FakeRaw({}, sentinel),
|
|
])
|
|
params = self._params(response_format={"type": "json_schema"})
|
|
with mock.patch.object(T, "openai", fake_openai(create)):
|
|
result = T.callMistral(params)
|
|
self.assertIs(result, sentinel)
|
|
self.assertEqual(create.calls[1]["response_format"], {"type": "json_object"})
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|