Update src/sql_generator.py

src/sql_generator.py

@@ -1,36 +1,33 @@
 import os
-import requests
 import re
 import json
 from dotenv import load_dotenv
+from huggingface_hub import InferenceClient
 
 class SQLGenerator:
     def __init__(self):
         load_dotenv()
         
-        # 1. AUTHENTICATION
+        # 1. CLEAN THE KEY (Fixes "Invalid Header" error)
         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. ACTIVE FREE TIER MODELS (2025)
-        # We prioritize "Showcase" models which are kept online by sponsors.
-        self.models = [
-            "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
-        self.base_url = "https://router.huggingface.co/models/"
+        # 2. SETUP CLIENT (Fixes "404/410" errors)
+        # The client automatically handles the complex routing logic
+        if self.api_key:
+            self.client = InferenceClient(api_key=self.api_key)
+        else:
+            self.client = None
+            print("   ❌ FATAL: API Key missing.")
+
+        # 3. USE THE BEST FREE MODEL
+        self.model_id = "Qwen/Qwen2.5-Coder-32B-Instruct"
 
     def generate_followup_questions(self, question, sql_query):
         return ["Visualize this result", "Export as CSV", "Compare with last year"]
 
     def generate_sql(self, question, context, history=None):
-        if history is None: history = []
-
-        if not self.api_key:
+        if not self.client:
             return "SELECT 'Error: HF_API_KEY Missing' as status", "Configuration Error", "Please add HF_API_KEY to your Space Secrets."
 
         # 🛡️ Safety Layer
@@ -39,53 +36,56 @@ class SQLGenerator:
              return "SELECT 'Error: Blocked by Safety Layer' as status", "Safety Alert", "I cannot execute commands that modify data."
 
         # Prompt
-        system_prompt = f"""You are an SQL Expert.
-        Schema:
-        {context}
-
-        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 = {
-            "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"}
+        messages = [
+            {"role": "system", "content": f"""You are an SQL Expert.
+            Database Schema:
+            {context}
+            
+            Rules:
+            1. Output valid JSON: {{ "sql": "SELECT ...", "message": "Short text", "explanation": "Brief summary" }}
+            2. Read-only SELECT queries only.
+            3. No markdown formatting.
+            """},
+            {"role": "user", "content": question}
+        ]
 
-        # 🔄 RETRY LOOP
-        errors = []
-        for model in self.models:
-            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")
+        try:
+            print(f"   ⚡ Generating SQL using {self.model_id}...")
+            
+            # 🚀 OFFICIAL CLIENT CALL (The Robust Way)
+            response = self.client.chat.completions.create(
+                model=self.model_id,
+                messages=messages,
+                max_tokens=500,
+                temperature=0.1,
+                stream=False
+            )
+            
+            raw_text = response.choices[0].message.content
+            return self._process_response(raw_text)
 
-        return f"SELECT 'Error: All models failed' as status", "System Error", f"Debug Info: {', '.join(errors)}"
+        except Exception as e:
+            print(f"   ❌ AI ERROR: {e}")
+            # Failover to backup model if Qwen is busy
+            if "404" in str(e) or "429" in str(e):
+                return self._fallback_generate(messages)
+            return f"SELECT 'Error: {str(e)[:50]}' as status", "System Error", "AI Model unavailable."
 
-    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)
+    def _fallback_generate(self, messages):
+        """Backup using a smaller model if the main one fails"""
+        try:
+            backup_model = "meta-llama/Llama-3.2-3B-Instruct"
+            print(f"   ⚠️ Switching to backup: {backup_model}...")
+            response = self.client.chat.completions.create(
+                model=backup_model,
+                messages=messages,
+                max_tokens=500
+            )
+            return self._process_response(response.choices[0].message.content)
+        except Exception as e:
+             return "SELECT 'Error: All models failed' as status", "System Error", "Please check your API Key permissions."
 
+    def _process_response(self, raw_text):
         sql_query = ""
         message = "Here is the data."
         explanation = "Query generated successfully."