Update app.py

app.py

@@ -87,12 +87,12 @@ def generate_generic_table_data(table_name, row_count=15):
             'position': [random.choice(['Engineer', 'Manager', 'Analyst', 'Developer', 'Designer']) 
                         for _ in range(row_count)]
         },
-        'departments': {
+        'departments': lambda: {
             'id': list(range(1, 6)),
             'name': ['Engineering', 'Sales', 'Marketing', 'HR', 'Finance'],
             'manager_id': [random.randint(1, 15) for _ in range(5)],
             'budget': [random.randint(100000, 1000000) for _ in range(5)]
-        }[:5],
+        },
         'books': {
             'book_id': gen_id(),
             'title': [f"Book Title {i}" for i in range(1, row_count + 1)],
@@ -230,11 +230,15 @@ def generate_generic_table_data(table_name, row_count=15):
     }
     
     # Return predefined schema if exists, otherwise create generic one
-    if table_name.lower() in table_schemas:
-        return table_schemas[table_name.lower()]
+    table_lower = table_name.lower()
+    if table_lower in table_schemas:
+        schema = table_schemas[table_lower]
+        # If it's a callable (lambda), execute it
+        if callable(schema):
+            return schema()
+        return schema
     
     # Generic fallback for unknown tables
-    # Try to infer structure from table name
     generic_data = {
         f'{table_name}_id': gen_id(),
         'name': gen_names(),
@@ -257,11 +261,11 @@ def create_database_from_tables(tables_used):
         table_name = table.lower().strip()
         
         # Generate appropriate sample data
-        table_dict = generate_generic_table_data(table_name, row_count=15)
-        
-        # Adjust row count for lookup tables
-        if table_name in ['departments']:
-            table_dict = {k: v[:5] if isinstance(v, list) else v for k, v in table_dict.items()}
+        # Special handling for departments (only 5 rows)
+        if table_name == 'departments':
+            table_dict = generate_generic_table_data(table_name, row_count=5)
+        else:
+            table_dict = generate_generic_table_data(table_name, row_count=15)
         
         df = pd.DataFrame(table_dict)
         df.to_sql(table_name, conn, index=False, if_exists='replace')
@@ -371,7 +375,8 @@ def process_nl_query(api_key, natural_query):
         # Display sample tables (show first 10 rows for readability)
         for table_name, df in sample_data.items():
             output_text += f"**📊 Sample `{table_name}` Table** ({len(df)} rows):\n\n"
-            output_text += df.head(10).to_markdown(index=False)
+            display_df = df.head(10)
+            output_text += display_df.to_markdown(index=False)
             if len(df) > 10:
                 output_text += f"\n\n*...and {len(df) - 10} more rows*"
             output_text += "\n\n"
@@ -406,7 +411,7 @@ def process_nl_query(api_key, natural_query):
         return error_msg, "", pd.DataFrame({"Error": [str(e)]}), ""
 
 # Create Gradio Interface
-with gr.Blocks(title="Natural Language to SQL Query Executor", theme=gr.themes.Ocean()) as demo:
+with gr.Blocks(title="Natural Language to SQL Query Executor", theme=gr.themes.Soft()) as demo:
     gr.Markdown("""
     # 🔍 Natural Language to SQL Query Executor