# 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 # 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() @classmethod def get_instance(cls): if cls._instance is None: with cls._instance_lock: if cls._instance is None: cls._instance = cls() return cls._instance def __init__(self): if self._instance is not None: raise Exception("Singleton instance already exists!") self.api_key = os.getenv("GEMINI_API_KEY") if not self.api_key: logger.warning( "⚠️ GEMINI_API_KEY not found in environment variables. AI features will fail." ) 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, } # 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._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.minute_window_start = time.time() self.requests_this_minute = 0 def _load_stats(self): default_stats = { "total_requests": 0, "successful_requests": 0, "rate_limit_hits": 0, "input_tokens": 0, "output_tokens": 0, "errors": 0, "local_model_requests": 0, "date": datetime.now().strftime("%Y-%m-%d"), "daily_requests_count": 0, } if os.path.exists(STATS_FILE): 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) except Exception as e: logger.error(f"Failed to save stats: {e}") def _check_daily_reset(self): with self._stats_lock: today = datetime.now().strftime("%Y-%m-%d") if self.stats.get("date") != today: logger.info("📅 New day detected. Resetting daily AI quotas.") self.stats["date"] = today self.stats["daily_requests_count"] = 0 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: self.minute_window_start = now self.requests_this_minute = 0 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 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_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: 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: import random self._check_daily_reset() self._check_rpm_window() if not self.api_key: logger.error("Cannot generate: Missing GEMINI_API_KEY") raise RuntimeError("MISSING_API_KEY") if not self._reserve_request_slot(): logger.error("❌ Daily Quota Exceeded (500 requests). 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, ) retries = 0 max_retries = 5 base_delay = 2 while retries <= max_retries: try: with self._stats_lock: self.stats["input_tokens"] += len(prompt) // 4 self._save_stats() response = model.generate_content(prompt) with self._stats_lock: self.stats["successful_requests"] += 1 if response.text: self.stats["output_tokens"] += len(response.text) // 4 self._save_stats() return response.text.strip() except Exception as e: error_str = str(e) if "429" in error_str or "quota" 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) logger.warning( f"⚠️ 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"Gemini generation failed: {e}") raise RuntimeError(f"GENERATION_FAILED: {e}") with self._stats_lock: self.stats["errors"] += 1 self._save_stats() logger.error("❌ Max retries reached. Request failed.") 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: 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, ) response = model.generate_content(prompt) return response.text except Exception as e: logger.error(f"Gemini text generation failed: {e}") return f"Error generating content: {str(e)}" llm_engine = LLMSingleton.get_instance()