Update app/core/model_loader.py

app/core/model_loader.py

@@ -7,15 +7,23 @@ import random
 from datetime import datetime
 from dotenv import load_dotenv
 from groq import Groq
+import google.generativeai as genai
 
 load_dotenv()
 logger = logging.getLogger(__name__)
 
 STATS_FILE = "usage_stats.json"
 
-# Fixed: Updated to active Groq production model IDs to prevent 404 errors
-DEFAULT_MODEL = "openai/gpt-oss-20b"
-FALLBACK_MODEL = "openai/gpt-oss-120b"
+# ====================== CONFIG ======================
+# Gemini (Primary)
+DEFAULT_GEMINI_MODEL = "gemini-2.5-pro"          # Higher quality. Use "gemini-2.5-flash" if you want faster + cheaper
+FALLBACK_GEMINI_MODEL = "gemini-2.0-flash"
+
+# Groq (Fallback)
+DEFAULT_GROQ_MODEL = "llama-3.3-70b-versatile"
+FALLBACK_GROQ_MODEL = "openai/gpt-oss-120b"
+# ====================================================
+
 
 class LLMSingleton:
     _instance = None
@@ -33,25 +41,36 @@ class LLMSingleton:
         if self._instance is not None:
             raise Exception("Singleton instance already exists!")
 
-        # 1. Load and sanitize API Key
-        raw_key = os.getenv("GROQ_API_KEY", "")
-        self.api_key = raw_key.strip().strip('"').strip("'") if raw_key else ""
+        # ---------- Gemini ----------
+        self.gemini_key = os.getenv("GOOGLE_API_KEY") or os.getenv("GEMINI_API_KEY") or ""
+        self.gemini_key = self.gemini_key.strip().strip('"').strip("'")
+        self.gemini_model_name = os.getenv("GEMINI_MODEL", DEFAULT_GEMINI_MODEL).strip() or DEFAULT_GEMINI_MODEL
 
-        # Debug check for API key loading status
-        if self.api_key:
-            logger.info(f"✅ Loaded GROQ_API_KEY successfully! (Length: {len(self.api_key)}, Starts with: {self.api_key[:4]})")
-            self.client = Groq(api_key=self.api_key)
+        if self.gemini_key:
+            try:
+                genai.configure(api_key=self.gemini_key)
+                self.gemini_model = genai.GenerativeModel(self.gemini_model_name)
+                logger.info(f"✅ Gemini initialized → {self.gemini_model_name}")
+            except Exception as e:
+                logger.error(f"❌ Failed to init Gemini: {e}")
+                self.gemini_model = None
         else:
-            logger.error("❌ GROQ_API_KEY is empty or missing from environment variables!")
-            self.client = None
+            logger.warning("⚠️ GOOGLE_API_KEY / GEMINI_API_KEY not found")
+            self.gemini_model = None
 
-        # 2. Load and sanitize Model Name
-        raw_model = os.getenv("GROQ_MODEL", DEFAULT_MODEL)
-        clean_model = raw_model.strip().strip('"').strip("'") if raw_model else ""
-        self.model_name = clean_model if clean_model else DEFAULT_MODEL
-        logger.info(f"🔑 Groq Client initialized with model target: {self.model_name}")
+        # ---------- Groq (Fallback) ----------
+        raw_key = os.getenv("GROQ_API_KEY", "")
+        self.groq_key = raw_key.strip().strip('"').strip("'") if raw_key else ""
+        self.groq_model_name = os.getenv("GROQ_MODEL", DEFAULT_GROQ_MODEL).strip() or DEFAULT_GROQ_MODEL
 
-        # 3. Threading locks and usage stats setup
+        if self.groq_key:
+            self.groq_client = Groq(api_key=self.groq_key)
+            logger.info(f"✅ Groq initialized → {self.groq_model_name}")
+        else:
+            logger.warning("⚠️ GROQ_API_KEY not found")
+            self.groq_client = None
+
+        # ---------- Stats & Rate limiting ----------
         self._stats_lock = threading.Lock()
         self._rpm_lock = threading.Lock()
         self.stats = self._load_stats()
@@ -60,6 +79,9 @@ class LLMSingleton:
         self.minute_window_start = time.time()
         self.requests_this_minute = 0
 
+    # ------------------------------------------------------------------
+    # Stats helpers
+    # ------------------------------------------------------------------
     def _load_stats(self):
         default_stats = {
             "total_requests": 0,
@@ -121,12 +143,6 @@ class LLMSingleton:
         stats["remaining_rpm"] = max(0, self.rpm_limit - requests_this_minute)
         return stats
 
-    def track_local_usage(self, input_chars: int = 0):
-        with self._stats_lock:
-            self.stats["local_model_requests"] += 1
-            self.stats["input_tokens"] += input_chars // 4
-            self._save_stats()
-
     def _reserve_request_slot(self) -> bool:
         with self._stats_lock:
             if self.stats["daily_requests_count"] >= 1000:
@@ -139,133 +155,130 @@ class LLMSingleton:
             self.requests_this_minute += 1
         return True
 
+    # ------------------------------------------------------------------
+    # Core generation methods
+    # ------------------------------------------------------------------
+    def _call_gemini(self, prompt: str, system_prompt: str, max_tokens: int = 2048, json_mode: bool = False) -> str:
+        if not self.gemini_model:
+            raise RuntimeError("Gemini not available")
+
+        full_prompt = f"{system_prompt}\n\n{prompt}"
+
+        generation_config = {
+            "max_output_tokens": max_tokens,
+            "temperature": 0.2 if json_mode else 0.3,
+        }
+        if json_mode:
+            generation_config["response_mime_type"] = "application/json"
+
+        response = self.gemini_model.generate_content(
+            full_prompt,
+            generation_config=generation_config,
+        )
+        return (response.text or "").strip()
+
+    def _call_groq(self, prompt: str, system_prompt: str, max_tokens: int = 2048, json_mode: bool = False) -> str:
+        if not self.groq_client:
+            raise RuntimeError("Groq not available")
+
+        messages = [
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": prompt},
+        ]
+
+        kwargs = {
+            "model": self.groq_model_name,
+            "messages": messages,
+            "max_tokens": max_tokens,
+            "temperature": 0.2 if json_mode else 0.3,
+        }
+        if json_mode:
+            kwargs["response_format"] = {"type": "json_object"}
+
+        completion = self.groq_client.chat.completions.create(**kwargs)
+        return (completion.choices[0].message.content or "").strip()
+
     def generate(self, prompt: str, max_tokens: int = 2048) -> str:
+        """Used for structured JSON responses (code review etc.)"""
         self._check_daily_reset()
         self._check_rpm_window()
 
-        if not self.api_key or not self.client:
-            logger.error("Cannot generate: Missing GROQ_API_KEY")
-            raise RuntimeError("MISSING_API_KEY")
-
         if not self._reserve_request_slot():
-            logger.error("❌ Daily Quota Exceeded. Request blocked.")
             raise RuntimeError("QUOTA_EXCEEDED")
 
-        logger.info(f"🤖 Generating with Groq ({self.model_name}). Prompt start: {prompt[:50]}...")
-
-        retries = 0
-        max_retries = 3
-        base_delay = 2
-        active_model = self.model_name
+        system_prompt = (
+            "You are a senior Android code reviewer. "
+            "You MUST return a valid JSON object matching the requested schema strictly."
+        )
 
-        while retries <= max_retries:
+        # Try Gemini first
+        if self.gemini_model:
             try:
+                logger.info(f"🤖 Generating with Gemini ({self.gemini_model_name})")
+                result = self._call_gemini(prompt, system_prompt, max_tokens, json_mode=True)
                 with self._stats_lock:
-                    self.stats["input_tokens"] += len(prompt) // 4
+                    self.stats["successful_requests"] += 1
+                    self.stats["output_tokens"] += len(result) // 4
                     self._save_stats()
+                return result
+            except Exception as e:
+                logger.warning(f"Gemini failed → falling back to Groq: {e}")
 
-                chat_completion = self.client.chat.completions.create(
-                    model=active_model,
-                    messages=[
-                        {
-                            "role": "system",
-                            "content": "You are a senior Android code reviewer. You MUST return a valid JSON object matching the requested schema strictly.",
-                        },
-                        {"role": "user", "content": prompt},
-                    ],
-                    response_format={"type": "json_object"},
-                    max_tokens=max_tokens,
-                    temperature=0.2,
-                )
-
-                response_text = chat_completion.choices[0].message.content or ""
-
+        # Fallback to Groq
+        if self.groq_client:
+            try:
+                logger.info(f"🤖 Generating with Groq ({self.groq_model_name})")
+                result = self._call_groq(prompt, system_prompt, max_tokens, json_mode=True)
                 with self._stats_lock:
                     self.stats["successful_requests"] += 1
-                    if response_text:
-                        self.stats["output_tokens"] += len(response_text) // 4
+                    self.stats["output_tokens"] += len(result) // 4
                     self._save_stats()
-
-                return response_text.strip()
-
+                return result
             except Exception as e:
-                error_str = str(e)
-
-                # Automatic Model Fallback handling if a 404/model_not_found occurs
-                if "404" in error_str or "model_not_found" in error_str.lower():
-                    if active_model != FALLBACK_MODEL:
-                        logger.warning(
-                            f"⚠️ Model target '{active_model}' rejected (404). Falling back to '{FALLBACK_MODEL}'..."
-                        )
-                        active_model = FALLBACK_MODEL
-                        continue
-
-                # Rate Limit handling (HTTP 429)
-                if "429" in error_str or "rate_limit_exceeded" in error_str.lower():
-                    with self._stats_lock:
-                        self.stats["rate_limit_hits"] += 1
-                        self._save_stats()
-
-                    wait_time = (base_delay * (2**retries)) + random.uniform(0.5, 1.5)
-                    logger.warning(
-                        f"⚠️ Groq rate limit hit. Retrying in {wait_time:.2f}s... (Attempt {retries + 1}/{max_retries})"
-                    )
-                    time.sleep(wait_time)
-                    retries += 1
-                else:
-                    with self._stats_lock:
-                        self.stats["errors"] += 1
-                        self._save_stats()
-                    logger.error(f"Groq generation failed: {e}")
-                    raise RuntimeError(f"GENERATION_FAILED: {e}")
-
-        with self._stats_lock:
-            self.stats["errors"] += 1
-            self._save_stats()
+                logger.error(f"Groq also failed: {e}")
+                raise RuntimeError(f"GENERATION_FAILED: {e}")
 
-        logger.error("❌ Max retries reached. Request failed.")
-        raise RuntimeError("RATE_LIMIT_EXCEEDED")
+        raise RuntimeError("No LLM provider available")
 
     def generate_text(self, prompt: str) -> str:
+        """Used for normal chat / explanations"""
         self._check_daily_reset()
         self._check_rpm_window()
 
-        if not self.api_key or not self.client:
-            return "Error: Missing API Key."
-
         if not self._reserve_request_slot():
             return "Error: Daily Quota Exceeded."
 
-        try:
-            with self._stats_lock:
-                self.stats["input_tokens"] += len(prompt) // 4
-                self._save_stats()
+        system_prompt = "You are GitGud AI, an expert software architect."
 
-            chat_completion = self.client.chat.completions.create(
-                model=self.model_name,
-                messages=[
-                    {
-                        "role": "system",
-                        "content": "You are GitGud AI, an expert software architect.",
-                    },
-                    {"role": "user", "content": prompt},
-                ],
-                max_tokens=2048,
-                temperature=0.3,
-            )
-
-            response_text = chat_completion.choices[0].message.content or ""
-
-            with self._stats_lock:
-                self.stats["successful_requests"] += 1
-                self.stats["output_tokens"] += len(response_text) // 4
-                self._save_stats()
+        # Try Gemini first
+        if self.gemini_model:
+            try:
+                logger.info(f"🤖 Chat with Gemini ({self.gemini_model_name})")
+                result = self._call_gemini(prompt, system_prompt, max_tokens=2048, json_mode=False)
+                with self._stats_lock:
+                    self.stats["successful_requests"] += 1
+                    self.stats["output_tokens"] += len(result) // 4
+                    self._save_stats()
+                return result
+            except Exception as e:
+                logger.warning(f"Gemini chat failed → falling back to Groq: {e}")
 
-            return response_text
+        # Fallback to Groq
+        if self.groq_client:
+            try:
+                logger.info(f"🤖 Chat with Groq ({self.groq_model_name})")
+                result = self._call_groq(prompt, system_prompt, max_tokens=2048, json_mode=False)
+                with self._stats_lock:
+                    self.stats["successful_requests"] += 1
+                    self.stats["output_tokens"] += len(result) // 4
+                    self._save_stats()
+                return result
+            except Exception as e:
+                logger.error(f"Groq chat also failed: {e}")
+                return f"Error generating content: {str(e)}"
+
+        return "Error: No LLM provider available (check API keys)."
 
-        except Exception as e:
-            logger.error(f"Groq text generation failed: {e}")
-            return f"Error generating content: {str(e)}"
 
 # Export global instance
 llm_engine = LLMSingleton.get_instance()