Update src/sql_generator.py

src/sql_generator.py

@@ -1,4 +1,5 @@
 import os
+import requests
 import re
 import json
 from dotenv import load_dotenv
@@ -8,21 +9,19 @@ class SQLGenerator:
     def __init__(self):
         load_dotenv()
         
-        # 1. CLEAN THE KEY (Fixes "Invalid Header" error)
+        # 1. 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. SETUP CLIENT (Fixes "404/410" errors)
-        # The client automatically handles the complex routing logic
+        # 2. SETUP CLIENT
         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
-       # Change this line in src/sql_generator.py if 32B gets slow:
-        self.model_id = "Qwen/Qwen2.5-Coder-7B-Instruct"
+        # 3. USE QWEN 2.5 (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"]
@@ -36,29 +35,29 @@ 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."
 
-        # Prompt
+        # 🧠 SMART PROMPT (Fixes the "No Such Table" error)
         messages = [
-            {"role": "system", "content": f"""You are an SQL Expert.
-            Database Schema:
-            {context}
+            {"role": "system", "content": f"""You are a precise SQL Expert.
+            
+            CRITICAL RULES:
+            1. You MUST use the EXACT table names and column names from the SCHEMA below.
+            2. Do NOT hallucinate table names (e.g., if schema says 'Employee', do NOT use 'employees').
+            3. Output valid JSON only.
             
-            Rules:
-            1. Output valid JSON: {{ "sql": "SELECT ...", "message": "Short text", "explanation": "Brief summary" }}
-            2. Read-only SELECT queries only.
-            3. No markdown formatting.
+            SCHEMA:
+            {context}
             """},
-            {"role": "user", "content": question}
+            {"role": "user", "content": f"Question: {question}\nReturn JSON format: {{ 'sql': 'SELECT ...', 'message': '...', 'explanation': '...' }}"}
         ]
 
         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,
+                temperature=0.1, # Low temp = More strict
                 stream=False
             )
             
@@ -67,15 +66,15 @@ class SQLGenerator:
 
         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):
+            # Failover to 7B if 32B is busy
+            if "404" in str(e) or "429" in str(e) or "503" in str(e):
                 return self._fallback_generate(messages)
             return f"SELECT 'Error: {str(e)[:50]}' as status", "System Error", "AI Model unavailable."
 
     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"
+            # Fallback to the smaller, faster model
+            backup_model = "Qwen/Qwen2.5-Coder-7B-Instruct" 
             print(f"   ⚠️ Switching to backup: {backup_model}...")
             response = self.client.chat.completions.create(
                 model=backup_model,