""" Groq API Key Pool Manager Manages multiple Groq API keys with round-robin and rate limiting """ import asyncio import time import logging from typing import List, Dict from dataclasses import dataclass, field logger = logging.getLogger(__name__) @dataclass class KeyUsage: """Track usage for a single API key""" api_key: str calls_this_minute: int = 0 last_reset_time: float = field(default_factory=time.time) total_calls: int = 0 total_failures: int = 0 is_blocked: bool = False blocked_until: float = 0 class GroqAPIPool: """ Round-robin API key pool with rate limiting Free tier: 30 requests/minute per key """ RATE_LIMIT_PER_MINUTE = 30 RATE_WINDOW_SECONDS = 60 def __init__(self, api_keys: List[str]): """ Initialize pool with API keys Args: api_keys: List of Groq API keys """ if not api_keys: raise ValueError("At least one Groq API key is required") self.api_keys = api_keys self.key_usage = {key: KeyUsage(api_key=key) for key in api_keys} self.current_index = 0 self._lock = asyncio.Lock() logger.info(f"GroqAPIPool initialized with {len(api_keys)} keys") logger.info(f"Total capacity: {len(api_keys) * self.RATE_LIMIT_PER_MINUTE} requests/minute") async def get_available_key(self) -> str: """ Get next available API key (rate-limit aware) Uses round-robin across keys, skipping rate-limited ones """ async with self._lock: now = time.time() attempts = 0 max_attempts = len(self.api_keys) * 2 while attempts < max_attempts: # Get next key in round-robin order key = self.api_keys[self.current_index % len(self.api_keys)] self.current_index += 1 usage = self.key_usage[key] # Reset counter if minute window has passed if now - usage.last_reset_time > self.RATE_WINDOW_SECONDS: usage.calls_this_minute = 0 usage.last_reset_time = now # Unblock if block period passed if usage.is_blocked and now > usage.blocked_until: usage.is_blocked = False logger.info(f"Key {key[:10]}... unblocked") # Check if usable if ( not usage.is_blocked and usage.calls_this_minute < self.RATE_LIMIT_PER_MINUTE ): usage.calls_this_minute += 1 usage.total_calls += 1 return key attempts += 1 # All keys exhausted, wait and retry wait_time = self._calculate_wait_time() logger.warning( f"All Groq keys rate-limited. Waiting {wait_time:.1f}s for capacity..." ) await asyncio.sleep(wait_time) return await self.get_available_key() def _calculate_wait_time(self) -> float: """Calculate minimum wait time until any key becomes available""" now = time.time() min_wait = float("inf") for usage in self.key_usage.values(): if usage.is_blocked: wait = max(0, usage.blocked_until - now) else: wait = max(0, self.RATE_WINDOW_SECONDS - (now - usage.last_reset_time)) min_wait = min(min_wait, wait) return max(1.0, min(min_wait, 60.0)) async def invoke( self, prompt: str, model: str = "llama-3.1-8b-instant", temperature: float = 0.2, max_tokens: int = 1000, system_prompt: str = None, ) -> str: """ Invoke Groq API with automatic key rotation Args: prompt: User prompt model: Groq model name (llama-3.1-8b-instant or llama-3.1-70b-versatile) temperature: Sampling temperature max_tokens: Max output tokens system_prompt: Optional system prompt Returns: Response text from Groq """ # Get available key api_key = await self.get_available_key() try: # Import here to avoid hard dependency from groq import AsyncGroq client = AsyncGroq(api_key=api_key) # Build messages messages = [] if system_prompt: messages.append({"role": "system", "content": system_prompt}) messages.append({"role": "user", "content": prompt}) # Make API call response = await client.chat.completions.create( model=model, messages=messages, temperature=temperature, max_tokens=max_tokens, ) # Extract response text return response.choices[0].message.content except Exception as e: # Mark key as failed usage = self.key_usage[api_key] usage.total_failures += 1 error_str = str(e).lower() # Detect rate limit errors if "rate" in error_str or "limit" in error_str or "429" in error_str: usage.is_blocked = True usage.blocked_until = time.time() + 60 logger.warning(f"Key {api_key[:10]}... blocked for 60s (rate limit)") # Retry with different key if usage.total_failures < 3: logger.warning(f"Retrying with different key: {str(e)[:100]}") return await self.invoke(prompt, model, temperature, max_tokens, system_prompt) raise Exception(f"All Groq keys failed: {str(e)}") async def call_async( self, model: str, messages: List[Dict], temperature: float = 0.2, max_tokens: int = 1000, ) -> str: """ Call Groq API with pre-built messages (for polishers, agents, etc.) Args: model: Groq model name messages: List of message dicts with 'role' and 'content' temperature: Sampling temperature max_tokens: Max output tokens Returns: Response text from Groq """ # Get available key api_key = await self.get_available_key() try: from groq import AsyncGroq client = AsyncGroq(api_key=api_key) # Make API call with provided messages response = await client.chat.completions.create( model=model, messages=messages, temperature=temperature, max_tokens=max_tokens, ) # Extract response text return response.choices[0].message.content except Exception as e: # Mark key as failed usage = self.key_usage[api_key] usage.total_failures += 1 error_str = str(e).lower() # Detect rate limit errors if "rate" in error_str or "limit" in error_str or "429" in error_str: usage.is_blocked = True usage.blocked_until = time.time() + 60 logger.warning(f"Key {api_key[:10]}... blocked for 60s (rate limit)") # Retry with different key if usage.total_failures < 3: logger.warning(f"Retrying with different key: {str(e)[:100]}") return await self.call_async(model, messages, temperature, max_tokens) raise Exception(f"All Groq keys failed: {str(e)}") def get_status(self) -> Dict: """Get current pool status""" now = time.time() return { "total_keys": len(self.api_keys), "active_keys": sum(1 for k in self.key_usage.values() if not k.is_blocked), "total_calls": sum(k.total_calls for k in self.key_usage.values()), "total_failures": sum(k.total_failures for k in self.key_usage.values()), "current_capacity": sum( self.RATE_LIMIT_PER_MINUTE - k.calls_this_minute for k in self.key_usage.values() if not k.is_blocked ), "keys": [ { "key": k.api_key[:10] + "...", "calls_this_minute": k.calls_this_minute, "total_calls": k.total_calls, "total_failures": k.total_failures, "is_blocked": k.is_blocked, } for k in self.key_usage.values() ], }