""" Gemini API Pool Manager Alternative to Groq when Groq is blocked in user's region Uses Gemini 2.x models for agent reasoning """ import asyncio import time import logging from collections import deque from typing import List, Dict from dataclasses import dataclass logger = logging.getLogger(__name__) @dataclass class GeminiKeyUsage: """Track usage for a Gemini API key""" api_key: str calls_today: int = 0 last_reset_date: str = "" total_calls: int = 0 total_failures: int = 0 is_blocked: bool = False blocked_until: float = 0 class GeminiAgentPool: """ Gemini-based agent pool (alternative to Groq) Uses Gemini 2.x models for fast agent reasoning Free tier: 1500 requests/day for gemini-2.0-flash """ # Available models - PRIORITIZE HIGHER RPM MODELS FOR FREE TIER # Free tier limits (per key): # gemini-2.5-flash: 10 RPM # gemini-2.0-flash: 15 RPM # gemini-2.5-flash-lite: 15 RPM # gemini-2.0-flash-lite: 30 RPM <-- BEST for high concurrency REASONING_MODEL = "gemini-2.0-flash" # 15 RPM, good reasoning FAST_MODEL = "gemini-2.0-flash-lite" # 30 RPM, fast routing LITE_MODEL = "gemini-2.0-flash-lite" # 30 RPM # Fallback model chain (highest RPM first) MODEL_FALLBACKS = [ "gemini-2.0-flash-lite", # 30 RPM "gemini-2.5-flash-lite", # 15 RPM "gemini-2.0-flash", # 15 RPM "gemini-2.5-flash", # 10 RPM "gemini-flash-latest", ] def __init__(self, api_keys: List[str]): """Initialize with Gemini API keys""" if not api_keys: raise ValueError("At least one Gemini API key is required") self.api_keys = api_keys self.key_usage = {key: GeminiKeyUsage(api_key=key) for key in api_keys} self.current_index = 0 self._lock = asyncio.Lock() self._configured_keys = set() self._working_model = None # Per-key rate limit tracking (timestamps of recent calls) self._key_call_times = {key: deque() for key in api_keys} # Concurrency limiter - avoid bursting too many concurrent calls # With 3 keys × 30 RPM = 90 calls/min capacity, allow 6 concurrent max_concurrent = min(6, len(api_keys) * 2) self._semaphore = asyncio.Semaphore(max_concurrent) logger.info( f"GeminiAgentPool initialized: {len(api_keys)} keys, " f"max concurrent={max_concurrent}" ) def _configure_key(self, api_key: str): """Configure Gemini SDK with this API key""" if api_key not in self._configured_keys: import google.generativeai as genai genai.configure(api_key=api_key) self._configured_keys.add(api_key) async def get_available_key(self) -> str: """Get next available API key (round-robin)""" async with self._lock: attempts = 0 max_attempts = len(self.api_keys) * 2 while attempts < max_attempts: key = self.api_keys[self.current_index % len(self.api_keys)] self.current_index += 1 usage = self.key_usage[key] # Reset block if expired now = time.time() if usage.is_blocked and now > usage.blocked_until: usage.is_blocked = False if not usage.is_blocked: usage.total_calls += 1 return key attempts += 1 await asyncio.sleep(2) return await self.get_available_key() async def invoke( self, prompt: str, model: str = "llama-3.1-8b-instant", # Ignored, mapped to Gemini temperature: float = 0.2, max_tokens: int = 1000, system_prompt: str = None, agent_name: str = "unknown", conversation_id: str = None, ) -> str: """ Invoke Gemini API (compatible with GroqAPIPool interface) """ from .llm_logger import llm_logger # Map Groq model names to Gemini models gemini_model = self._map_model(model) api_key = await self.get_available_key() start_time = time.time() try: import google.generativeai as genai self._configure_key(api_key) generation_config = { "temperature": temperature, "max_output_tokens": max_tokens, } models_to_try = [self._working_model] if self._working_model else [] models_to_try.extend([m for m in self.MODEL_FALLBACKS if m != self._working_model]) models_to_try = [m for m in models_to_try if m] last_error = None for model_name in models_to_try: try: gen_model = genai.GenerativeModel( model_name=model_name, system_instruction=system_prompt if system_prompt else None, generation_config=generation_config, ) response = await asyncio.to_thread( gen_model.generate_content, prompt, ) self._working_model = model_name response_text = response.text if hasattr(response, "text") else str(response) duration_ms = (time.time() - start_time) * 1000 # Log this call llm_logger.log_call( agent_name=agent_name, provider="gemini", model=model_name, prompt=prompt, response=response_text, tokens_input=len(prompt) // 4, tokens_output=len(response_text) // 4, duration_ms=duration_ms, success=True, conversation_id=conversation_id, metadata={"system_prompt_chars": len(system_prompt) if system_prompt else 0}, ) return response_text except Exception as e: last_error = e err_str = str(e).lower() # If model not found, try next if "404" in err_str or "not found" in err_str: continue # If rate limited, mark key as blocked and retry with next key if "429" in err_str or "rate" in err_str or "quota" in err_str: usage = self.key_usage[api_key] usage.is_blocked = True # Parse retry_delay from error if possible retry_seconds = 30 try: if "retry_delay" in err_str: import re m = re.search(r"seconds:\s*(\d+)", err_str) if m: retry_seconds = int(m.group(1)) except Exception: pass usage.blocked_until = time.time() + retry_seconds logger.warning( f"Gemini key {api_key[:10]}... rate limited (retry in {retry_seconds}s). " f"Trying next key from pool ({len(self.api_keys)} total)." ) # Try next key (recursion will pick a different key) return await self.invoke( prompt, model, temperature, max_tokens, system_prompt, agent_name, conversation_id ) # Other errors - try next model continue raise Exception(f"All Gemini models failed. Last error: {last_error}") except Exception as e: usage = self.key_usage[api_key] usage.total_failures += 1 logger.error(f"Gemini invocation failed: {e}") raise def _map_model(self, groq_model: str) -> str: """Map Groq-style model name to Gemini equivalent""" # 70B = needs reasoning -> use gemini-2.5-flash # 8B = fast routing -> use gemini-2.0-flash if "70b" in groq_model.lower() or "versatile" in groq_model.lower(): return self.REASONING_MODEL elif "8b" in groq_model.lower() or "instant" in groq_model.lower(): return self.FAST_MODEL return self.FAST_MODEL # Default def get_status(self) -> Dict: """Get current pool status""" return { "provider": "gemini", "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()), "working_model": self._working_model, }