Update src/sql_generator.py

src/sql_generator.py

@@ -1,15 +1,25 @@
 import os
+import requests
 import re
 import json
-from huggingface_hub import InferenceClient
 from dotenv import load_dotenv
 
 class SQLGenerator:
-    def __init__(self, api_key=None):
+    def __init__(self):
         load_dotenv()
-        self.api_token = os.getenv("HUGGINGFACEHUB_API_TOKEN")
-        self.repo_id = "Qwen/Qwen2.5-Coder-32B-Instruct"
-        self.client = InferenceClient(token=self.api_token, timeout=25.0)
+        
+        # 1. ROBUSTLY FETCH & CLEAN THE KEY
+        # We check both names and use .strip() to remove invisible newlines
+        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 A RELIABLE FREE TIER MODEL
+        # We switch to Mistral-7B-Instruct-v0.3 (Supported on Free Tier)
+        # Qwen-32B is too big and causes 404 errors.
+        self.repo_id = "mistralai/Mistral-7B-Instruct-v0.3"
+        
+        # 3. USE THE NEW ROUTER URL (Fixes 410 Gone errors)
+        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"]
@@ -17,53 +27,95 @@ class SQLGenerator:
     def generate_sql(self, question, context, history=None):
         if history is None: history = []
 
+        # 🚨 ERROR CHECK
+        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 (Keyword Block)
         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['user']}\nSQL: {h['sql']}" for h in history[-2:]])
+            history_text = "PREVIOUS CONVERSATION:\n" + "\n".join([f"User: {h.get('user')}\nSQL: {h.get('sql')}" for h in history[-2:]])
+
+        # System Prompt (Optimized for JSON output)
+        system_prompt = f"""<s>[INST] You are an elite SQL Expert.
+        DATABASE SCHEMA:
+        {context}
 
-        system_prompt = f"""You are an elite SQL Expert.
-Schema:
-{context}
+        {history_text}
 
-{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.
+        
+        QUESTION: {question} [/INST]"""
+        
+        # Payload
+        payload = {
+            "inputs": system_prompt,
+            "parameters": {
+                "max_new_tokens": 1024,
+                "temperature": 0.1,
+                "return_full_text": False
+            }
+        }
 
-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.
-"""
-        messages = [{"role": "system", "content": system_prompt}, {"role": "user", "content": question}]
+        headers = {
+            "Authorization": f"Bearer {self.api_key}",
+            "Content-Type": "application/json"
+        }
 
         try:
-            print(f"   ⚡ Generating SQL...")
-            response = self.client.chat_completion(messages=messages, model=self.repo_id, max_tokens=1024, temperature=0.1)
-            raw_text = response.choices[0].message.content
+            print(f"   ⚡ Generating SQL via Direct API ({self.repo_id})...")
             
+            # 🚀 DIRECT REQUEST (Bypasses library auth issues)
+            response = requests.post(self.api_url, headers=headers, json=payload, timeout=25)
+            
+            # Handle Errors clearly
+            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", f"Model unavailable ({response.status_code})"
+
+            # Parse Response
+            result = response.json()
+            
+            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)
+
+            # Clean JSON
             sql_query = ""
             message = "Here is the data."
             explanation = "Query generated successfully."
 
             try:
+                # Remove Markdown and extra text
+                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)
 
+            # Final Cleanup
             sql_query = sql_query.strip().replace("\n", " ")
             if sql_query and not sql_query.endswith(";"): sql_query += ";"
             
-            # ✅ FIX: Strip comments and whitespace before validation
+            # 🛡️ Final Validation (Allows SELECT or WITH)
             clean_check = re.sub(r"/\*.*?\*/|--.*?\n", "", sql_query, flags=re.DOTALL).strip().upper()
-
-            # ✅ FIX: Allow SELECT or WITH clauses
             if not clean_check.startswith("SELECT") and not clean_check.startswith("WITH"):
                 print(f"   ⚠️ Invalid SQL Blocked: {sql_query}")
                 return "SELECT 'Error: Invalid Query Type (Non-SELECT)' as status", "Safety Error", "I can only perform read-only operations."