gemini_key_manager.py

"""
gemini_key_manager.py
----------------------
Round-robin API key rotation untuk Gemini via LiteLLM, dengan fallback
otomatis dan cooldown khusus untuk limit RPD (Requests Per Day).

CARA PAKAI DI HUGGING FACE SPACES:
- Buka Settings > Variables and secrets pada Space Anda.
- Tambahkan secret dengan Name & Value:
    GEMINI_API_KEY_1  = AIza...key1
    GEMINI_API_KEY_2  = AIza...key2
    GEMINI_API_KEY_3  = AIza...key3
    GEMINI_API_KEY_4  = AIza...key4
    GEMINI_API_KEY_5  = AIza...key5
  (boleh 1 sampai berapa pun key, penomoran harus berurutan mulai dari 1,
   tidak boleh ada nomor yang loncat/kosong di tengah).
- Modul ini otomatis mendeteksi semua GEMINI_API_KEY_<n> yang ada.
"""

import os
import re
import threading
import logging
from datetime import datetime, timedelta, timezone
from typing import Dict, List, Optional

logger = logging.getLogger(__name__)

# Jam reset kuota harian Gemini (RPD) menurut Google: 08:05 UTC
RPD_RESET_HOUR_UTC = 8
RPD_RESET_MINUTE_UTC = 5


class GeminiKeyManager:
    def __init__(self, prefix: str = "GEMINI_API_KEY_"):
        keys: List[str] = []
        i = 1
        while True:
            val = os.environ.get(f"{prefix}{i}")
            if not val:
                break
            keys.append(val.strip())
            i += 1

        # fallback: kalau tidak ada GEMINI_API_KEY_1, coba GEMINI_API_KEY tunggal
        if not keys:
            single = os.environ.get("GEMINI_API_KEY")
            if single:
                keys = [single.strip()]

        if not keys:
            raise RuntimeError(
                "Tidak ada Gemini API key ditemukan. "
                "Set GEMINI_API_KEY_1, GEMINI_API_KEY_2, ... di Variables and secrets."
            )

        self.keys = keys
        self.n = len(keys)
        self._lock = threading.Lock()
        self._idx = 0
        # key -> datetime (UTC) sampai kapan key ini "istirahat" karena RPD habis
        self._cooldown_until: Dict[str, Optional[datetime]] = {k: None for k in keys}

        logger.info(f"GeminiKeyManager aktif dengan {self.n} API key.")

    # ------------------------------------------------------------------
    # Helper waktu reset
    # ------------------------------------------------------------------
    def _next_rpd_reset(self) -> datetime:
        now = datetime.now(timezone.utc)
        reset = now.replace(
            hour=RPD_RESET_HOUR_UTC, minute=RPD_RESET_MINUTE_UTC,
            second=0, microsecond=0,
        )
        if now >= reset:
            reset += timedelta(days=1)
        return reset

    def _is_available(self, key: str) -> bool:
        until = self._cooldown_until.get(key)
        if until is None:
            return True
        if datetime.now(timezone.utc) >= until:
            self._cooldown_until[key] = None  # cooldown habis, key aktif lagi
            return True
        return False

    # ------------------------------------------------------------------
    # API utama
    # ------------------------------------------------------------------
    def get_rotation_order(self) -> List[str]:
        """
        Kembalikan daftar key sesuai urutan giliran, dimulai dari key
        berikutnya dalam rotasi. Pointer selalu maju setiap kali method
        ini dipanggil (baik request berhasil maupun gagal), sehingga
        beban merata di semua user/panggilan.
        """
        with self._lock:
            order = [self.keys[(self._idx + i) % self.n] for i in range(self.n)]
            self._idx = (self._idx + 1) % self.n
        return order

    def mark_rpd_exhausted(self, key: str) -> None:
        """Tandai key kena limit harian (RPD) -> istirahat sampai reset berikutnya."""
        with self._lock:
            reset_at = self._next_rpd_reset()
            self._cooldown_until[key] = reset_at
        logger.warning(f"Key ...{key[-6:]} kena limit RPD, istirahat sampai {reset_at.isoformat()}")

    def status(self) -> Dict[str, str]:
        now = datetime.now(timezone.utc)
        out = {}
        for k in self.keys:
            until = self._cooldown_until.get(k)
            if until and now < until:
                out[f"...{k[-6:]}"] = f"cooldown until {until.isoformat()}"
            else:
                out[f"...{k[-6:]}"] = "available"
        return out


# ------------------------------------------------------------------
# Deteksi jenis error dari Gemini / LiteLLM
# ------------------------------------------------------------------
# Catatan: format pesan error Google untuk RPD biasanya mengandung kata
# "PerDay" pada quotaId, mis. "GenerateRequestsPerDayPerProjectPerModel-FreeTier".
# Untuk RPM/TPM mengandung "PerMinute". Kalau tidak yakin (ambigu), kita
# anggap BUKAN rpd (lebih aman -> key tidak dikunci seharian secara keliru).
RPD_PATTERNS = [r"PerDay", r"per[_ ]?day", r"daily limit", r"requests per day"]
RATE_LIMIT_PATTERNS = [r"RESOURCE_EXHAUSTED", r"429", r"rate limit", r"quota"]


def classify_gemini_error(exc: Exception) -> str:
    """Return 'rpd', 'rpm', atau 'other'."""
    msg = str(exc)
    is_rate_limit = any(re.search(p, msg, re.IGNORECASE) for p in RATE_LIMIT_PATTERNS)
    if not is_rate_limit:
        return "other"
    if any(re.search(p, msg, re.IGNORECASE) for p in RPD_PATTERNS):
        return "rpd"
    return "rpm"