Update src/sql_generator.py

src/sql_generator.py

@@ -8,30 +8,21 @@ class SQLGenerator:
     def __init__(self):
         load_dotenv()
         
-        # 1. AUTHENTICATION (With Debugging)
+        # 1. AUTHENTICATION
         raw_key = os.getenv("HF_API_KEY") or os.getenv("HUGGINGFACEHUB_API_TOKEN")
         self.api_key = raw_key.strip() if raw_key else None
         
-        if not self.api_key:
-            print("   ❌ FATAL: API Key missing.")
-        else:
-            print(f"   ✅ API Key loaded: {self.api_key[:5]}...")
-
-        # 2. THE "SHOTGUN" MODEL LIST
-        # We try these 5 models in order. One WILL work.
+        # 2. ACTIVE FREE TIER MODELS (2025)
+        # We prioritize "Showcase" models which are kept online by sponsors.
         self.models = [
-            "microsoft/Phi-3-mini-4k-instruct",      # High availability, very fast
-            "google/gemma-1.1-7b-it",                # Google's open model (very stable)
-            "mistralai/Mistral-7B-Instruct-v0.3",    # Standard free tier workhorse
-            "HuggingFaceH4/zephyr-7b-beta",          # Reliable fallback
-            "Qwen/Qwen2.5-Coder-7B-Instruct"         # Excellent coder (if online)
+            "Qwen/Qwen2.5-72B-Instruct",        # Currently the #1 Free Showcase Model
+            "Qwen/Qwen2.5-7B-Instruct",         # Reliable Backup
+            "microsoft/Phi-3.5-mini-instruct",  # Newest Microsoft Model (Active)
+            "mistralai/Mistral-Nemo-Instruct-2407" # New Mistral Standard
         ]
         
-        # 3. ENDPOINTS (Router + Legacy)
-        self.endpoints = [
-            "https://router.huggingface.co/models/",
-            "https://api-inference.huggingface.co/models/"
-        ]
+        # 3. ENDPOINTS
+        self.base_url = "https://router.huggingface.co/models/"
 
     def generate_followup_questions(self, question, sql_query):
         return ["Visualize this result", "Export as CSV", "Compare with last year"]
@@ -42,12 +33,12 @@ class SQLGenerator:
         if not self.api_key:
             return "SELECT 'Error: HF_API_KEY Missing' as status", "Configuration Error", "Please add HF_API_KEY to your Space Secrets."
 
-        # 🛡️ Safety
+        # 🛡️ Safety Layer
         forbidden = ["DROP", "DELETE", "UPDATE", "INSERT", "ALTER", "TRUNCATE", "GRANT"]
         if any(word in question.upper() for word in forbidden):
              return "SELECT 'Error: Blocked by Safety Layer' as status", "Safety Alert", "I cannot execute commands that modify data."
 
-        # Simple Prompt
+        # Prompt
         system_prompt = f"""You are an SQL Expert.
         Schema:
         {context}
@@ -60,35 +51,31 @@ class SQLGenerator:
         Question: {question}"""
         
         payload = {
-            "inputs": f"<|user|>\n{system_prompt}\n<|end|>\n<|assistant|>\n",
+            "inputs": f"<|im_start|>system\n{system_prompt}<|im_end|>\n<|im_start|>user\n{question}<|im_end|>\n<|im_start|>assistant\n",
             "parameters": {"max_new_tokens": 512, "temperature": 0.1, "return_full_text": False}
         }
 
         headers = {"Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json"}
 
-        # 🔄 ULTRA-ROBUST RETRY LOOP
+        # 🔄 RETRY LOOP
         errors = []
-        
         for model in self.models:
-            for base_url in self.endpoints:
-                api_url = f"{base_url}{model}"
-                try:
-                    print(f"   ⚡ Attempting: {model}...")
-                    response = requests.post(api_url, headers=headers, json=payload, timeout=15)
-                    
-                    if response.status_code == 200:
-                        print(f"   ✅ SUCCESS with {model}!")
-                        return self._process_response(response.json())
-                    
-                    # Log failure and continue immediately
-                    print(f"   ❌ Failed ({response.status_code})")
-                    errors.append(f"{model}: {response.status_code}")
-                    
-                except Exception as e:
-                    print(f"   ⚠️ Connection Error: {e}")
-                    errors.append(f"{model}: Error")
+            api_url = f"{self.base_url}{model}"
+            try:
+                print(f"   ⚡ Attempting: {model}...")
+                response = requests.post(api_url, headers=headers, json=payload, timeout=20)
+                
+                if response.status_code == 200:
+                    print(f"   ✅ SUCCESS with {model}!")
+                    return self._process_response(response.json())
+                
+                print(f"   ❌ Failed ({response.status_code})")
+                errors.append(f"{model}: {response.status_code}")
+                
+            except Exception as e:
+                print(f"   ⚠️ Connection Error: {e}")
+                errors.append(f"{model}: Error")
 
-        # If we get here, literally everything failed (Rare)
         return f"SELECT 'Error: All models failed' as status", "System Error", f"Debug Info: {', '.join(errors)}"
 
     def _process_response(self, result):
@@ -104,9 +91,7 @@ class SQLGenerator:
         explanation = "Query generated successfully."
 
         try:
-            # Clean and Extract
             clean_json = re.sub(r"```json|```", "", raw_text).strip()
-            # Regex to find JSON
             json_match = re.search(r"\{.*\}", clean_json, re.DOTALL)
             if json_match:
                 data = json.loads(json_match.group(0))
@@ -123,8 +108,7 @@ class SQLGenerator:
         sql_query = sql_query.strip().replace("\n", " ")
         if sql_query and not sql_query.endswith(";"): sql_query += ";"
         
-        # Fallback for empty SQL
         if not sql_query:
-            sql_query = "SELECT 'Error: AI generated empty query' as status"
+            sql_query = "SELECT 'Error: Empty Query' as status"
 
         return sql_query, explanation, message