import os import asyncio import aiohttp import backoff import logging from typing import Optional, Dict, Any, List from dataclasses import dataclass from utils.error_handling import APIError, RateLimitError, error_handler @dataclass class APIConfig: api_key: str base_url: str = "https://generativelanguage.googleapis.com" timeout: int = 120 max_retries: int = 3 class GoogleAIStudioClient: def __init__(self, api_key: Optional[str] = None): self.api_key = api_key or os.getenv("GOOGLE_API_KEY") or os.getenv("GEMINI_API_KEY") if not self.api_key: raise ValueError("Google AI Studio API key required. Set GOOGLE_API_KEY or GEMINI_API_KEY environment variable.") self.config = APIConfig(api_key=self.api_key) self.session: Optional[aiohttp.ClientSession] = None self.logger = logging.getLogger(__name__) async def __aenter__(self): timeout = aiohttp.ClientTimeout(total=self.config.timeout) self.session = aiohttp.ClientSession( timeout=timeout, headers={ "Content-Type": "application/json", "x-goog-api-key": self.api_key } ) return self async def __aexit__(self, exc_type, exc_val, exc_tb): if self.session: await self.session.close() @backoff.on_exception( backoff.expo, (aiohttp.ClientError, asyncio.TimeoutError), max_tries=3, max_time=300 ) async def generate_content(self, prompt: str, model: str = "gemini-2.5-flash", **kwargs) -> Dict[str, Any]: """Generate content using Google AI models""" endpoint = f"/v1beta/models/{model}:generateContent" payload = { "contents": [{ "role": "user", "parts": [{"text": prompt}] }], "generationConfig": { "temperature": kwargs.get("temperature", 0.7), "maxOutputTokens": kwargs.get("max_tokens", 8192), "topP": kwargs.get("top_p", 0.95), "topK": kwargs.get("top_k", 40) } } # Add safety settings payload["safetySettings"] = [ { "category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_MEDIUM_AND_ABOVE" }, { "category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_MEDIUM_AND_ABOVE" }, { "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "BLOCK_MEDIUM_AND_ABOVE" }, { "category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_MEDIUM_AND_ABOVE" } ] return await self._make_request(endpoint, payload) async def generate_image(self, prompt: str, model: str = "gemini-2.5-flash", **kwargs) -> Dict[str, Any]: """Generate image using Google AI models""" # For image generation, we modify the prompt to be more specific image_prompt = f"Generate an image: {prompt}" # Add image-specific parameters enhanced_kwargs = { **kwargs, "temperature": kwargs.get("temperature", 0.7), "max_tokens": kwargs.get("max_tokens", 1000) } return await self.generate_content(image_prompt, model, **enhanced_kwargs) async def generate_video(self, prompt: str, model: str = "gemini-2.5-flash", duration: float = 5.0, **kwargs) -> Dict[str, Any]: """Generate video description using available models""" video_prompt = f"Create a detailed video description for: {prompt}. Duration: {duration} seconds." enhanced_kwargs = { **kwargs, "temperature": kwargs.get("temperature", 0.8), "max_tokens": kwargs.get("max_tokens", 2000) } return await self.generate_content(video_prompt, model, **enhanced_kwargs) @error_handler() async def _make_request(self, endpoint: str, payload: Dict[str, Any]) -> Dict[str, Any]: """Make HTTP request to Google AI Studio API""" if not self.session: raise RuntimeError("Client session not initialized. Use 'async with' context manager.") url = f"{self.config.base_url}{endpoint}" self.logger.debug(f"Making request to: {url}") try: async with self.session.post(url, json=payload) as response: if response.status == 429: # Rate limit handling retry_after = int(response.headers.get("retry-after", "60")) error_msg = f"Rate limited. Retry after {retry_after}s" self.logger.warning(error_msg) raise RateLimitError(error_msg, retry_after) if response.status == 403: error_msg = "API key invalid or insufficient permissions" raise APIError(error_msg) if response.status == 400: error_text = await response.text() error_msg = f"Bad request: {error_text}" raise APIError(error_msg) response.raise_for_status() result = await response.json() self.logger.debug(f"API Response: {result}") return result except aiohttp.ClientError as e: raise APIError(f"Network error: {str(e)}") except asyncio.TimeoutError: raise APIError("Request timeout") def validate_model(self, model: str) -> bool: """Validate if model is supported""" supported_models = [ "gemini-2.5-flash", "gemini-2.5-pro", "gemini-1.5-flash", "gemini-1.5-pro" ] return model in supported_models def estimate_tokens(self, text: str) -> int: """Rough estimation of token count""" # Simple estimation: ~4 characters per token return len(text) // 4 class RateLimiter: """Rate limiter for Google AI Studio API calls""" def __init__(self): self.requests = {} self.model_limits = { "gemini-2.5-flash": {"rpm": 1000, "tpm": 2000000}, # Tier 1 "gemini-2.5-pro": {"rpm": 50, "tpm": 2000000}, "gemini-1.5-flash": {"rpm": 1000, "tpm": 1000000}, "gemini-1.5-pro": {"rpm": 50, "tpm": 32000} } async def acquire(self, model: str, estimated_tokens: int = 1000) -> None: """Acquire rate limit permission before making request""" if model not in self.model_limits: return # No limits defined for this model now = asyncio.get_event_loop().time() # Initialize tracking for this model if needed if model not in self.requests: self.requests[model] = [] # Clean old requests (older than 1 minute) self.requests[model] = [ req_time for req_time in self.requests[model] if now - req_time <= 60 ] limits = self.model_limits[model] # Check RPM limit if len(self.requests[model]) >= limits.get("rpm", float('inf')): wait_time = 60 - (now - self.requests[model][0]) if wait_time > 0: await asyncio.sleep(wait_time) # Record this request self.requests[model].append(now)