Mirrowel commited on
Commit
818ba14
·
1 Parent(s): 543f871

feat(antigravity): ✨ add credential availability tracking for models

Browse files

Add 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)}/{len(credentials_for_provider)}"
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)}/{len(credentials_for_provider)}"
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.