src/copilot/api_pool.py

"""
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()
            ],
        }