using Gemini NL2SQL

pipeline.py

@@ -2,54 +2,63 @@ import os
 import requests
 import time
 import re
+import json
 from db_utils import get_schema, execute_sql
 
-# Hugging Face Inference API endpoint for Qwen2.5-Coder
-API_URL = "https://api-inference.huggingface.co/models/Qwen/Qwen2.5-Coder-7B-Instruct"
-
-def query_huggingface_api(prompt, max_retries=3):
-    """Query the Hugging Face Inference API"""
-    hf_token = os.getenv("HF_TOKEN")
-    if not hf_token:
-        raise ValueError("HF_TOKEN not found in environment variables")
+def query_gemini_api(prompt, max_retries=3):
+    """Query the Google Gemini API"""
+    api_key = os.getenv("GOOGLE_API_KEY")
+    if not api_key:
+        raise ValueError("GOOGLE_API_KEY not found in environment variables")
+    
+    # Gemini API endpoint
+    url = f"https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash:generateContent?key={api_key}"
+    
+    headers = {
+        "Content-Type": "application/json"
+    }
     
-    headers = {"Authorization": f"Bearer {hf_token}"}
     payload = {
-        "inputs": prompt,
-        "parameters": {
-            "max_new_tokens": 200,
+        "contents": [{
+            "parts": [{
+                "text": prompt
+            }]
+        }],
+        "generationConfig": {
             "temperature": 0.1,
-            "do_sample": False,
-            "return_full_text": False,
-            "stop": ["###", "\n\n"]
+            "topK": 1,
+            "topP": 0.8,
+            "maxOutputTokens": 200,
+            "stopSequences": ["```", "\n\n"]
         }
     }
     
     for attempt in range(max_retries):
         try:
-            print(f"=== DEBUG: API attempt {attempt + 1}")
-            response = requests.post(API_URL, headers=headers, json=payload, timeout=30)
+            print(f"=== DEBUG: Gemini API attempt {attempt + 1}")
+            response = requests.post(url, headers=headers, json=payload, timeout=30)
             print(f"=== DEBUG: API Response Status: {response.status_code}")
             
             if response.status_code == 200:
                 result = response.json()
                 print(f"=== DEBUG: API Response: {result}")
                 
-                if isinstance(result, list) and len(result) > 0:
-                    generated_text = result[0].get("generated_text", "").strip()
-                else:
-                    generated_text = str(result).strip()
-                    
-                return generated_text
+                if "candidates" in result and len(result["candidates"]) > 0:
+                    candidate = result["candidates"][0]
+                    if "content" in candidate and "parts" in candidate["content"]:
+                        generated_text = candidate["content"]["parts"][0]["text"].strip()
+                        return generated_text
+                
+                return "No valid response generated"
             
-            elif response.status_code == 503:
-                wait_time = 20 * (attempt + 1)
-                print(f"=== DEBUG: Model loading, waiting {wait_time} seconds...")
+            elif response.status_code == 429:
+                wait_time = 60 * (attempt + 1)  # Rate limit - wait longer
+                print(f"=== DEBUG: Rate limited, waiting {wait_time} seconds...")
                 time.sleep(wait_time)
                 continue
                 
             else:
-                error_msg = f"API Error {response.status_code}: {response.text}"
+                error_msg = f"Gemini API Error {response.status_code}: {response.text}"
                 print(f"=== DEBUG: {error_msg}")
                 if attempt == max_retries - 1:
                     raise Exception(error_msg)
@@ -90,9 +99,18 @@ def clean_sql_output(sql_text, user_limit=None):
     # Remove markdown formatting
     if sql_text.startswith("```"):
         lines = sql_text.split('\n')
-        sql_text = '\n'.join(lines[1:-1]) if len(lines) > 2 else sql_text
+        # Find SQL content between backticks
+        sql_lines = []
+        in_sql = False
+        for line in lines:
+            if line.strip().startswith("```"):
+                in_sql = not in_sql
+                continue
+            if in_sql:
+                sql_lines.append(line)
+        sql_text = '\n'.join(sql_lines)
     
-    # Handle multiple lines - take the SQL part
+    # Handle multiple lines - extract the main SELECT query
     lines = sql_text.split('\n')
     sql = ""
     for line in lines:
@@ -103,10 +121,10 @@ def clean_sql_output(sql_text, user_limit=None):
                 break
     
     if not sql:
-        # If no SELECT found, take the first non-empty line
+        # If no SELECT found, take the first non-empty line that looks like SQL
         for line in lines:
             line = line.strip()
-            if line:
+            if line and any(keyword in line.upper() for keyword in ['SELECT', 'WITH', 'FROM']):
                 sql = line
                 break
     
@@ -120,7 +138,7 @@ def clean_sql_output(sql_text, user_limit=None):
     return sql
 
 def text_to_sql(nl_query):
-    """Convert natural language to SQL using Qwen2.5-Coder via HF Inference API"""
+    """Convert natural language to SQL using Google Gemini"""
     try:
         print(f"=== DEBUG: Starting text_to_sql with query: {nl_query}")
         
@@ -136,37 +154,34 @@ def text_to_sql(nl_query):
         user_limit = extract_user_requested_limit(nl_query)
         print(f"=== DEBUG: Extracted user limit: {user_limit}")
         
-        # Create optimized prompt for Qwen2.5-Coder
-        prompt = f"""<|im_start|>system
-You are an expert SQL developer. Generate PostgreSQL queries based on natural language questions.
+        # Create optimized prompt for Gemini
+        prompt = f"""You are an expert PostgreSQL developer. Convert this natural language question to a precise SQL query.
+
+Question: {nl_query}
 
 Database Schema:
 {schema[:1500]}
 
-Rules:
-- Return ONLY the SQL query
+Requirements:
+- Generate ONLY the SQL query, no explanation
 - Use PostgreSQL syntax
-- Be precise with table and column names
-- Do not include explanations or markdown formatting
-<|im_end|>
-<|im_start|>user
-{nl_query}
-<|im_end|>
-<|im_start|>assistant
-"""
+- Be precise with table and column names from the schema
+- Return a single SELECT statement
+
+SQL Query:"""
 
-        print("=== DEBUG: Calling Qwen2.5-Coder API...")
-        generated_sql = query_huggingface_api(prompt)
+        print("=== DEBUG: Calling Google Gemini API...")
+        generated_sql = query_gemini_api(prompt)
         print(f"=== DEBUG: Generated SQL raw: {generated_sql}")
         
-        if not generated_sql:
-            return "Error: No SQL generated from the model", []
+        if not generated_sql or "No valid response" in generated_sql:
+            return "Error: No SQL generated from Gemini", []
         
         # Clean the SQL output
         sql = clean_sql_output(generated_sql, user_limit)
         print(f"=== DEBUG: Final cleaned SQL: {sql}")
         
-        if not sql or not sql.upper().startswith('SELECT'):
+        if not sql or not sql.upper().strip().startswith('SELECT'):
             return f"Error: Invalid SQL generated: {sql}", []
         
         # Execute SQL