Update app.py

app.py

@@ -61,11 +61,13 @@ def sql_engine(query: str) -> list:
     try:
         with engine.connect() as con:
             rows = con.execute(text(query))
-            return [tuple(row) for row in rows] or []
+            columns = rows.keys()
+            data = [tuple(row) for row in rows]
+            return {"columns": columns, "data": data} if data else {"columns": [], "data": []}
     except Exception as e:
-        return [f"⚠️ SQL Error: {str(e)}"]
+        return {"columns": ["Error"], "data": [[str(e)]]}
 
-# Table descriptions
+# --- Table descriptions ---
 inspector = inspect(engine)
 updated_description = "Ask natural language questions about the student marks database.\n"
 for table in ["students", "marks"]:
@@ -87,17 +89,15 @@ agent = CodeAgent(
 )
 
 # --- Gradio functions ---
-def ask_agent(question: str) -> str:
-    try:
-        result = agent.run(question)
-        if isinstance(result, list):
-            df = pd.DataFrame(result)
-            return df.to_string(index=False)
-        return str(result)
-    except Exception as e:
-        return f"⚠️ Error: {str(e)}"
-
-# Sample prompts
+def ask_agent(question: str):
+    result = agent.run(question)
+    if isinstance(result, dict) and "columns" in result and "data" in result:
+        if not result["data"]:
+            return pd.DataFrame(columns=result["columns"])
+        return pd.DataFrame(result["data"], columns=result["columns"])
+    return pd.DataFrame([[str(result)]], columns=["Result"])
+
+# --- Sample prompts ---
 sample_prompts = [
     "Which student has the highest marks in Math?",
     "List all students with their subjects and grades.",
@@ -108,7 +108,7 @@ sample_prompts = [
 def set_prompt(prompt):
     return prompt
 
-# Display database tables
+# --- Display database tables ---
 def get_tables():
     with engine.connect() as con:
         students_df = pd.read_sql_table("students", con)
@@ -119,18 +119,19 @@ students_df, marks_df = get_tables()
 
 # --- Gradio App ---
 with gr.Blocks() as demo:
-    gr.Markdown("## Text-TO-SQL Agent")
+    gr.Markdown("## 🎓 Student Marks Database Explorer")
+    
     with gr.Row():
-        gr.Column():
+        with gr.Column():
             gr.Markdown("### Students Table")
             gr.Dataframe(students_df, interactive=False)
-        gr.Column():
+        with gr.Column():
             gr.Markdown("### Marks Table")
             gr.Dataframe(marks_df, interactive=False)
 
     gr.Markdown("### Ask the AI Agent")
     question = gr.Textbox(label="Ask a question")
-    output = gr.Textbox(label="Agent response")
+    output = gr.Dataframe(headers=None)  # Result displayed as DataFrame
     ask_btn = gr.Button("Ask")
     ask_btn.click(ask_agent, inputs=question, outputs=output)