Update app.py

app.py

@@ -1,15 +1,15 @@
 import os
 import gradio as gr
 from sqlalchemy import text
-from smolagents import tool, CodeAgent, HfApiModel
+from smolagents import CodeAgent, HfApiModel
 import pandas as pd
 from io import StringIO
+import tempfile
 from database import (
     engine,
     create_dynamic_table,
     clear_database,
-    insert_rows_into_table,
-    get_table_schema
+    insert_rows_into_table
 )
 
 # Initialize the AI agent
@@ -47,7 +47,6 @@ def process_txt_file(file_path):
         with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
             content = f.read()
         
-        # Structure detection with enhanced prompts
         structure_prompt = f"""
         Convert this text into valid CSV format:
         {content}
@@ -62,7 +61,6 @@ def process_txt_file(file_path):
         """
         csv_output = agent.run(structure_prompt)
         
-        # Robust CSV parsing
         try:
             df = pd.read_csv(
                 StringIO(csv_output),
@@ -76,7 +74,6 @@ def process_txt_file(file_path):
         if df.empty or len(df.columns) == 0:
             return False, "No structured data found", pd.DataFrame()
 
-        # Database operations
         clear_database()
         table = create_dynamic_table(df)
         insert_rows_into_table(df.to_dict('records'), table)
@@ -89,20 +86,34 @@ def process_txt_file(file_path):
 def handle_upload(file_obj):
     """Handle file upload and processing"""
     if file_obj is None:
-        return "Please upload a text file.", None, "No schema", gr.update(visible=True), gr.update(visible=False)
+        return [
+            "Please upload a text file.", 
+            None, 
+            "No schema", 
+            gr.update(visible=True), 
+            gr.update(visible=False),
+            gr.update(visible=False)
+        ]
     
     success, message, df = process_txt_file(file_obj)
     if success:
-        column_info = {col: {'type': 'text'} for col in df.columns}
         schema = "\n".join([f"- {col} (text)" for col in df.columns])
-        return (
+        return [
             message,
             df,
             f"### Detected Schema:\n```\n{schema}\n```",
             gr.update(visible=False),
+            gr.update(visible=True),
             gr.update(visible=True)
-        )
-    return message, None, "No schema", gr.update(visible=True), gr.update(visible=False)
+        ]
+    return [
+        message, 
+        None, 
+        "No schema", 
+        gr.update(visible=True), 
+        gr.update(visible=False),
+        gr.update(visible=False)
+    ]
 
 def query_analysis(user_query: str) -> str:
     """Handle natural language queries about the data"""
@@ -130,6 +141,16 @@ def query_analysis(user_query: str) -> str:
     except Exception as e:
         return f"Query error: {str(e)}"
 
+def download_csv():
+    """Generate CSV file for download"""
+    df = get_data_table()
+    if not df.empty:
+        temp_dir = tempfile.gettempdir()
+        file_path = os.path.join(temp_dir, "processed_data.csv")
+        df.to_csv(file_path, index=False)
+        return file_path
+    return None
+
 # Gradio interface setup
 with gr.Blocks() as demo:
     with gr.Group() as upload_group:
@@ -147,7 +168,9 @@ with gr.Blocks() as demo:
     with gr.Group(visible=False) as query_group:
         with gr.Row():
             with gr.Column(scale=1):
-                user_input = gr.Textbox(label="Ask about the data")
+                with gr.Row():
+                    user_input = gr.Textbox(label="Ask about the data", scale=4)
+                    submit_btn = gr.Button("Submit", scale=1)
                 query_output = gr.Markdown(label="Analysis Results")
             with gr.Column(scale=2):
                 gr.Markdown("### Extracted Data Preview")
@@ -155,6 +178,10 @@ with gr.Blocks() as demo:
                     label="Structured Data",
                     interactive=False
                 )
+                download_btn = gr.DownloadButton(
+                    "Download as CSV",
+                    visible=False
+                )
         schema_display = gr.Markdown()
         refresh_btn = gr.Button("Refresh View")
 
@@ -162,9 +189,15 @@ with gr.Blocks() as demo:
     file_input.upload(
         fn=handle_upload,
         inputs=file_input,
-        outputs=[status, data_table, schema_display, upload_group, query_group]
+        outputs=[status, data_table, schema_display, upload_group, query_group, download_btn]
     )
 
+    submit_btn.click(
+        fn=query_analysis,
+        inputs=user_input,
+        outputs=query_output
+    )
+    
     user_input.submit(
         fn=query_analysis,
         inputs=user_input,
@@ -175,6 +208,11 @@ with gr.Blocks() as demo:
         fn=lambda: (get_data_table().head(10), "Schema refreshed"),
         outputs=[data_table, schema_display]
     )
+    
+    download_btn.click(
+        fn=download_csv,
+        outputs=download_btn
+    )
 
 if __name__ == "__main__":
     demo.launch(