Error Fix

app.py

@@ -65,10 +65,9 @@ Important guidelines:
 4. For percentiles, use window functions or approximate methods
 5. Keep queries efficient and focused on answering the specific question
 6. Always use 'data_tab' as the table name
+7. IMPORTANT: Return ONLY the SQL query without any markdown formatting, explanations, or code blocks
 
 Question: {question}
-
-SQL Query:
 """)
 
 # Define the prompt for interpreting the SQL query result
@@ -104,12 +103,51 @@ Important guidelines for SQLite syntax:
 
 5. Always use 'data_tab' as the table name
 
+6. IMPORTANT: Return ONLY the SQL query without any markdown formatting, explanations, or code blocks
+
 Question: {question}
 Visualization type: {viz_type}
-
-SQL Query:
 """)
 
+# Add this helper function to clean SQL queries
+def clean_sql_query(query_text):
+    """Clean SQL query text by removing markdown formatting and comments"""
+    # Remove markdown code blocks
+    if "```sql" in query_text.lower():
+        # Extract content between code blocks
+        pattern = r"```(?:sql)?(.*?)```"
+        matches = re.findall(pattern, query_text, re.DOTALL)
+        if matches:
+            query_text = matches[0].strip()
+    
+    # Remove any "Here is the SQL query" text that might precede the query
+    prefixes = [
+        "here is the sql query",
+        "here is the sqlite query",
+        "here is a query",
+        "here's the sql query",
+        "the sql query is"
+    ]
+    
+    for prefix in prefixes:
+        if query_text.lower().startswith(prefix):
+            # Find the first occurrence of "SELECT", "WITH", etc.
+            sql_keywords = ["select", "with", "create", "insert", "update", "delete"]
+            positions = [query_text.lower().find(keyword) for keyword in sql_keywords]
+            positions = [pos for pos in positions if pos != -1]
+            
+            if positions:
+                start_pos = min(positions)
+                query_text = query_text[start_pos:]
+    
+    # Remove SQL comments
+    query_text = re.sub(r'--.*?(\n|$)', ' ', query_text)
+    
+    # Remove trailing semicolon if present (optional)
+    query_text = query_text.strip().rstrip(';')
+    
+    return query_text
+
 def process_text_query(query, history):
     """Process a text query and update chat history"""
     if not query:
@@ -137,7 +175,10 @@ def process_text_query(query, history):
             
             # Generate SQL query using LLM
             ai_msg = query_prompt | llm
-            sql_query = ai_msg.invoke({"question": question_with_context}).content.strip()
+            raw_sql_query = ai_msg.invoke({"question": question_with_context}).content.strip()
+            
+            # Clean the SQL query
+            sql_query = clean_sql_query(raw_sql_query)
             
             print(f"Generated SQL Query: {sql_query}")