Update app.py

app.py

@@ -6,6 +6,7 @@ import matplotlib.pyplot as plt
 import seaborn as sns
 import os
 from typing import Optional, Tuple
+import re
 
 # OpenRouter API Key (Replace with yours)
 OPENROUTER_API_KEY = "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"
@@ -38,6 +39,14 @@ def fetch_schema(db_path: str) -> str:
     conn.close()
     return schema
 
+# Function: Extract SQL query from LLM response
+def extract_sql_query(response: str) -> str:
+    # Use regex to find content between ```sql and ```
+    match = re.search(r"```sql(.*?)```", response, re.DOTALL)
+    if match:
+        return match.group(1).strip()  # Extract and return the SQL query
+    return response  # Fallback: return the entire response if no SQL block is found
+
 # Function: Convert text to SQL
 def text_to_sql(query: str, schema: str) -> str:
     prompt = (
@@ -52,11 +61,12 @@ def text_to_sql(query: str, schema: str) -> str:
             model=OPENROUTER_MODEL,
             messages=[{"role": "system", "content": "You are an SQL expert."}, {"role": "user", "content": prompt}]
         )
-        sql_query = response.choices[0].message.content.strip()
-        return sql_query
+        sql_response = response.choices[0].message.content.strip()
+        return extract_sql_query(sql_response)  # Extract SQL query from the response
     except Exception as e:
         return f"Error: {e}"
 
+
 # Function: Execute SQL on SQLite database
 def execute_sql(sql_query: str) -> Tuple[Optional[pd.DataFrame], Optional[str]]:
     try: