app/key_manager.py

import time
from typing import List, Optional
from threading import Lock

class APIKeyManager:
    def __init__(self, keys: List[str]):
        self.lock = Lock()
        self.keys = keys
        self.usage = {
            key: {
                "count": 0,
                "minute_window": [],
                "total_requests": 0
            } for key in keys
        }
        self.rate_limit = 20  # per minute
        self.total_limit = 50

    def _cleanup_old_timestamps(self, timestamps: List[float]) -> List[float]:
        one_minute_ago = time.time() - 60
        return [t for t in timestamps if t > one_minute_ago]

    def get_available_key(self) -> Optional[str]:
        with self.lock:
            for key in self.keys:
                record = self.usage[key]
                record["minute_window"] = self._cleanup_old_timestamps(record["minute_window"])

                if record["total_requests"] >= self.total_limit:
                    continue

                if len(record["minute_window"]) >= self.rate_limit:
                    continue

                return key
        return None

    def record_usage(self, key: str):
        with self.lock:
            self.usage[key]["minute_window"].append(time.time())
            self.usage[key]["total_requests"] += 1