Update src/sql_generator.py

src/sql_generator.py

@@ -12,10 +12,12 @@ class SQLGenerator:
         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
-        self.repo_id = "Qwen/Qwen2.5-Coder-32B-Instruct"
+        # 2. USE A MODEL THAT IS DEFINITELY FREE & ONLINE
+        # The 32B model caused the 404. We switch to Mistral-7B-Instruct-v0.3
+        # It is excellent for SQL and 100% supported on the free tier.
+        self.repo_id = "mistralai/Mistral-7B-Instruct-v0.3"
         
-        # ✅ FIX: Updated to the new Hugging Face Router URL
+        # 3. USE THE NEW ROUTER URL
         self.api_url = f"https://router.huggingface.co/models/{self.repo_id}"
 
     def generate_followup_questions(self, question, sql_query):
@@ -38,22 +40,23 @@ class SQLGenerator:
         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
-        system_prompt = f"""You are an elite SQL Expert.
-        Schema:
+        # System Prompt (Optimized for Mistral)
+        system_prompt = f"""<s>[INST] You are an expert SQL Assistant.
+        DATABASE SCHEMA:
         {context}
 
         {history_text}
 
-        Rules:
-        1. Output JSON: {{ "sql": "SELECT ...", "message": "Friendly text", "explanation": "Brief summary" }}
-        2. Query MUST be Read-Only (SELECT).
-        3. Do not include markdown formatting like ```json.
-        """
+        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.
+        
+        QUESTION: {question} [/INST]"""
         
         # Payload
         payload = {
-            "inputs": f"<|im_start|>system\n{system_prompt}<|im_end|>\n<|im_start|>user\n{question}<|im_end|>\n<|im_start|>assistant\n",
+            "inputs": system_prompt,
             "parameters": {
                 "max_new_tokens": 1024,
                 "temperature": 0.1,
@@ -67,14 +70,14 @@ class SQLGenerator:
         }
 
         try:
-            print(f"   ⚡ Generating SQL via Direct API...")
+            print(f"   ⚡ Generating SQL via Direct API ({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 Status: {response.status_code} - {response.text}")
-                return f"SELECT 'Error: API returned {response.status_code}' as status", "API Error", "The AI model is currently unavailable."
+                return f"SELECT 'Error: API returned {response.status_code}' as status", "API Error", f"Model Error: {response.status_code}"
 
             # Parse Response
             result = response.json()
@@ -92,12 +95,17 @@ class SQLGenerator:
             explanation = "Query generated successfully."
 
             try:
+                # Mistral sometimes keeps the prompt. Split by [/INST] if present.
+                if "[/INST]" in raw_text:
+                    raw_text = raw_text.split("[/INST]")[-1]
+
                 clean_json = re.sub(r"```json|```", "", raw_text).strip()
                 data = json.loads(clean_json)
                 sql_query = data.get("sql", "")
                 message = data.get("message", message)
                 explanation = data.get("explanation", explanation)
             except:
+                # Fallback: Regex to find SQL if JSON parsing fails
                 match = re.search(r"(SELECT[\s\S]+?;)", raw_text, re.IGNORECASE)
                 if match: sql_query = match.group(1)