Update app.py

app.py

@@ -12,9 +12,11 @@ from database import (
     get_table_schema
 )
 
-def get_data_table():
+def get_table_info():
     """
-    Fetches all data from the current table and returns it as a Pandas DataFrame.
+    Gets the current table name and column information.
+    Returns:
+        tuple: (table_name, list of column names, column info)
     """
     try:
         # Get list of tables
@@ -24,24 +26,30 @@ def get_data_table():
             )).fetchall()
             
         if not tables:
-            return pd.DataFrame()
+            return None, [], {}
             
         # Use the first table found
         table_name = tables[0][0]
-            
+        
+        # Get column information
         with engine.connect() as con:
-            result = con.execute(text(f"SELECT * FROM {table_name}"))
-            rows = result.fetchall()
+            columns = con.execute(text(f"PRAGMA table_info({table_name})")).fetchall()
             
-            if not rows:
-                return pd.DataFrame()
-                
-            columns = result.keys()
-            df = pd.DataFrame(rows, columns=columns)
-            return df
-
+        # Extract column names and types
+        column_names = [col[1] for col in columns]
+        column_info = {
+            col[1]: {
+                'type': col[2],
+                'is_primary': bool(col[5])
+            }
+            for col in columns
+        }
+        
+        return table_name, column_names, column_info
+        
     except Exception as e:
-        return pd.DataFrame({"Error": [str(e)]})
+        print(f"Error getting table info: {str(e)}")
+        return None, [], {}
 
 def process_sql_file(file_path):
     """
@@ -152,14 +160,22 @@ def query_sql(user_query: str) -> str:
     """
     Converts natural language input to an SQL query using CodeAgent.
     """
-    schema = get_table_schema()
-    if not schema:
+    # Get current table information
+    table_name, column_names, column_info = get_table_info()
+    
+    if not table_name:
         return "Error: No data table exists. Please upload a file first."
-        
+    
+    # Create schema information with actual column names
     schema_info = (
-        f"The database has the following schema:\n"
-        f"{schema}\n"
-        "Generate a valid SQL SELECT query using ONLY these column names.\n"
+        f"The database has a table named '{table_name}' with the following columns:\n"
+        + "\n".join([
+            f"- {col} ({info['type']}){' primary key' if info['is_primary'] else ''}"
+            for col, info in column_info.items()
+        ])
+        + "\n\nGenerate a valid SQL SELECT query using ONLY these column names.\n"
+        "The table name is '" + table_name + "'.\n"
+        "If column names contain spaces, they must be quoted.\n"
         "DO NOT explain your reasoning, and DO NOT return anything other than the SQL query itself."
     )
 
@@ -171,6 +187,17 @@ def query_sql(user_query: str) -> str:
     if not generated_sql.strip().lower().startswith(("select", "show", "pragma")):
         return "Error: Only SELECT queries are allowed."
 
+    # Fix table names
+    for wrong_name in ['table_name', 'customers', 'main']:
+        if wrong_name in generated_sql:
+            generated_sql = generated_sql.replace(wrong_name, table_name)
+    
+    # Add quotes around column names that need them
+    for col in column_names:
+        if ' ' in col:  # If column name contains spaces
+            if col in generated_sql and f'"{col}"' not in generated_sql and f'`{col}`' not in generated_sql:
+                generated_sql = generated_sql.replace(col, f'"{col}"')
+
     result = sql_engine(generated_sql)
     
     try: