Update src/sql_generator.py

src/sql_generator.py

@@ -1,76 +1,121 @@
-import os
-import re
-import json
-from huggingface_hub import InferenceClient
-from dotenv import load_dotenv
-
-class SQLGenerator:
-    def __init__(self, api_key=None):
-        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)
-
-    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 = []
-
-        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."
-
-        history_text = ""
-        if history:
-            history_text = "PREVIOUS CONVERSATION:\n" + "\n".join([f"User: {h['user']}\nSQL: {h['sql']}" for h in history[-2:]])
-
-        system_prompt = f"""You are an elite SQL Expert.
-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.
-"""
-        messages = [{"role": "system", "content": system_prompt}, {"role": "user", "content": question}]
-
-        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
-            
-            sql_query = ""
-            message = "Here is the data."
-            explanation = "Query generated successfully."
-
-            try:
-                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:
-                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 += ";"
-            
-            # ✅ 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."
-
-            return sql_query, explanation, message
-
-        except Exception as e:
-            print(f"   ❌ Model Error: {e}")
-            safe_e = str(e).replace("'", "").replace('"', "")
+import os
+import requests
+import re
+import json
+from dotenv import load_dotenv
+
+class SQLGenerator:
+    def __init__(self):
+        load_dotenv()
+        
+        # 1. Robustly fetch the API Key
+        # Try HF_API_KEY first (our standard), then fall back to others if user renamed it
+        self.api_key = os.getenv("HF_API_KEY") or os.getenv("HUGGINGFACEHUB_API_TOKEN")
+        
+        # 2. Use the powerful Qwen 2.5 Coder model
+        self.repo_id = "Qwen/Qwen2.5-Coder-32B-Instruct"
+        self.api_url = f"https://api-inference.huggingface.co/models/{self.repo_id}"
+
+    def generate_followup_questions(self, question, sql_query):
+        # Basic heuristics to suggest next steps
+        return ["Visualize this result", "Export as CSV", "Compare with last year"]
+
+    def generate_sql(self, question, context, history=None):
+        if history is None: history = []
+
+        # 🚨 ERROR CHECK: Stop early if key is missing
+        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: Block dangerous keywords
+        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 for the AI
+        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"""You are an elite SQL Expert.
+        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.
+        """
+        
+        # Prepare the payload (Direct HTTP Request)
+        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": 1024,
+                "temperature": 0.1,  # Low temp = precise SQL
+                "return_full_text": False
+            }
+        }
+
+        headers = {
+            "Authorization": f"Bearer {self.api_key}",
+            "Content-Type": "application/json"
+        }
+
+        try:
+            print(f"   ⚡ Generating SQL via Direct API...")
+            
+            # 🚀 DIRECT REQUEST (Bypasses library auth issues)
+            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."
+
+            # Parse Response
+            result = response.json()
+            
+            # Handle different response formats (sometimes list, sometimes dict)
+            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)
+
+            # JSON Parsing & Cleanup
+            sql_query = ""
+            message = "Here is the data."
+            explanation = "Query generated successfully."
+
+            try:
+                # Remove markdown code blocks if AI added them
+                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 SQL Cleanup
+            sql_query = sql_query.strip().replace("\n", " ")
+            if sql_query and not sql_query.endswith(";"): sql_query += ";"
+            
+            # 🛡️ Final Validation (Allow SELECT or WITH)
+            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"):
+                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."
+
+            return sql_query, explanation, message
+
+        except Exception as e:
+            print(f"   ❌ Model Error: {e}")
+            safe_e = str(e).replace("'", "").replace('"', "")
             return f"SELECT 'Error: {safe_e}' as status", "System Error", "An unexpected error occurred."