Update app/core/model_loader.py

app/core/model_loader.py

@@ -1,22 +1,17 @@
-# app/core/model_loader.py
 import os
 import json
 import logging
 import threading
 import time
 from datetime import datetime
-import google.generativeai as genai
-from google.generativeai.types import HarmCategory, HarmBlockThreshold
 from dotenv import load_dotenv
+from groq import Groq
 
-# Force load .env immediately upon module import
 load_dotenv()
 
 logger = logging.getLogger(__name__)
-
 STATS_FILE = "usage_stats.json"
 
-
 class LLMSingleton:
     _instance = None
     _instance_lock = threading.Lock()
@@ -33,45 +28,21 @@ class LLMSingleton:
         if self._instance is not None:
             raise Exception("Singleton instance already exists!")
 
-        self.api_key = os.getenv("GEMINI_API_KEY")
+        self.api_key = os.getenv("GROQ_API_KEY")
         if not self.api_key:
-            logger.warning(
-                "⚠️ GEMINI_API_KEY not found in environment variables. AI features will fail."
-            )
+            logger.warning("⚠️ GROQ_API_KEY not found in environment variables. AI features will fail.")
+            self.client = None
         else:
-            genai.configure(api_key=self.api_key)
-
-        # Gemini Flash Latest: Stable alias (likely 1.5 Flash)
-        # Explicitly pinning to 1.5-flash to guarantee the 1500 RPD free tier
-        self.model_name = "gemini-flash-latest"
-        self.generation_config = {
-            "temperature": 0.3,
-            "top_p": 0.95,
-            "top_k": 64,
-            "max_output_tokens": 8192,
-            "response_mime_type": "application/json",
-        }
-        self.safety_settings = {
-            HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_NONE,
-            HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_NONE,
-            HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_NONE,
-            HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_NONE,
-        }
+            self.client = Groq(api_key=self.api_key)
 
-        # Guards self.stats (read/modify/write) and the JSON file writes.
-        # Needed now that multiple review batches can call generate()
-        # concurrently from different threads.
+        self.model_name = os.getenv("GROQ_MODEL", "llama-3.1-8b-instant")
         self._stats_lock = threading.Lock()
-
-        # Guards RPM window tracking for the same reason.
         self._rpm_lock = threading.Lock()
 
-        # Load persistent stats
         self.stats = self._load_stats()
         self._check_daily_reset()
 
-        # RPM Tracking
-        self.rpm_limit = 15
+        self.rpm_limit = 30
         self.minute_window_start = time.time()
         self.requests_this_minute = 0
 
@@ -91,16 +62,12 @@ class LLMSingleton:
             try:
                 with open(STATS_FILE, "r") as f:
                     data = json.load(f)
-                    # Merge defaults for backward compatibility
                     return {**default_stats, **data}
             except Exception as e:
                 logger.error(f"Failed to load stats: {e}")
         return default_stats
 
     def _save_stats(self):
-        """
-        Caller must hold self._stats_lock before calling this.
-        """
         try:
             with open(STATS_FILE, "w") as f:
                 json.dump(self.stats, f, indent=2)
@@ -117,9 +84,6 @@ class LLMSingleton:
                 self._save_stats()
 
     def _check_rpm_window(self):
-        """
-        Resets the minute counter if 60 seconds have passed.
-        """
         with self._rpm_lock:
             now = time.time()
             if now - self.minute_window_start >= 60:
@@ -129,81 +93,55 @@ class LLMSingleton:
     def get_usage_stats(self):
         self._check_daily_reset()
         self._check_rpm_window()
-
         with self._stats_lock:
             stats = self.stats.copy()
-
         with self._rpm_lock:
             requests_this_minute = self.requests_this_minute
 
-        # Limits for Gemini 3 Flash (Free Tier)
-        daily_limit = 500
-
+        daily_limit = 1000
         stats["limits"] = {
             "requests_per_minute": self.rpm_limit,
             "requests_per_day": daily_limit,
         }
-        stats["remaining_daily_requests"] = max(
-            0, daily_limit - stats["daily_requests_count"]
-        )
+        stats["remaining_daily_requests"] = max(0, daily_limit - stats["daily_requests_count"])
         stats["remaining_rpm"] = max(0, self.rpm_limit - requests_this_minute)
-
         return stats
 
     def track_local_usage(self, input_chars: int = 0):
-        """
-        Track usage of local models (like CodeBERT).
-        """
         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:
-        """
-        Atomically checks quota and increments counters for one outgoing
-        request. Returns False if the daily quota is exceeded (caller should
-        raise before touching the network). This closes the old race where
-        two threads could both pass the quota check before either one
-        incremented daily_requests_count.
-        """
         with self._stats_lock:
-            if self.stats["daily_requests_count"] >= 500:
+            if self.stats["daily_requests_count"] >= 1000:
                 return False
             self.stats["total_requests"] += 1
             self.stats["daily_requests_count"] += 1
             self._save_stats()
-
         with self._rpm_lock:
             self.requests_this_minute += 1
-
         return True
 
-    def generate(self, prompt: str, max_tokens: int = 8192) -> str:
+    def generate(self, prompt: str, max_tokens: int = 2048) -> str:
         import random
 
         self._check_daily_reset()
         self._check_rpm_window()
 
-        if not self.api_key:
-            logger.error("Cannot generate: Missing GEMINI_API_KEY")
+        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 (500 requests). Request blocked.")
+            logger.error("❌ Daily Quota Exceeded. Request blocked.")
             raise RuntimeError("QUOTA_EXCEEDED")
 
-        # Log the prompt source (first 50 chars) to identify the caller
-        logger.info(f"🤖 Generating with Gemini. Prompt start: {prompt[:50]}...")
-
-        model = genai.GenerativeModel(
-            model_name=self.model_name,
-            generation_config=self.generation_config,
-            safety_settings=self.safety_settings,
-        )
+        logger.info(f"🤖 Generating with Groq ({self.model_name}). Prompt start: {prompt[:50]}...")
 
         retries = 0
-        max_retries = 5
+        max_retries = 3
         base_delay = 2
 
         while retries <= max_retries:
@@ -212,24 +150,39 @@ class LLMSingleton:
                     self.stats["input_tokens"] += len(prompt) // 4
                     self._save_stats()
 
-                response = model.generate_content(prompt)
+                chat_completion = self.client.chat.completions.create(
+                    model=self.model_name,
+                    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 ""
 
                 with self._stats_lock:
                     self.stats["successful_requests"] += 1
-                    if response.text:
-                        self.stats["output_tokens"] += len(response.text) // 4
+                    if response_text:
+                        self.stats["output_tokens"] += len(response_text) // 4
                     self._save_stats()
 
-                return response.text.strip()
+                return response_text.strip()
+
             except Exception as e:
                 error_str = str(e)
-                if "429" in error_str or "quota" in error_str.lower():
+                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, 1)
+                    wait_time = (base_delay * (2**retries)) + random.uniform(0.5, 1.5)
                     logger.warning(
-                        f"⚠️ Rate limit hit. Retrying in {wait_time:.2f}s... (Attempt {retries+1}/{max_retries})"
+                        f"⚠️ Groq rate limit hit. Retrying in {wait_time:.2f}s... (Attempt {retries + 1}/{max_retries})"
                     )
                     time.sleep(wait_time)
                     retries += 1
@@ -237,7 +190,7 @@ class LLMSingleton:
                     with self._stats_lock:
                         self.stats["errors"] += 1
                         self._save_stats()
-                    logger.error(f"Gemini generation failed: {e}")
+                    logger.error(f"Groq generation failed: {e}")
                     raise RuntimeError(f"GENERATION_FAILED: {e}")
 
         with self._stats_lock:
@@ -247,33 +200,43 @@ class LLMSingleton:
         raise RuntimeError("RATE_LIMIT_EXCEEDED")
 
     def generate_text(self, prompt: str) -> str:
-        """
-        Helper for non-JSON text generation (like Guides).
-        """
         self._check_daily_reset()
         self._check_rpm_window()
 
-        if not self.api_key:
+        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:
-            # Override config for text
-            config = self.generation_config.copy()
-            config["response_mime_type"] = "text/plain"
-
-            model = genai.GenerativeModel(
-                model_name=self.model_name,
-                generation_config=config,
-                safety_settings=self.safety_settings,
+            with self._stats_lock:
+                self.stats["input_tokens"] += len(prompt) // 4
+                self._save_stats()
+
+            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 = model.generate_content(prompt)
-            return response.text
+
+            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()
+
+            return response_text
         except Exception as e:
-            logger.error(f"Gemini text generation failed: {e}")
+            logger.error(f"Groq text generation failed: {e}")
             return f"Error generating content: {str(e)}"
 
-
 llm_engine = LLMSingleton.get_instance()