Spaces:
Paused
Paused
Mirrowel commited on
Commit ·
818ba14
1
Parent(s): 543f871
feat(antigravity): ✨ add credential availability tracking for models
Browse filesAdd get_available_credentials_for_model method to filter credentials not on cooldown.
Update client logging to display available credentials count alongside total count.
src/rotator_library/client.py
CHANGED
|
@@ -1033,8 +1033,17 @@ class RotatingClient:
|
|
| 1033 |
if not creds_to_try:
|
| 1034 |
break
|
| 1035 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1036 |
lib_logger.info(
|
| 1037 |
-
f"Acquiring key for model {model}. Tried keys: {len(tried_creds)}/{
|
| 1038 |
)
|
| 1039 |
max_concurrent = self.max_concurrent_requests_per_key.get(provider, 1)
|
| 1040 |
current_cred = await self.usage_manager.acquire_key(
|
|
@@ -1757,8 +1766,17 @@ class RotatingClient:
|
|
| 1757 |
)
|
| 1758 |
break
|
| 1759 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1760 |
lib_logger.info(
|
| 1761 |
-
f"Acquiring credential for model {model}. Tried credentials: {len(tried_creds)}/{
|
| 1762 |
)
|
| 1763 |
max_concurrent = self.max_concurrent_requests_per_key.get(
|
| 1764 |
provider, 1
|
|
|
|
| 1033 |
if not creds_to_try:
|
| 1034 |
break
|
| 1035 |
|
| 1036 |
+
# Get count of credentials not on cooldown for this model
|
| 1037 |
+
available_creds = (
|
| 1038 |
+
await self.usage_manager.get_available_credentials_for_model(
|
| 1039 |
+
creds_to_try, model
|
| 1040 |
+
)
|
| 1041 |
+
)
|
| 1042 |
+
available_count = len(available_creds)
|
| 1043 |
+
total_count = len(credentials_for_provider)
|
| 1044 |
+
|
| 1045 |
lib_logger.info(
|
| 1046 |
+
f"Acquiring key for model {model}. Tried keys: {len(tried_creds)}/{available_count}({total_count})"
|
| 1047 |
)
|
| 1048 |
max_concurrent = self.max_concurrent_requests_per_key.get(provider, 1)
|
| 1049 |
current_cred = await self.usage_manager.acquire_key(
|
|
|
|
| 1766 |
)
|
| 1767 |
break
|
| 1768 |
|
| 1769 |
+
# Get count of credentials not on cooldown for this model
|
| 1770 |
+
available_creds = (
|
| 1771 |
+
await self.usage_manager.get_available_credentials_for_model(
|
| 1772 |
+
creds_to_try, model
|
| 1773 |
+
)
|
| 1774 |
+
)
|
| 1775 |
+
available_count = len(available_creds)
|
| 1776 |
+
total_count = len(credentials_for_provider)
|
| 1777 |
+
|
| 1778 |
lib_logger.info(
|
| 1779 |
+
f"Acquiring credential for model {model}. Tried credentials: {len(tried_creds)}/{available_count}({total_count})"
|
| 1780 |
)
|
| 1781 |
max_concurrent = self.max_concurrent_requests_per_key.get(
|
| 1782 |
provider, 1
|
src/rotator_library/usage_manager.py
CHANGED
|
@@ -600,6 +600,43 @@ class UsageManager:
|
|
| 600 |
async with self._data_lock:
|
| 601 |
return dict(self._usage_data) if self._usage_data else {}
|
| 602 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 603 |
async def _reset_daily_stats_if_needed(self):
|
| 604 |
"""
|
| 605 |
Checks if usage stats need to be reset for any key.
|
|
|
|
| 600 |
async with self._data_lock:
|
| 601 |
return dict(self._usage_data) if self._usage_data else {}
|
| 602 |
|
| 603 |
+
async def get_available_credentials_for_model(
|
| 604 |
+
self, credentials: List[str], model: str
|
| 605 |
+
) -> List[str]:
|
| 606 |
+
"""
|
| 607 |
+
Get credentials that are not on cooldown for a specific model.
|
| 608 |
+
|
| 609 |
+
Filters out credentials where:
|
| 610 |
+
- key_cooldown_until > now (key-level cooldown)
|
| 611 |
+
- model_cooldowns[model] > now (model-specific cooldown, includes quota exhausted)
|
| 612 |
+
|
| 613 |
+
Args:
|
| 614 |
+
credentials: List of credential identifiers to check
|
| 615 |
+
model: Model name to check cooldowns for
|
| 616 |
+
|
| 617 |
+
Returns:
|
| 618 |
+
List of credentials that are available (not on cooldown) for this model
|
| 619 |
+
"""
|
| 620 |
+
await self._lazy_init()
|
| 621 |
+
now = time.time()
|
| 622 |
+
available = []
|
| 623 |
+
|
| 624 |
+
async with self._data_lock:
|
| 625 |
+
for key in credentials:
|
| 626 |
+
key_data = self._usage_data.get(key, {})
|
| 627 |
+
|
| 628 |
+
# Skip if key-level cooldown is active
|
| 629 |
+
if (key_data.get("key_cooldown_until") or 0) > now:
|
| 630 |
+
continue
|
| 631 |
+
|
| 632 |
+
# Skip if model-specific cooldown is active
|
| 633 |
+
if (key_data.get("model_cooldowns", {}).get(model) or 0) > now:
|
| 634 |
+
continue
|
| 635 |
+
|
| 636 |
+
available.append(key)
|
| 637 |
+
|
| 638 |
+
return available
|
| 639 |
+
|
| 640 |
async def _reset_daily_stats_if_needed(self):
|
| 641 |
"""
|
| 642 |
Checks if usage stats need to be reset for any key.
|