Update src/sql_generator.py

src/sql_generator.py

@@ -12,18 +12,13 @@ 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. MODEL LIST (The "Shotgun" Strategy)
-        # We try 4 different models. If one fails (404/429), it instantly tries the next.
-        self.models = [
-            "mistralai/Mistral-7B-Instruct-v0.3",    # Most popular free model
-            "google/gemma-1.1-7b-it",                # Highly reliable backup
-            "microsoft/Phi-3-mini-4k-instruct",      # Very fast, almost always online
-            "HuggingFaceH4/zephyr-7b-beta"           # Old reliable (last resort)
-        ]
+        # 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"
         
-        # 3. BASE URLs (Try Router first, then Legacy)
-        self.router_url = "https://router.huggingface.co/models/"
-        self.legacy_url = "https://api-inference.huggingface.co/models/"
+        # 3. ROUTER URL
+        self.api_url = f"https://router.huggingface.co/models/{self.repo_id}"
 
     def generate_followup_questions(self, question, sql_query):
         return ["Visualize this result", "Export as CSV", "Compare with last year"]
@@ -39,13 +34,13 @@ class SQLGenerator:
         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
-        system_prompt = f"""<|system|>
-        You are an elite SQL Expert.
+        # System Prompt (Qwen ChatML Format)
+        system_prompt = f"""You are an elite SQL Expert.
         DATABASE SCHEMA:
         {context}
 
@@ -55,14 +50,13 @@ class SQLGenerator:
         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} </s>
-        <|assistant|>"""
-        
+        # Payload (Qwen Specific Format)
         payload = {
-            "inputs": system_prompt,
+            "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,
+                "max_new_tokens": 1024,
                 "temperature": 0.1,
                 "return_full_text": False
             }
@@ -73,72 +67,58 @@ class SQLGenerator:
             "Content-Type": "application/json"
         }
 
-        # 🔄 ROBUST RETRY LOOP
-        last_error = ""
-        
-        for model in self.models:
-            # Try Router URL first, then Legacy URL
-            urls_to_try = [self.router_url + model, self.legacy_url + model]
+        try:
+            print(f"   ⚡ Generating SQL using {self.repo_id}...")
             
-            for api_url in urls_to_try:
-                try:
-                    print(f"   ⚡ Trying {model} at {api_url}...")
-                    
-                    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())
-                    
-                    # If 404/410/500, we log and continue to next
-                    print(f"   ❌ Failed ({response.status_code}). Trying next...")
-                    last_error = f"{response.status_code}: {response.text}"
-                    
-                except Exception as e:
-                    print(f"   ⚠️ Connection Error: {e}")
-                    last_error = str(e)
-
-        # If ALL models fail
-        return f"SELECT 'Error: {last_error}' as status", "System Error", "All AI models are currently unavailable. Check your HF_API_KEY."
-
-    def _process_response(self, result):
-        """Helper to parse the AI response cleanly"""
-        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."
+            # 🚀 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}"
 
-        try:
-            # Clean Markdown
-            clean_json = re.sub(r"```json|```", "", raw_text).strip()
-            # Attempt to find JSON object
-            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)
+            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', '')
             else:
-                 # Fallback regex
-                 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)
-
-        # Final Cleanup
-        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"):
-             # Last resort fallback for clean "SELECT"
-             return "SELECT 'Error: Invalid Query Type' as status", "Safety Error", "I can only perform read-only operations."
+                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
 
-        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."