keycycle 0.3.2__tar.gz → 0.5.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {keycycle-0.3.2 → keycycle-0.5.0}/PKG-INFO +1 -1
- {keycycle-0.3.2 → keycycle-0.5.0}/keycycle/adapters/generic_adapter.py +51 -26
- {keycycle-0.3.2 → keycycle-0.5.0}/keycycle/adapters/openai_adapter.py +56 -27
- {keycycle-0.3.2 → keycycle-0.5.0}/keycycle/config/constants.py +7 -0
- {keycycle-0.3.2 → keycycle-0.5.0}/keycycle/core/__init__.py +2 -0
- {keycycle-0.3.2 → keycycle-0.5.0}/keycycle/core/utils.py +23 -0
- {keycycle-0.3.2 → keycycle-0.5.0}/keycycle/key_rotation/rotating_mixin.py +37 -19
- {keycycle-0.3.2 → keycycle-0.5.0}/keycycle/key_rotation/rotation_manager.py +107 -6
- {keycycle-0.3.2 → keycycle-0.5.0}/keycycle/legacy_multi_provider_wrapper.py +30 -4
- {keycycle-0.3.2 → keycycle-0.5.0}/keycycle/multi_client_wrapper.py +39 -3
- {keycycle-0.3.2 → keycycle-0.5.0}/keycycle.egg-info/PKG-INFO +1 -1
- {keycycle-0.3.2 → keycycle-0.5.0}/keycycle.egg-info/SOURCES.txt +4 -1
- {keycycle-0.3.2 → keycycle-0.5.0}/pyproject.toml +1 -1
- {keycycle-0.3.2 → keycycle-0.5.0}/tests/test_api_key_pinning.py +4 -2
- keycycle-0.5.0/tests/test_dead_key_rotation.py +124 -0
- keycycle-0.5.0/tests/test_key_wait.py +78 -0
- keycycle-0.5.0/tests/test_usage_tracking_toggle.py +503 -0
- {keycycle-0.3.2 → keycycle-0.5.0}/MANIFEST.in +0 -0
- {keycycle-0.3.2 → keycycle-0.5.0}/keycycle/__init__.py +0 -0
- {keycycle-0.3.2 → keycycle-0.5.0}/keycycle/adapters/__init__.py +0 -0
- {keycycle-0.3.2 → keycycle-0.5.0}/keycycle/config/__init__.py +0 -0
- {keycycle-0.3.2 → keycycle-0.5.0}/keycycle/config/dataclasses.py +0 -0
- {keycycle-0.3.2 → keycycle-0.5.0}/keycycle/config/enums.py +0 -0
- {keycycle-0.3.2 → keycycle-0.5.0}/keycycle/config/loader.py +0 -0
- {keycycle-0.3.2 → keycycle-0.5.0}/keycycle/config/log_config.py +0 -0
- {keycycle-0.3.2 → keycycle-0.5.0}/keycycle/config/models/cerebras.yaml +0 -0
- {keycycle-0.3.2 → keycycle-0.5.0}/keycycle/config/models/cohere.yaml +0 -0
- {keycycle-0.3.2 → keycycle-0.5.0}/keycycle/config/models/gemini.yaml +0 -0
- {keycycle-0.3.2 → keycycle-0.5.0}/keycycle/config/models/groq.yaml +0 -0
- {keycycle-0.3.2 → keycycle-0.5.0}/keycycle/config/models/mistral.yaml +0 -0
- {keycycle-0.3.2 → keycycle-0.5.0}/keycycle/config/models/moonshot.yaml +0 -0
- {keycycle-0.3.2 → keycycle-0.5.0}/keycycle/config/models/openrouter.yaml +0 -0
- {keycycle-0.3.2 → keycycle-0.5.0}/keycycle/config/models/openrouter_models.yaml +0 -0
- {keycycle-0.3.2 → keycycle-0.5.0}/keycycle/config/models.py +0 -0
- {keycycle-0.3.2 → keycycle-0.5.0}/keycycle/core/backoff.py +0 -0
- {keycycle-0.3.2 → keycycle-0.5.0}/keycycle/core/exceptions.py +0 -0
- {keycycle-0.3.2 → keycycle-0.5.0}/keycycle/key_rotation/__init__.py +0 -0
- {keycycle-0.3.2 → keycycle-0.5.0}/keycycle/py.typed +0 -0
- {keycycle-0.3.2 → keycycle-0.5.0}/keycycle/usage/__init__.py +0 -0
- {keycycle-0.3.2 → keycycle-0.5.0}/keycycle/usage/db_logic.py +0 -0
- {keycycle-0.3.2 → keycycle-0.5.0}/keycycle/usage/usage_logger.py +0 -0
- {keycycle-0.3.2 → keycycle-0.5.0}/keycycle/utils.py +0 -0
- {keycycle-0.3.2 → keycycle-0.5.0}/keycycle.egg-info/dependency_links.txt +0 -0
- {keycycle-0.3.2 → keycycle-0.5.0}/keycycle.egg-info/requires.txt +0 -0
- {keycycle-0.3.2 → keycycle-0.5.0}/keycycle.egg-info/top_level.txt +0 -0
- {keycycle-0.3.2 → keycycle-0.5.0}/setup.cfg +0 -0
- {keycycle-0.3.2 → keycycle-0.5.0}/tests/test_api_key_retrieval.py +0 -0
- {keycycle-0.3.2 → keycycle-0.5.0}/tests/test_cerebras_direct.py +0 -0
- {keycycle-0.3.2 → keycycle-0.5.0}/tests/test_custom_key_limits.py +0 -0
- {keycycle-0.3.2 → keycycle-0.5.0}/tests/test_multi_client_wrapper.py +0 -0
- {keycycle-0.3.2 → keycycle-0.5.0}/tests/test_payment_required_rotation.py +0 -0
- {keycycle-0.3.2 → keycycle-0.5.0}/tests/test_rate_limit_logic.py +0 -0
- {keycycle-0.3.2 → keycycle-0.5.0}/tests/test_specific_key_retrieval.py +0 -0
- {keycycle-0.3.2 → keycycle-0.5.0}/tests/test_streams_recording.py +0 -0
|
@@ -35,7 +35,7 @@ from ..config.constants import (
|
|
|
35
35
|
from ..core.utils import (
|
|
36
36
|
is_rate_limit_error,
|
|
37
37
|
is_temporary_rate_limit_error,
|
|
38
|
-
|
|
38
|
+
is_dead_key_error,
|
|
39
39
|
get_key_suffix,
|
|
40
40
|
)
|
|
41
41
|
from ..core.backoff import ExponentialBackoff, BackoffConfig
|
|
@@ -275,13 +275,20 @@ class BaseGenericRotatingClient(Generic[T]):
|
|
|
275
275
|
|
|
276
276
|
return self.config.client_class(**final_kwargs)
|
|
277
277
|
|
|
278
|
-
def _record_usage(
|
|
278
|
+
def _record_usage(
|
|
279
|
+
self,
|
|
280
|
+
key_usage: KeyUsage,
|
|
281
|
+
model_id: str,
|
|
282
|
+
actual_tokens: int,
|
|
283
|
+
track_usage: Optional[bool] = None,
|
|
284
|
+
) -> None:
|
|
279
285
|
"""Record usage for a key."""
|
|
280
286
|
self.manager.record_usage(
|
|
281
287
|
key_obj=key_usage,
|
|
282
288
|
model_id=model_id,
|
|
283
289
|
actual_tokens=actual_tokens,
|
|
284
290
|
estimated_tokens=self.config.estimated_tokens,
|
|
291
|
+
track_usage=track_usage,
|
|
285
292
|
)
|
|
286
293
|
|
|
287
294
|
def _extract_usage(self, response: Any) -> int:
|
|
@@ -301,13 +308,16 @@ class SyncGenericRotatingClient(BaseGenericRotatingClient[T]):
|
|
|
301
308
|
|
|
302
309
|
def _execute(self, path: List[str], args: tuple, kwargs: dict) -> Any:
|
|
303
310
|
"""Execute a method call with key rotation."""
|
|
311
|
+
# Per-call tracking override - popped before the kwargs reach the real
|
|
312
|
+
# client, which may reject unknown keyword arguments.
|
|
313
|
+
track_usage = kwargs.pop("track_usage", None)
|
|
304
314
|
model_id = self._get_model_id(kwargs)
|
|
305
315
|
limits = self.limit_resolver(model_id, None)
|
|
306
316
|
|
|
307
317
|
for attempt in range(self.config.max_retries + 1):
|
|
308
|
-
key_usage = self.manager.
|
|
318
|
+
key_usage = self.manager.acquire_key(model_id, limits, self.config.estimated_tokens)
|
|
309
319
|
if not key_usage:
|
|
310
|
-
raise RuntimeError(
|
|
320
|
+
raise RuntimeError(self.manager.unavailable_message(model_id))
|
|
311
321
|
|
|
312
322
|
# Create backoff for temporary rate limits
|
|
313
323
|
temp_backoff = ExponentialBackoff(BackoffConfig(
|
|
@@ -331,9 +341,11 @@ class SyncGenericRotatingClient(BaseGenericRotatingClient[T]):
|
|
|
331
341
|
if hasattr(result, "__iter__") and not isinstance(result, (str, bytes, dict, list)):
|
|
332
342
|
# Check if it looks like a generator/iterator
|
|
333
343
|
if hasattr(result, "__next__") or inspect.isgenerator(result):
|
|
334
|
-
return self._wrap_stream(result, key_usage, model_id)
|
|
344
|
+
return self._wrap_stream(result, key_usage, model_id, track_usage)
|
|
335
345
|
|
|
336
|
-
self._record_usage(
|
|
346
|
+
self._record_usage(
|
|
347
|
+
key_usage, model_id, self._extract_usage(result), track_usage
|
|
348
|
+
)
|
|
337
349
|
return result
|
|
338
350
|
|
|
339
351
|
except Exception as e:
|
|
@@ -348,11 +360,11 @@ class SyncGenericRotatingClient(BaseGenericRotatingClient[T]):
|
|
|
348
360
|
time.sleep(delay)
|
|
349
361
|
continue # Retry with SAME key
|
|
350
362
|
|
|
351
|
-
#
|
|
363
|
+
# Dead key (401/403 revoked/denied, 402 credits exhausted) - key is dead
|
|
352
364
|
# for this process, rotate to next key
|
|
353
|
-
if
|
|
365
|
+
if is_dead_key_error(e) and attempt < self.config.max_retries:
|
|
354
366
|
logger.warning(
|
|
355
|
-
"402/
|
|
367
|
+
"Dead key (401/402/403) hit for %s on key ...%s. Marking key dead and rotating. (Attempt %d/%d)",
|
|
356
368
|
model_id, get_key_suffix(key_usage.api_key),
|
|
357
369
|
attempt + 1, self.config.max_retries + 1
|
|
358
370
|
)
|
|
@@ -373,7 +385,7 @@ class SyncGenericRotatingClient(BaseGenericRotatingClient[T]):
|
|
|
373
385
|
time.sleep(KEY_ROTATION_DELAY_SECONDS)
|
|
374
386
|
break # Break inner loop, continue outer loop with new key
|
|
375
387
|
|
|
376
|
-
self._record_usage(key_usage, model_id, 0)
|
|
388
|
+
self._record_usage(key_usage, model_id, 0, track_usage)
|
|
377
389
|
raise
|
|
378
390
|
else:
|
|
379
391
|
# Inner loop exhausted without success - continue to next key
|
|
@@ -389,7 +401,11 @@ class SyncGenericRotatingClient(BaseGenericRotatingClient[T]):
|
|
|
389
401
|
raise RuntimeError(f"All retry attempts exhausted for {model_id}")
|
|
390
402
|
|
|
391
403
|
def _wrap_stream(
|
|
392
|
-
self,
|
|
404
|
+
self,
|
|
405
|
+
generator: Generator,
|
|
406
|
+
key_usage: KeyUsage,
|
|
407
|
+
model_id: str,
|
|
408
|
+
track_usage: Optional[bool] = None,
|
|
393
409
|
) -> Generator:
|
|
394
410
|
"""Wrap a streaming response to track usage and handle errors."""
|
|
395
411
|
final_tokens = 0
|
|
@@ -401,9 +417,9 @@ class SyncGenericRotatingClient(BaseGenericRotatingClient[T]):
|
|
|
401
417
|
final_tokens = chunk_tokens
|
|
402
418
|
yield chunk
|
|
403
419
|
except Exception as e:
|
|
404
|
-
if
|
|
420
|
+
if is_dead_key_error(e):
|
|
405
421
|
logger.warning(
|
|
406
|
-
"402/
|
|
422
|
+
"Dead key (401/402/403) hit during streaming for %s on key ...%s. Marking key dead.",
|
|
407
423
|
model_id, get_key_suffix(key_usage.api_key)
|
|
408
424
|
)
|
|
409
425
|
key_usage.mark_dead()
|
|
@@ -417,7 +433,7 @@ class SyncGenericRotatingClient(BaseGenericRotatingClient[T]):
|
|
|
417
433
|
self.manager.force_rotate_index()
|
|
418
434
|
raise
|
|
419
435
|
finally:
|
|
420
|
-
self._record_usage(key_usage, model_id, final_tokens)
|
|
436
|
+
self._record_usage(key_usage, model_id, final_tokens, track_usage)
|
|
421
437
|
|
|
422
438
|
|
|
423
439
|
class SyncGenericProxyHelper:
|
|
@@ -442,13 +458,16 @@ class AsyncGenericRotatingClient(BaseGenericRotatingClient[T]):
|
|
|
442
458
|
|
|
443
459
|
async def _execute(self, path: List[str], args: tuple, kwargs: dict) -> Any:
|
|
444
460
|
"""Execute a method call with key rotation (async)."""
|
|
461
|
+
# Per-call tracking override - popped before the kwargs reach the real
|
|
462
|
+
# client, which may reject unknown keyword arguments.
|
|
463
|
+
track_usage = kwargs.pop("track_usage", None)
|
|
445
464
|
model_id = self._get_model_id(kwargs)
|
|
446
465
|
limits = self.limit_resolver(model_id, None)
|
|
447
466
|
|
|
448
467
|
for attempt in range(self.config.max_retries + 1):
|
|
449
|
-
key_usage = self.manager.
|
|
468
|
+
key_usage = await self.manager.acquire_key_async(model_id, limits, self.config.estimated_tokens)
|
|
450
469
|
if not key_usage:
|
|
451
|
-
raise RuntimeError(
|
|
470
|
+
raise RuntimeError(self.manager.unavailable_message(model_id))
|
|
452
471
|
|
|
453
472
|
# Create backoff for temporary rate limits
|
|
454
473
|
temp_backoff = ExponentialBackoff(BackoffConfig(
|
|
@@ -470,9 +489,11 @@ class AsyncGenericRotatingClient(BaseGenericRotatingClient[T]):
|
|
|
470
489
|
|
|
471
490
|
# Handle async streaming responses
|
|
472
491
|
if hasattr(result, "__aiter__"):
|
|
473
|
-
return self._wrap_stream(result, key_usage, model_id)
|
|
492
|
+
return self._wrap_stream(result, key_usage, model_id, track_usage)
|
|
474
493
|
|
|
475
|
-
self._record_usage(
|
|
494
|
+
self._record_usage(
|
|
495
|
+
key_usage, model_id, self._extract_usage(result), track_usage
|
|
496
|
+
)
|
|
476
497
|
return result
|
|
477
498
|
|
|
478
499
|
except Exception as e:
|
|
@@ -487,11 +508,11 @@ class AsyncGenericRotatingClient(BaseGenericRotatingClient[T]):
|
|
|
487
508
|
await asyncio.sleep(delay)
|
|
488
509
|
continue # Retry with SAME key
|
|
489
510
|
|
|
490
|
-
#
|
|
511
|
+
# Dead key (401/403 revoked/denied, 402 credits exhausted) - key is dead
|
|
491
512
|
# for this process, rotate to next key
|
|
492
|
-
if
|
|
513
|
+
if is_dead_key_error(e) and attempt < self.config.max_retries:
|
|
493
514
|
logger.warning(
|
|
494
|
-
"402/
|
|
515
|
+
"Dead key (401/402/403) hit for %s on key ...%s. Marking key dead and rotating. (Attempt %d/%d)",
|
|
495
516
|
model_id, get_key_suffix(key_usage.api_key),
|
|
496
517
|
attempt + 1, self.config.max_retries + 1
|
|
497
518
|
)
|
|
@@ -512,7 +533,7 @@ class AsyncGenericRotatingClient(BaseGenericRotatingClient[T]):
|
|
|
512
533
|
await asyncio.sleep(KEY_ROTATION_DELAY_SECONDS)
|
|
513
534
|
break # Break inner loop, continue outer loop with new key
|
|
514
535
|
|
|
515
|
-
self._record_usage(key_usage, model_id, 0)
|
|
536
|
+
self._record_usage(key_usage, model_id, 0, track_usage)
|
|
516
537
|
raise
|
|
517
538
|
else:
|
|
518
539
|
# Inner loop exhausted without success - continue to next key
|
|
@@ -528,7 +549,11 @@ class AsyncGenericRotatingClient(BaseGenericRotatingClient[T]):
|
|
|
528
549
|
raise RuntimeError(f"All retry attempts exhausted for {model_id}")
|
|
529
550
|
|
|
530
551
|
async def _wrap_stream(
|
|
531
|
-
self,
|
|
552
|
+
self,
|
|
553
|
+
generator: AsyncGenerator,
|
|
554
|
+
key_usage: KeyUsage,
|
|
555
|
+
model_id: str,
|
|
556
|
+
track_usage: Optional[bool] = None,
|
|
532
557
|
) -> AsyncGenerator:
|
|
533
558
|
"""Wrap an async streaming response to track usage and handle errors."""
|
|
534
559
|
final_tokens = 0
|
|
@@ -540,9 +565,9 @@ class AsyncGenericRotatingClient(BaseGenericRotatingClient[T]):
|
|
|
540
565
|
final_tokens = chunk_tokens
|
|
541
566
|
yield chunk
|
|
542
567
|
except Exception as e:
|
|
543
|
-
if
|
|
568
|
+
if is_dead_key_error(e):
|
|
544
569
|
logger.warning(
|
|
545
|
-
"402/
|
|
570
|
+
"Dead key (401/402/403) hit during streaming for %s on key ...%s. Marking key dead.",
|
|
546
571
|
model_id, get_key_suffix(key_usage.api_key)
|
|
547
572
|
)
|
|
548
573
|
key_usage.mark_dead()
|
|
@@ -556,7 +581,7 @@ class AsyncGenericRotatingClient(BaseGenericRotatingClient[T]):
|
|
|
556
581
|
self.manager.force_rotate_index()
|
|
557
582
|
raise
|
|
558
583
|
finally:
|
|
559
|
-
self._record_usage(key_usage, model_id, final_tokens)
|
|
584
|
+
self._record_usage(key_usage, model_id, final_tokens, track_usage)
|
|
560
585
|
|
|
561
586
|
|
|
562
587
|
class AsyncGenericProxyHelper:
|
|
@@ -14,7 +14,7 @@ from ..config.constants import (
|
|
|
14
14
|
from ..core.utils import (
|
|
15
15
|
is_rate_limit_error,
|
|
16
16
|
is_temporary_rate_limit_error,
|
|
17
|
-
|
|
17
|
+
is_dead_key_error,
|
|
18
18
|
get_key_suffix,
|
|
19
19
|
)
|
|
20
20
|
from ..core.backoff import ExponentialBackoff, BackoffConfig
|
|
@@ -97,12 +97,19 @@ class BaseRotatingClient:
|
|
|
97
97
|
if self.base_url:
|
|
98
98
|
self.client_kwargs['base_url'] = self.base_url
|
|
99
99
|
|
|
100
|
-
def _record_usage(
|
|
100
|
+
def _record_usage(
|
|
101
|
+
self,
|
|
102
|
+
key_usage: KeyUsage,
|
|
103
|
+
model_id: str,
|
|
104
|
+
actual_tokens: int,
|
|
105
|
+
track_usage: Optional[bool] = None,
|
|
106
|
+
) -> None:
|
|
101
107
|
self.manager.record_usage(
|
|
102
108
|
key_obj=key_usage,
|
|
103
109
|
model_id=model_id,
|
|
104
110
|
actual_tokens=actual_tokens,
|
|
105
|
-
estimated_tokens=self.estimated_tokens
|
|
111
|
+
estimated_tokens=self.estimated_tokens,
|
|
112
|
+
track_usage=track_usage
|
|
106
113
|
)
|
|
107
114
|
|
|
108
115
|
def _extract_usage(self, response: Any) -> int:
|
|
@@ -123,6 +130,9 @@ class RotatingOpenAIClient(BaseRotatingClient):
|
|
|
123
130
|
return SyncProxyHelper(self, [name])
|
|
124
131
|
|
|
125
132
|
def _execute(self, path: List[str], args, kwargs):
|
|
133
|
+
# Per-call tracking override - popped before the kwargs reach the real
|
|
134
|
+
# OpenAI client, which rejects unknown keyword arguments.
|
|
135
|
+
track_usage = kwargs.pop('track_usage', None)
|
|
126
136
|
model_id = kwargs.get('model', self.default_model)
|
|
127
137
|
if 'model' not in kwargs:
|
|
128
138
|
kwargs['model'] = model_id
|
|
@@ -132,9 +142,9 @@ class RotatingOpenAIClient(BaseRotatingClient):
|
|
|
132
142
|
kwargs['stream_options'] = {"include_usage": True}
|
|
133
143
|
|
|
134
144
|
for attempt in range(self.max_retries + 1):
|
|
135
|
-
key_usage = self.manager.
|
|
145
|
+
key_usage = self.manager.acquire_key(model_id, limits, self.estimated_tokens)
|
|
136
146
|
if not key_usage:
|
|
137
|
-
raise RuntimeError(
|
|
147
|
+
raise RuntimeError(self.manager.unavailable_message(model_id))
|
|
138
148
|
|
|
139
149
|
# Create backoff for temporary rate limits
|
|
140
150
|
temp_backoff = ExponentialBackoff(BackoffConfig(
|
|
@@ -154,9 +164,11 @@ class RotatingOpenAIClient(BaseRotatingClient):
|
|
|
154
164
|
result = target(*args, **kwargs)
|
|
155
165
|
|
|
156
166
|
if kwargs.get('stream', False):
|
|
157
|
-
return self._wrap_stream(result, key_usage, model_id)
|
|
167
|
+
return self._wrap_stream(result, key_usage, model_id, track_usage)
|
|
158
168
|
|
|
159
|
-
self._record_usage(
|
|
169
|
+
self._record_usage(
|
|
170
|
+
key_usage, model_id, self._extract_usage(result), track_usage
|
|
171
|
+
)
|
|
160
172
|
return result
|
|
161
173
|
|
|
162
174
|
except Exception as e:
|
|
@@ -171,11 +183,11 @@ class RotatingOpenAIClient(BaseRotatingClient):
|
|
|
171
183
|
time.sleep(delay)
|
|
172
184
|
continue # Retry with SAME key
|
|
173
185
|
|
|
174
|
-
#
|
|
186
|
+
# Dead key (401/403 revoked/denied, 402 credits exhausted) - key is dead
|
|
175
187
|
# for this process, rotate to next key
|
|
176
|
-
if
|
|
188
|
+
if is_dead_key_error(e) and attempt < self.max_retries:
|
|
177
189
|
logger.warning(
|
|
178
|
-
"402/
|
|
190
|
+
"Dead key (401/402/403) hit for %s on key ...%s. Marking key dead and rotating. (Attempt %d/%d)",
|
|
179
191
|
model_id, get_key_suffix(key_usage.api_key), attempt + 1, self.max_retries + 1
|
|
180
192
|
)
|
|
181
193
|
key_usage.mark_dead()
|
|
@@ -194,7 +206,7 @@ class RotatingOpenAIClient(BaseRotatingClient):
|
|
|
194
206
|
time.sleep(KEY_ROTATION_DELAY_SECONDS)
|
|
195
207
|
break # Break inner loop, continue outer loop with new key
|
|
196
208
|
|
|
197
|
-
self._record_usage(key_usage, model_id, 0)
|
|
209
|
+
self._record_usage(key_usage, model_id, 0, track_usage)
|
|
198
210
|
raise
|
|
199
211
|
else:
|
|
200
212
|
# Inner loop exhausted without success - continue to next key
|
|
@@ -209,7 +221,13 @@ class RotatingOpenAIClient(BaseRotatingClient):
|
|
|
209
221
|
|
|
210
222
|
raise RuntimeError(f"All retry attempts exhausted for {model_id}")
|
|
211
223
|
|
|
212
|
-
def _wrap_stream(
|
|
224
|
+
def _wrap_stream(
|
|
225
|
+
self,
|
|
226
|
+
generator: Generator,
|
|
227
|
+
key_usage: KeyUsage,
|
|
228
|
+
model_id: str,
|
|
229
|
+
track_usage: Optional[bool] = None,
|
|
230
|
+
):
|
|
213
231
|
final_tokens = 0
|
|
214
232
|
try:
|
|
215
233
|
for chunk in generator:
|
|
@@ -217,9 +235,9 @@ class RotatingOpenAIClient(BaseRotatingClient):
|
|
|
217
235
|
final_tokens = chunk.usage.total_tokens
|
|
218
236
|
yield chunk
|
|
219
237
|
except Exception as e:
|
|
220
|
-
if
|
|
238
|
+
if is_dead_key_error(e):
|
|
221
239
|
logger.warning(
|
|
222
|
-
"402/
|
|
240
|
+
"Dead key (401/402/403) hit during streaming for %s on key ...%s. Marking key dead.",
|
|
223
241
|
model_id, get_key_suffix(key_usage.api_key)
|
|
224
242
|
)
|
|
225
243
|
key_usage.mark_dead()
|
|
@@ -233,7 +251,7 @@ class RotatingOpenAIClient(BaseRotatingClient):
|
|
|
233
251
|
self.manager.force_rotate_index()
|
|
234
252
|
raise
|
|
235
253
|
finally:
|
|
236
|
-
self._record_usage(key_usage, model_id, final_tokens)
|
|
254
|
+
self._record_usage(key_usage, model_id, final_tokens, track_usage)
|
|
237
255
|
|
|
238
256
|
class SyncProxyHelper:
|
|
239
257
|
def __init__(self, client: RotatingOpenAIClient, path: List[str]):
|
|
@@ -257,6 +275,9 @@ class RotatingAsyncOpenAIClient(BaseRotatingClient):
|
|
|
257
275
|
return AsyncProxyHelper(self, [name])
|
|
258
276
|
|
|
259
277
|
async def _execute(self, path: List[str], args, kwargs: dict):
|
|
278
|
+
# Per-call tracking override - popped before the kwargs reach the real
|
|
279
|
+
# OpenAI client, which rejects unknown keyword arguments.
|
|
280
|
+
track_usage = kwargs.pop('track_usage', None)
|
|
260
281
|
model_id = kwargs.get('model', self.default_model)
|
|
261
282
|
if 'model' not in kwargs:
|
|
262
283
|
kwargs['model'] = model_id
|
|
@@ -266,9 +287,9 @@ class RotatingAsyncOpenAIClient(BaseRotatingClient):
|
|
|
266
287
|
kwargs['stream_options'] = {"include_usage": True}
|
|
267
288
|
|
|
268
289
|
for attempt in range(self.max_retries + 1):
|
|
269
|
-
key_usage = self.manager.
|
|
290
|
+
key_usage = await self.manager.acquire_key_async(model_id, limits, self.estimated_tokens)
|
|
270
291
|
if not key_usage:
|
|
271
|
-
raise RuntimeError(
|
|
292
|
+
raise RuntimeError(self.manager.unavailable_message(model_id))
|
|
272
293
|
|
|
273
294
|
# Create backoff for temporary rate limits
|
|
274
295
|
temp_backoff = ExponentialBackoff(BackoffConfig(
|
|
@@ -288,9 +309,11 @@ class RotatingAsyncOpenAIClient(BaseRotatingClient):
|
|
|
288
309
|
result = await target(*args, **kwargs)
|
|
289
310
|
|
|
290
311
|
if kwargs.get('stream', False):
|
|
291
|
-
return self._wrap_stream(result, key_usage, model_id)
|
|
312
|
+
return self._wrap_stream(result, key_usage, model_id, track_usage)
|
|
292
313
|
|
|
293
|
-
self._record_usage(
|
|
314
|
+
self._record_usage(
|
|
315
|
+
key_usage, model_id, self._extract_usage(result), track_usage
|
|
316
|
+
)
|
|
294
317
|
return result
|
|
295
318
|
|
|
296
319
|
except Exception as e:
|
|
@@ -305,11 +328,11 @@ class RotatingAsyncOpenAIClient(BaseRotatingClient):
|
|
|
305
328
|
await asyncio.sleep(delay)
|
|
306
329
|
continue # Retry with SAME key
|
|
307
330
|
|
|
308
|
-
#
|
|
331
|
+
# Dead key (401/403 revoked/denied, 402 credits exhausted) - key is dead
|
|
309
332
|
# for this process, rotate to next key
|
|
310
|
-
if
|
|
333
|
+
if is_dead_key_error(e) and attempt < self.max_retries:
|
|
311
334
|
logger.warning(
|
|
312
|
-
"402/
|
|
335
|
+
"Dead key (401/402/403) hit for %s on key ...%s. Marking key dead and rotating. (Attempt %d/%d)",
|
|
313
336
|
model_id, get_key_suffix(key_usage.api_key), attempt + 1, self.max_retries + 1
|
|
314
337
|
)
|
|
315
338
|
key_usage.mark_dead()
|
|
@@ -328,7 +351,7 @@ class RotatingAsyncOpenAIClient(BaseRotatingClient):
|
|
|
328
351
|
await asyncio.sleep(KEY_ROTATION_DELAY_SECONDS)
|
|
329
352
|
break # Break inner loop, continue outer loop with new key
|
|
330
353
|
|
|
331
|
-
self._record_usage(key_usage, model_id, 0)
|
|
354
|
+
self._record_usage(key_usage, model_id, 0, track_usage)
|
|
332
355
|
raise
|
|
333
356
|
else:
|
|
334
357
|
# Inner loop exhausted without success - continue to next key
|
|
@@ -343,7 +366,13 @@ class RotatingAsyncOpenAIClient(BaseRotatingClient):
|
|
|
343
366
|
|
|
344
367
|
raise RuntimeError(f"All retry attempts exhausted for {model_id}")
|
|
345
368
|
|
|
346
|
-
async def _wrap_stream(
|
|
369
|
+
async def _wrap_stream(
|
|
370
|
+
self,
|
|
371
|
+
generator: AsyncGenerator,
|
|
372
|
+
key_usage: KeyUsage,
|
|
373
|
+
model_id: str,
|
|
374
|
+
track_usage: Optional[bool] = None,
|
|
375
|
+
):
|
|
347
376
|
final_tokens = 0
|
|
348
377
|
try:
|
|
349
378
|
async for chunk in generator:
|
|
@@ -351,9 +380,9 @@ class RotatingAsyncOpenAIClient(BaseRotatingClient):
|
|
|
351
380
|
final_tokens = chunk.usage.total_tokens
|
|
352
381
|
yield chunk
|
|
353
382
|
except Exception as e:
|
|
354
|
-
if
|
|
383
|
+
if is_dead_key_error(e):
|
|
355
384
|
logger.warning(
|
|
356
|
-
"402/
|
|
385
|
+
"Dead key (401/402/403) hit during streaming for %s on key ...%s. Marking key dead.",
|
|
357
386
|
model_id, get_key_suffix(key_usage.api_key)
|
|
358
387
|
)
|
|
359
388
|
key_usage.mark_dead()
|
|
@@ -367,7 +396,7 @@ class RotatingAsyncOpenAIClient(BaseRotatingClient):
|
|
|
367
396
|
self.manager.force_rotate_index()
|
|
368
397
|
raise
|
|
369
398
|
finally:
|
|
370
|
-
self._record_usage(key_usage, model_id, final_tokens)
|
|
399
|
+
self._record_usage(key_usage, model_id, final_tokens, track_usage)
|
|
371
400
|
|
|
372
401
|
class AsyncProxyHelper:
|
|
373
402
|
def __init__(self, client: RotatingAsyncOpenAIClient, path: List[str]):
|
|
@@ -31,5 +31,12 @@ TEMP_RATE_LIMIT_MULTIPLIER = 2.0
|
|
|
31
31
|
# Key rotation delay (seconds to wait after rotating to a new key)
|
|
32
32
|
KEY_ROTATION_DELAY_SECONDS = 0.5
|
|
33
33
|
|
|
34
|
+
# How long acquire_key() polls for a live key to leave its rate window /
|
|
35
|
+
# cooldown before giving up. A single production key at its per-minute cap
|
|
36
|
+
# is the common case (Cohere, 2026-09-07): the right answer is to wait for
|
|
37
|
+
# the window, not to abort the caller's whole run.
|
|
38
|
+
KEY_WAIT_MAX_SECONDS = 120.0
|
|
39
|
+
KEY_WAIT_POLL_SECONDS = 1.0
|
|
40
|
+
|
|
34
41
|
# Usage logging batch size
|
|
35
42
|
USAGE_LOG_BATCH_SIZE = 50
|
|
@@ -14,6 +14,7 @@ from .utils import (
|
|
|
14
14
|
get_key_suffix,
|
|
15
15
|
is_rate_limit_error,
|
|
16
16
|
is_auth_error,
|
|
17
|
+
is_dead_key_error,
|
|
17
18
|
validate_api_key,
|
|
18
19
|
)
|
|
19
20
|
from .backoff import ExponentialBackoff, BackoffConfig
|
|
@@ -33,6 +34,7 @@ __all__ = [
|
|
|
33
34
|
"get_key_suffix",
|
|
34
35
|
"is_rate_limit_error",
|
|
35
36
|
"is_auth_error",
|
|
37
|
+
"is_dead_key_error",
|
|
36
38
|
"validate_api_key",
|
|
37
39
|
# Backoff
|
|
38
40
|
"ExponentialBackoff",
|
|
@@ -241,6 +241,29 @@ def is_auth_error(e: Exception) -> bool:
|
|
|
241
241
|
return False
|
|
242
242
|
|
|
243
243
|
|
|
244
|
+
# A key that answers 401/402/403 is unusable for good (revoked, unpaid,
|
|
245
|
+
# or its project denied) — no cooldown will fix it, so rotate off it.
|
|
246
|
+
DEAD_KEY_STATUS_CODES = AUTH_STATUS_CODES | PAYMENT_REQUIRED_STATUS_CODES
|
|
247
|
+
|
|
248
|
+
|
|
249
|
+
def is_dead_key_error(e: Exception) -> bool:
|
|
250
|
+
"""
|
|
251
|
+
Detect errors that mean THIS KEY is dead, not merely busy: HTTP 401/403
|
|
252
|
+
(invalid, revoked, or project-denied key — e.g. Gemini's
|
|
253
|
+
"Your project has been denied access") and 402 (credits exhausted).
|
|
254
|
+
|
|
255
|
+
A status_code, when present, is authoritative: 429s are never dead keys
|
|
256
|
+
whatever their message says. Without one, fall back to the string
|
|
257
|
+
heuristics of is_payment_required_error / is_auth_error.
|
|
258
|
+
"""
|
|
259
|
+
status = getattr(e, "status_code", None)
|
|
260
|
+
if status in DEAD_KEY_STATUS_CODES:
|
|
261
|
+
return True
|
|
262
|
+
if status == 429:
|
|
263
|
+
return False
|
|
264
|
+
return is_payment_required_error(e) or is_auth_error(e)
|
|
265
|
+
|
|
266
|
+
|
|
244
267
|
def validate_api_key(api_key: str) -> bool:
|
|
245
268
|
"""
|
|
246
269
|
Basic validation of API key format.
|
|
@@ -15,7 +15,7 @@ from ..config.constants import (
|
|
|
15
15
|
from ..core.utils import (
|
|
16
16
|
is_rate_limit_error,
|
|
17
17
|
is_temporary_rate_limit_error,
|
|
18
|
-
|
|
18
|
+
is_dead_key_error,
|
|
19
19
|
get_key_suffix,
|
|
20
20
|
)
|
|
21
21
|
from ..core.backoff import ExponentialBackoff, BackoffConfig
|
|
@@ -115,7 +115,12 @@ class RotatingCredentialsMixin:
|
|
|
115
115
|
def _get_retry_limit(self) -> int:
|
|
116
116
|
return min(self._max_retries, len(self.wrapper.manager.keys) - 1)
|
|
117
117
|
|
|
118
|
-
def _record_usage(
|
|
118
|
+
def _record_usage(
|
|
119
|
+
self,
|
|
120
|
+
key_obj: KeyUsage,
|
|
121
|
+
response: Optional["ModelResponse"],
|
|
122
|
+
track_usage: Optional[bool] = None,
|
|
123
|
+
):
|
|
119
124
|
"""
|
|
120
125
|
Extracts usage from the response and reports it to the manager.
|
|
121
126
|
Falls back to estimated_tokens if response is None or usage is missing.
|
|
@@ -136,7 +141,8 @@ class RotatingCredentialsMixin:
|
|
|
136
141
|
key_obj=key_obj,
|
|
137
142
|
model_id=self.model_id,
|
|
138
143
|
actual_tokens=actual_tokens,
|
|
139
|
-
estimated_tokens=self._estimated_tokens
|
|
144
|
+
estimated_tokens=self._estimated_tokens,
|
|
145
|
+
track_usage=track_usage
|
|
140
146
|
)
|
|
141
147
|
|
|
142
148
|
def _create_temp_backoff(self) -> ExponentialBackoff:
|
|
@@ -148,6 +154,9 @@ class RotatingCredentialsMixin:
|
|
|
148
154
|
))
|
|
149
155
|
|
|
150
156
|
def invoke(self, *args, **kwargs) -> "ModelResponse":
|
|
157
|
+
# Per-call tracking override - popped before the kwargs reach the
|
|
158
|
+
# underlying model, which does not know this keyword.
|
|
159
|
+
track_usage = kwargs.pop("track_usage", None)
|
|
151
160
|
limit = self._get_retry_limit()
|
|
152
161
|
|
|
153
162
|
for attempt in range(limit + 1):
|
|
@@ -157,7 +166,7 @@ class RotatingCredentialsMixin:
|
|
|
157
166
|
for temp_attempt in range(TEMP_RATE_LIMIT_MAX_RETRIES + 1):
|
|
158
167
|
try:
|
|
159
168
|
response = super().invoke(*args, **kwargs)
|
|
160
|
-
self._record_usage(key_usage, response)
|
|
169
|
+
self._record_usage(key_usage, response, track_usage)
|
|
161
170
|
return response
|
|
162
171
|
except Exception as e:
|
|
163
172
|
# Check for temporary rate limit first - retry with SAME key
|
|
@@ -171,11 +180,11 @@ class RotatingCredentialsMixin:
|
|
|
171
180
|
time.sleep(delay)
|
|
172
181
|
continue # Retry with SAME key
|
|
173
182
|
|
|
174
|
-
#
|
|
183
|
+
# Dead key (401/403 revoked/denied, 402 credits exhausted) - key is dead
|
|
175
184
|
# for this process, rotate to next key
|
|
176
|
-
if
|
|
185
|
+
if is_dead_key_error(e) and attempt < limit:
|
|
177
186
|
self.logger.warning(
|
|
178
|
-
"402
|
|
187
|
+
"Dead key (401/402/403) on key %s (Sync) [%s]. Marking key dead and rotating (%d/%d).",
|
|
179
188
|
get_key_suffix(self.api_key), self.model_id, attempt + 1, limit
|
|
180
189
|
)
|
|
181
190
|
key_usage.mark_dead()
|
|
@@ -205,6 +214,9 @@ class RotatingCredentialsMixin:
|
|
|
205
214
|
raise
|
|
206
215
|
|
|
207
216
|
async def ainvoke(self, *args, **kwargs) -> "ModelResponse":
|
|
217
|
+
# Per-call tracking override - popped before the kwargs reach the
|
|
218
|
+
# underlying model, which does not know this keyword.
|
|
219
|
+
track_usage = kwargs.pop("track_usage", None)
|
|
208
220
|
limit = self._get_retry_limit()
|
|
209
221
|
|
|
210
222
|
for attempt in range(limit + 1):
|
|
@@ -214,7 +226,7 @@ class RotatingCredentialsMixin:
|
|
|
214
226
|
for temp_attempt in range(TEMP_RATE_LIMIT_MAX_RETRIES + 1):
|
|
215
227
|
try:
|
|
216
228
|
response = await super().ainvoke(*args, **kwargs)
|
|
217
|
-
self._record_usage(key_usage, response)
|
|
229
|
+
self._record_usage(key_usage, response, track_usage)
|
|
218
230
|
return response
|
|
219
231
|
except Exception as e:
|
|
220
232
|
# Check for temporary rate limit first - retry with SAME key
|
|
@@ -228,11 +240,11 @@ class RotatingCredentialsMixin:
|
|
|
228
240
|
await asyncio.sleep(delay)
|
|
229
241
|
continue # Retry with SAME key
|
|
230
242
|
|
|
231
|
-
#
|
|
243
|
+
# Dead key (401/403 revoked/denied, 402 credits exhausted) - key is dead
|
|
232
244
|
# for this process, rotate to next key
|
|
233
|
-
if
|
|
245
|
+
if is_dead_key_error(e) and attempt < limit:
|
|
234
246
|
self.logger.warning(
|
|
235
|
-
"402
|
|
247
|
+
"Dead key (401/402/403) on key %s (Async) [%s]. Marking key dead and rotating (%d/%d).",
|
|
236
248
|
get_key_suffix(self.api_key), self.model_id, attempt + 1, limit
|
|
237
249
|
)
|
|
238
250
|
key_usage.mark_dead()
|
|
@@ -262,6 +274,9 @@ class RotatingCredentialsMixin:
|
|
|
262
274
|
raise
|
|
263
275
|
|
|
264
276
|
def invoke_stream(self, *args, **kwargs) -> Iterator["ModelResponse"]:
|
|
277
|
+
# Per-call tracking override - popped before the kwargs reach the
|
|
278
|
+
# underlying model, which does not know this keyword.
|
|
279
|
+
track_usage = kwargs.pop("track_usage", None)
|
|
265
280
|
limit = self._get_retry_limit()
|
|
266
281
|
|
|
267
282
|
for attempt in range(limit + 1):
|
|
@@ -280,10 +295,10 @@ class RotatingCredentialsMixin:
|
|
|
280
295
|
yield chunk
|
|
281
296
|
if final_usage:
|
|
282
297
|
dummy_response = _UsageResponse(final_usage)
|
|
283
|
-
self._record_usage(key_usage, dummy_response)
|
|
298
|
+
self._record_usage(key_usage, dummy_response, track_usage)
|
|
284
299
|
else:
|
|
285
300
|
# If the stream didn't return usage data, fallback to estimation
|
|
286
|
-
self._record_usage(key_usage, None)
|
|
301
|
+
self._record_usage(key_usage, None, track_usage)
|
|
287
302
|
|
|
288
303
|
return
|
|
289
304
|
except Exception as e:
|
|
@@ -300,9 +315,9 @@ class RotatingCredentialsMixin:
|
|
|
300
315
|
time.sleep(delay)
|
|
301
316
|
continue
|
|
302
317
|
|
|
303
|
-
if
|
|
318
|
+
if is_dead_key_error(e) and attempt < limit:
|
|
304
319
|
self.logger.warning(
|
|
305
|
-
"402
|
|
320
|
+
"Dead key (401/402/403) on key %s (Sync Stream) [%s]. Marking key dead and rotating (%d/%d).",
|
|
306
321
|
get_key_suffix(self.api_key), self.model_id, attempt + 1, limit
|
|
307
322
|
)
|
|
308
323
|
key_usage.mark_dead()
|
|
@@ -326,6 +341,9 @@ class RotatingCredentialsMixin:
|
|
|
326
341
|
raise
|
|
327
342
|
|
|
328
343
|
async def ainvoke_stream(self, *args, **kwargs) -> AsyncIterator["ModelResponse"]:
|
|
344
|
+
# Per-call tracking override - popped before the kwargs reach the
|
|
345
|
+
# underlying model, which does not know this keyword.
|
|
346
|
+
track_usage = kwargs.pop("track_usage", None)
|
|
329
347
|
limit = self._get_retry_limit()
|
|
330
348
|
|
|
331
349
|
for attempt in range(limit + 1):
|
|
@@ -347,9 +365,9 @@ class RotatingCredentialsMixin:
|
|
|
347
365
|
# Stream completed successfully. Record usage.
|
|
348
366
|
if final_usage:
|
|
349
367
|
dummy_response = _UsageResponse(final_usage)
|
|
350
|
-
self._record_usage(key_usage, dummy_response)
|
|
368
|
+
self._record_usage(key_usage, dummy_response, track_usage)
|
|
351
369
|
else:
|
|
352
|
-
self._record_usage(key_usage, None)
|
|
370
|
+
self._record_usage(key_usage, None, track_usage)
|
|
353
371
|
|
|
354
372
|
return
|
|
355
373
|
except Exception as e:
|
|
@@ -366,9 +384,9 @@ class RotatingCredentialsMixin:
|
|
|
366
384
|
await asyncio.sleep(delay)
|
|
367
385
|
continue
|
|
368
386
|
|
|
369
|
-
if
|
|
387
|
+
if is_dead_key_error(e) and attempt < limit:
|
|
370
388
|
self.logger.warning(
|
|
371
|
-
"402
|
|
389
|
+
"Dead key (401/402/403) on key %s (Async Stream) [%s]. Marking key dead and rotating (%d/%d).",
|
|
372
390
|
get_key_suffix(self.api_key), self.model_id, attempt + 1, limit
|
|
373
391
|
)
|
|
374
392
|
key_usage.mark_dead()
|