Update src/sql_generator.py

src/sql_generator.py

@@ -1,24 +1,15 @@
 import os
-import requests
 import re
 import json
+from huggingface_hub import InferenceClient
 from dotenv import load_dotenv
 
 class SQLGenerator:
-    def __init__(self):
+    def __init__(self, api_key=None):
         load_dotenv()
-        
-        # 1. ROBUSTLY FETCH & CLEAN THE KEY
-        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 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"
-        
-        # 3. USE THE NEW ROUTER URL
-        self.api_url = f"https://router.huggingface.co/models/{self.repo_id}"
+        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)
 
     def generate_followup_questions(self, question, sql_query):
         return ["Visualize this result", "Export as CSV", "Compare with last year"]
@@ -26,94 +17,53 @@ 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
         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 (Optimized for Mistral)
-        system_prompt = f"""<s>[INST] You are an expert SQL Assistant.
-        DATABASE SCHEMA:
-        {context}
+            history_text = "PREVIOUS CONVERSATION:\n" + "\n".join([f"User: {h['user']}\nSQL: {h['sql']}" for h in history[-2:]])
 
-        {history_text}
+        system_prompt = f"""You are an elite SQL Expert.
+Schema:
+{context}
 
-        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
-            }
-        }
+{history_text}
 
-        headers = {
-            "Authorization": f"Bearer {self.api_key}",
-            "Content-Type": "application/json"
-        }
+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}]
 
         try:
-            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", f"Model Error: {response.status_code}"
-
-            # Parse Response
-            result = response.json()
+            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
             
-            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:
-                # 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)
 
-            # 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
             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."