Rahul-Samedavar/CodeCanvas
0
1"""2AI model management for the AI Web Visualization Generator.3 4This module handles interactions with multiple AI providers (Gemini and Requesty),5including API key rotation, fallback mechanisms, and streaming response generation.6"""7 8import itertools9from typing import AsyncGenerator, List, Optional10 11import openai12import google.generativeai as genai13from fastapi import Request14 15from config import AppSettings16 17 18class GeminiModelManager:19 """20 Manages Gemini API interactions with support for multiple API keys.21 22 This class handles API key rotation and provides streaming content generation23 using Google's Generative AI models. If one API key fails, it automatically24 tries the next available key.25 26 Attributes:27 model_name: Name of the Gemini model to use28 keys: List of Gemini API keys for rotation29 generation_config: Configuration for text generation30 """31 32 def __init__(self, config: AppSettings):33 """34 Initialize the Gemini model manager.35 36 Args:37 config: Application settings containing API keys and model name38 39 Raises:40 ValueError: If no Gemini API keys are provided41 """42 self.model_name = config.primary_model_name43 self.keys = config.gemini_api_keys_list44 45 if not self.keys:46 raise ValueError(47 "GeminiModelManager initialized but no GEMINI_API_KEYS were provided."48 )49 50 self.key_cycler = itertools.cycle(self.keys)51 self.generation_config = genai.GenerationConfig(52 temperature=0.7,53 top_p=1,54 top_k=155 )56 57 print(58 f"Gemini Manager initialized for model: {self.model_name} "59 f"with {len(self.keys)} API key(s)."60 )61 62 async def generate_content_streaming_with_key(63 self, 64 prompt: str, 65 api_key: str66 ) -> AsyncGenerator[str, None]:67 """68 Generate streaming content using a specific API key.69 70 Args:71 prompt: The prompt to send to the model72 api_key: The Gemini API key to use73 74 Yields:75 str: Chunks of generated text76 77 Raises:78 Exception: If the API call fails79 """80 genai.configure(api_key=api_key)81 model = genai.GenerativeModel(self.model_name)82 83 print(84 f"[Gemini] Attempting generation with model {self.model_name} "85 f"using key ending in ...{api_key[-4:]}"86 )87 88 stream = await model.generate_content_async(89 prompt,90 stream=True,91 generation_config=self.generation_config92 )93 94 async for chunk in stream:95 if chunk.text:96 yield chunk.text97 98 print(99 f"[Gemini] Successfully generated response "100 f"with key ending in ...{api_key[-4:]}"101 )102 103 async def try_all_keys_streaming(104 self, 105 prompt: str106 ) -> AsyncGenerator[str, None]:107 """108 Attempt to generate content using all available API keys.109 110 Tries each API key in sequence until one succeeds. If all keys fail,111 raises the last exception encountered.112 113 Args:114 prompt: The prompt to send to the model115 116 Yields:117 str: Chunks of generated text118 119 Raises:120 Exception: If all API keys fail121 """122 last_exception = None123 124 for i, api_key in enumerate(self.keys):125 try:126 print(127 f"[Gemini] Trying key {i+1}/{len(self.keys)} "128 f"(ending in ...{api_key[-4:]})"129 )130 131 async for chunk in self.generate_content_streaming_with_key(132 prompt, 133 api_key134 ):135 yield chunk136 137 # If we got here, generation succeeded138 return139 140 except Exception as e:141 last_exception = e142 print(f"[Gemini] Key {i+1}/{len(self.keys)} failed: {str(e)}")143 continue144 145 # All keys failed146 print(147 f"[Gemini] All {len(self.keys)} API keys failed. "148 f"Last error: {last_exception}"149 )150 raise last_exception or Exception("All Gemini API keys failed")151 152 153class RequestyModelManager:154 """155 Manages Requesty API interactions as a fallback provider.156 157 This class provides a fallback mechanism when Gemini API is unavailable158 or all API keys have been exhausted. It uses the Requesty router service159 with OpenAI-compatible API.160 161 Attributes:162 model_name: Name of the model to use via Requesty163 client: Async OpenAI client configured for Requesty164 """165 166 def __init__(self, config: AppSettings):167 """168 Initialize the Requesty model manager.169 170 Args:171 config: Application settings containing API key and site info172 """173 self.model_name = config.fallback_model_name174 175 # Build headers for Requesty service176 headers = {177 "HTTP-Referer": config.requesty_site_url,178 "X-Title": config.requesty_site_name179 }180 181 # Filter out empty header values182 headers = {k: v for k, v in headers.items() if v}183 184 self.client = openai.AsyncOpenAI(185 api_key=config.requesty_api_key,186 base_url="https://router.requesty.ai/v1",187 default_headers=headers188 )189 190 print(f"Requesty Fallback Manager initialized for model: {self.model_name}")191 192 async def generate_content_streaming(193 self, 194 prompt: str, 195 request: Request196 ) -> AsyncGenerator[str, None]:197 """198 Generate streaming content using Requesty API.199 200 Args:201 prompt: The prompt to send to the model202 request: FastAPI request object (for disconnect detection)203 204 Yields:205 str: Chunks of generated text206 """207 print(f"[Requesty] Attempting generation with model {self.model_name}")208 209 stream = await self.client.chat.completions.create(210 model=self.model_name,211 messages=[{"role": "user", "content": prompt}],212 stream=True213 )214 215 async for chunk in stream:216 # Check if client disconnected217 if await request.is_disconnected():218 print("[Requesty] Client disconnected. Cancelling stream.")219 break220 221 content = chunk.choices[0].delta.content222 if content:223 yield content224 225 print("[Requesty] Successfully generated response.")226 227 228class MultiModelManager:229 """230 Orchestrates multiple AI model providers with fallback logic.231 232 This class manages the coordination between primary (Gemini) and fallback233 (Requesty) AI providers. It attempts to use Gemini first with all available234 API keys, then falls back to Requesty if all Gemini attempts fail.235 236 Attributes:237 gemini_manager: Optional Gemini model manager238 requesty_manager: Requesty model manager (always available)239 """240 241 def __init__(self, config: AppSettings):242 """243 Initialize the multi-model manager.244 245 Args:246 config: Application settings for all providers247 """248 self.gemini_manager: Optional[GeminiModelManager] = None249 250 # Try to initialize Gemini manager if API keys are available251 if config.gemini_api_keys_list:252 try:253 self.gemini_manager = GeminiModelManager(config)254 except ValueError as e:255 print(f"Warning: Could not initialize Gemini Manager. {e}")256 257 # Always initialize Requesty as fallback258 self.requesty_manager = RequestyModelManager(config)259 260 async def generate_content_streaming(261 self, 262 prompt: str, 263 request: Request264 ) -> AsyncGenerator[str, None]:265 """266 Generate content with automatic fallback between providers.267 268 First attempts to use Gemini with all available API keys. If all fail269 or Gemini is not available, falls back to Requesty. Sends a special270 [STREAM_RESTART] marker when switching providers.271 272 Args:273 prompt: The prompt to send to the model274 request: FastAPI request object275 276 Yields:277 str: Chunks of generated text278 """279 # Try Gemini first if available280 if self.gemini_manager:281 try:282 print(283 "[Orchestrator] Attempting generation with all "284 "available Gemini keys..."285 )286 287 async for chunk in self.gemini_manager.try_all_keys_streaming(prompt):288 yield chunk289 290 # If we got here, generation succeeded291 return292 293 except Exception as e:294 print(295 f"[Orchestrator] All Gemini keys failed with final error: {e}. "296 f"Falling back to Requesty."297 )298 # Send restart marker to indicate provider switch299 yield "[STREAM_RESTART]\n"300 301 # Use Requesty fallback302 print("[Orchestrator] Using fallback: Requesty.")303 async for chunk in self.requesty_manager.generate_content_streaming(304 prompt, 305 request306 ):307 yield chunk