Update src/sql_generator.py

src/sql_generator.py

@@ -8,17 +8,30 @@ class SQLGenerator:
     def __init__(self):
         load_dotenv()
         
-        # 1. GET KEY (Cleaned)
+        # 1. AUTHENTICATION (With Debugging)
         raw_key = os.getenv("HF_API_KEY") or os.getenv("HUGGINGFACEHUB_API_TOKEN")
         self.api_key = raw_key.strip() if raw_key else None
         
-        # 2. USE QWEN 2.5 CODER (7B)
-        # This is the Best "Ungated" Free Model right now.
-        # It does not require a license agreement click, unlike Mistral/Llama.
-        self.repo_id = "Qwen/Qwen2.5-Coder-7B-Instruct"
+        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.
+        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)
+        ]
         
-        # 3. ROUTER URL
-        self.api_url = f"https://router.huggingface.co/models/{self.repo_id}"
+        # 3. ENDPOINTS (Router + Legacy)
+        self.endpoints = [
+            "https://router.huggingface.co/models/",
+            "https://api-inference.huggingface.co/models/"
+        ]
 
     def generate_followup_questions(self, question, sql_query):
         return ["Visualize this result", "Export as CSV", "Compare with last year"]
@@ -29,96 +42,89 @@ 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 Layer
+        # 🛡️ Safety
         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."
 
-        # Format History
-        history_text = ""
-        if history:
-            history_text = "PREVIOUS CONVERSATION:\n" + "\n".join([f"User: {h.get('user')}\nSQL: {h.get('sql')}" for h in history[-2:]])
-
-        # System Prompt (Qwen ChatML Format)
-        system_prompt = f"""You are an elite SQL Expert.
-        DATABASE SCHEMA:
+        # Simple Prompt
+        system_prompt = f"""You are an SQL Expert.
+        Schema:
         {context}
 
-        {history_text}
-
-        RULES:
-        1. Output ONLY a valid JSON object. Format: {{ "sql": "SELECT ...", "message": "Friendly text", "explanation": "Brief summary" }}
-        2. The SQL query MUST be Read-Only (SELECT).
-        3. Do not use markdown formatting.
-        """
+        Rules:
+        1. Output valid JSON: {{ "sql": "SELECT ...", "message": "Short text", "explanation": "Brief summary" }}
+        2. Read-only SELECT queries only.
+        3. No markdown.
+        
+        Question: {question}"""
         
-        # Payload (Qwen Specific Format)
         payload = {
-            "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": 1024,
-                "temperature": 0.1,
-                "return_full_text": False
-            }
+            "inputs": f"<|user|>\n{system_prompt}\n<|end|>\n<|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"
-        }
+        headers = {"Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json"}
+
+        # 🔄 ULTRA-ROBUST 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")
+
+        # 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):
+        if isinstance(result, list) and len(result) > 0:
+            raw_text = result[0].get('generated_text', '')
+        elif isinstance(result, dict):
+            raw_text = result.get('generated_text', '')
+        else:
+            raw_text = str(result)
+
+        sql_query = ""
+        message = "Here is the data."
+        explanation = "Query generated successfully."
 
         try:
-            print(f"   ⚡ Generating SQL using {self.repo_id}...")
-            
-            # 🚀 DIRECT REQUEST
-            response = requests.post(self.api_url, headers=headers, json=payload, timeout=25)
-            
-            if response.status_code != 200:
-                print(f"   ❌ API FAILURE: {response.status_code}")
-                print(f"   ❌ RESPONSE: {response.text}")
-                return f"SELECT 'Error: API returned {response.status_code}' as status", "API Error", f"Model Error: {response.status_code}"
-
-            result = response.json()
-            
-            # Helper to extract text
-            if isinstance(result, list) and len(result) > 0:
-                raw_text = result[0].get('generated_text', '')
-            elif isinstance(result, dict):
-                raw_text = result.get('generated_text', '')
+            # 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))
+                sql_query = data.get("sql", "")
+                message = data.get("message", message)
+                explanation = data.get("explanation", explanation)
             else:
-                raw_text = str(result)
-
-            # JSON Parsing
-            sql_query = ""
-            message = "Here is the data."
-            explanation = "Query generated successfully."
-
-            try:
-                clean_json = re.sub(r"```json|```", "", raw_text).strip()
-                json_match = re.search(r"\{.*\}", clean_json, re.DOTALL)
-                if json_match:
-                    data = json.loads(json_match.group(0))
-                    sql_query = data.get("sql", "")
-                    message = data.get("message", message)
-                    explanation = data.get("explanation", explanation)
-                else:
-                    # Regex Fallback
-                    match = re.search(r"(SELECT[\s\S]+?;)", raw_text, re.IGNORECASE)
-                    if match: sql_query = match.group(1)
-            except:
-                match = re.search(r"(SELECT[\s\S]+?;)", raw_text, re.IGNORECASE)
-                if match: sql_query = match.group(1)
-
-            sql_query = sql_query.strip().replace("\n", " ")
-            if sql_query and not sql_query.endswith(";"): sql_query += ";"
-            
-            clean_check = re.sub(r"/\*.*?\*/|--.*?\n", "", sql_query, flags=re.DOTALL).strip().upper()
-            if not clean_check.startswith("SELECT") and not clean_check.startswith("WITH"):
-                return "SELECT 'Error: Invalid Query Type' as status", "Safety Error", "I can only perform read-only operations."
-
-            return sql_query, explanation, message
-
-        except Exception as e:
-            print(f"   ❌ SYSTEM EXCEPTION: {e}")
-            safe_e = str(e).replace("'", "").replace('"', "")
-            return f"SELECT 'Error: {safe_e}' as status", "System Error", "An unexpected error occurred."
+                 match = re.search(r"(SELECT[\s\S]+?;)", raw_text, re.IGNORECASE)
+                 if match: sql_query = match.group(1)
+        except:
+            match = re.search(r"(SELECT[\s\S]+?;)", raw_text, re.IGNORECASE)
+            if match: sql_query = match.group(1)
+
+        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"
+
+        return sql_query, explanation, message