api_pool.py

import os
import time
from threading import Lock
from collections import defaultdict

class APIKeyPool:
    def __init__(self):
        self.keys = self._load_keys()
        self.usage = defaultdict(int)  # key -> request_count
        self.total_usage = defaultdict(int)  # key -> total requests ever
        self.last_request = defaultdict(float)  # key -> last request time
        self.lock = Lock()
        self.rate_limit = 25  # Adjusted for 12 teams
        self.window = 60  # seconds
        self.min_delay = 2.0  # Minimum 2 seconds between requests per key
        self.current_index = 0  # Strict round-robin
        self.last_reset = time.time()
        print(f"API Pool initialized with {len(self.keys)} keys")
    
    def _load_keys(self):
        keys = []
        # Load multiple keys from environment
        for i in range(1, 11):  # Support up to 10 keys
            key = os.getenv(f"GROQ_API_KEY_{i}")
            if key:
                keys.append(key)
                print(f"✅ Loaded GROQ_API_KEY_{i}")
            else:
                print(f"❌ GROQ_API_KEY_{i} not found")
        
        # Fallback to single key
        if not keys:
            main_key = os.getenv("GROQ_API_KEY")
            if main_key:
                keys.append(main_key)
                print(f"✅ Loaded fallback GROQ_API_KEY")
        
        return keys
    
    def get_available_key(self):
        with self.lock:
            current_time = time.time()
            
            # Reset counters every minute
            if current_time - self.last_reset > self.window:
                self._print_usage_stats()
                self.usage.clear()
                self.last_reset = current_time
            
            # Strict round-robin - always use next key in sequence
            key_index = self.current_index
            key = self.keys[key_index]
            
            # Check if current key is under limit AND enough time has passed
            time_since_last = current_time - self.last_request[key]
            if (self.usage[key] < self.rate_limit and time_since_last >= self.min_delay):
                self.usage[key] += 1
                self.total_usage[key] += 1
                self.last_request[key] = current_time
                # Always move to next key
                self.current_index = (self.current_index + 1) % len(self.keys)
                print(f"✅ API Key {key_index + 1}: {self.usage[key]}/{self.rate_limit} (Total: {self.total_usage[key]})")
                return key
            else:
                print(f"⏳ Key {key_index + 1} blocked: usage={self.usage[key]}/{self.rate_limit}, delay={time_since_last:.1f}s/{self.min_delay}s")
            
            # If current key exhausted, try others
            for i in range(1, len(self.keys)):
                key_index = (self.current_index + i) % len(self.keys)
                key = self.keys[key_index]
                time_since_last = current_time - self.last_request[key]
                if (self.usage[key] < self.rate_limit and time_since_last >= self.min_delay):
                    self.usage[key] += 1
                    self.total_usage[key] += 1
                    self.last_request[key] = current_time
                    self.current_index = (key_index + 1) % len(self.keys)
                    print(f"✅ API Key {key_index + 1}: {self.usage[key]}/{self.rate_limit} (Total: {self.total_usage[key]})")
                    return key
                else:
                    print(f"⏳ Key {key_index + 1} blocked: usage={self.usage[key]}/{self.rate_limit}, delay={time_since_last:.1f}s/{self.min_delay}s")
            
            print("⚠️ ALL API KEYS EXHAUSTED - Using fallback")
            return None
    
    def _print_usage_stats(self):
        print("\n📊 API Usage Distribution (Last Minute):")
        for i, key in enumerate(self.keys):
            print(f"Key {i+1}: {self.usage[key]} requests (Total: {self.total_usage[key]})")
        print()

# Global instance
api_pool = APIKeyPool()