Update app.py

app.py

@@ -11,13 +11,16 @@ from e2b_code_interpreter import Sandbox
 # Configuration and Setup
 # =========================
 
+# Load environment variables
+load_dotenv()
+
 # Initialize OpenAI API key
 openai.api_key = os.getenv("OPENAI_API_KEY")
 if not openai.api_key:
     raise ValueError("Please set the OPENAI_API_KEY environment variable.")
 
-load_dotenv()
-sbx = Sandbox() # By default, the sandbox is alive for 5 minutes
+# Initialize the Sandbox
+sbx = Sandbox()  # By default, the sandbox is alive for 5 minutes
 
 # Path to your Parquet dataset
 DATASET_PATH = 'hsas.parquet'  # Update with your Parquet file path
@@ -67,8 +70,8 @@ def parse_query(nl_query):
     ]
 
     try:
-        response = openai.chat.completions.create(
-            model="gpt-4o-mini",  # Use a valid and accessible model
+        response = openai.ChatCompletion.create(
+            model="gpt-3.5-turbo",  # Use a valid and accessible model
             messages=messages,
             temperature=0,
             max_tokens=150,
@@ -122,23 +125,27 @@ def execute_sql_query(sql_query):
 with gr.Blocks(css="""
     .error-message {
         color: red;
+        font-weight: bold;
     }
     .gradio-container {
-        max-width: 1000px;
+        max-width: 1200px;
         margin: auto;
+        font-family: -apple-system, BlinkMacSystemFont, 'San Francisco', 'Helvetica Neue', Helvetica, Arial, sans-serif;
     }
     .header {
         text-align: center;
-        padding: 20px;
+        padding: 30px 0;
     }
     .instructions {
-        margin-bottom: 20px;
+        margin: 20px 0;
+        font-size: 18px;
+        line-height: 1.6;
     }
     .example-queries {
         margin-bottom: 20px;
     }
     .button-row {
-        margin-top: 10px;
+        margin-top: 20px;
     }
     .input-area {
         margin-bottom: 20px;
@@ -146,41 +153,49 @@ with gr.Blocks(css="""
     .schema-tab {
         padding: 20px;
     }
+    .results {
+        margin-top: 20px;
+    }
+    .copy-button {
+        margin-top: 10px;
+    }
 """) as demo:
-    with gr.Row():
-        gr.Markdown("# 🏥 Text-to-SQL Healthcare Data Analyst Agent", elem_classes="header")
-    
+    # Header
     gr.Markdown("""
-    # Analyze data from the U.S. Center of Medicare and Medicaid
+    # 🏥 Text-to-SQL Healthcare Data Analyst Agent
 
-    ## Demonstrate how to use AI for data analysis and visualization in an isolated sandbox
+    Analyze data from the U.S. Center of Medicare and Medicaid using natural language queries.
 
-    ## Use SQL directly in the browser to answer complex business questions starting with plain English.
+    """, elem_classes="header")
 
     # Instructions
+    gr.Markdown("""
+    ### Instructions
+
+    1. **Describe the data you want**: e.g., *"Show total days of care by zip code"*
+    2. **Use Example Queries**: Click on any example query button below to execute
+    3. **Generate SQL**: Or, enter your own query and click **Generate SQL Query**
+    4. **Execute the Query**: Review the generated SQL and click **Execute Query** to see the results
 
-    # 1.  **Describe the data you want**: e.g., `Show total days of care by zip`
-    # 2.  **Use Example Queries**: Click on any example query button below to execute
-    # 3.  **Generate SQL**: Or, enter your own query and click "Generate SQL"
     """, elem_classes="instructions")
 
     with gr.Row():
-        with gr.Column(scale=1, min_width=300):
-            gr.Markdown("# 💡 Example Queries:", elem_classes="example-queries")
+        with gr.Column(scale=1, min_width=350):
+            gr.Markdown("### 💡 Example Queries", elem_classes="example-queries")
             query_buttons = [
                 "Calculate the average total_charges by zip_cd_of_residence",
                 "For each zip_cd_of_residence, calculate the sum of total_charges",
                 "SELECT * FROM hsa_data WHERE total_days_of_care > 40 LIMIT 30;",
             ]
-            btn_queries = [gr.Button(q, variant="secondary", size="sm") for q in query_buttons]
-
-            with gr.Group():
-                query_input = gr.Textbox(
-                    label="🔍 Your Query",
-                    placeholder='e.g., "Show total charges over 1M by state"',
-                    lines=2,
-                    interactive=True,
-                )
+            btn_queries = [gr.Button(q, variant="secondary") for q in query_buttons]
+
+            query_input = gr.Textbox(
+                label="🔍 Your Query",
+                placeholder='e.g., "Show total charges over 1M by zip code"',
+                lines=2,
+                interactive=True,
+                elem_classes="input-area"
+            )
 
             with gr.Row(elem_classes="button-row"):
                 btn_generate_sql = gr.Button("Generate SQL Query", variant="primary")
@@ -189,28 +204,28 @@ with gr.Blocks(css="""
             sql_query_out = gr.Code(label="📝 Generated SQL Query", language="sql")
             error_out = gr.HTML(elem_classes="error-message", visible=False)
 
-        with gr.Column(scale=2, min_width=600):
+        with gr.Column(scale=2, min_width=650):
             gr.Markdown("### 📊 Query Results", elem_classes="results")
             results_out = gr.Dataframe(label="Query Results", interactive=False)
-            btn_copy_results = gr.Button("Copy to Clipboard", variant="secondary", size="sm")
 
-            # JavaScript for copying to clipboard - updated approach
+            # Copy to Clipboard Button
+            btn_copy_results = gr.Button("Copy Results to Clipboard", variant="secondary", elem_classes="copy-button")
+
+            # JavaScript for copying to clipboard
             copy_script = gr.HTML("""
                 <script>
                     function copyToClipboard() {
-                        const resultsContainer = document.querySelector('.dataframe-output table');
+                        const resultsContainer = document.querySelector('div[data-testid="dataframe"] table');
                         if (resultsContainer) {
                             const text = Array.from(resultsContainer.rows)
                                 .map(row => Array.from(row.cells)
-                                    .map(cell => cell.innerText).join("\t"))
-                                .join("\n");
-                            const tempTextArea = document.createElement('textarea');
-                            tempTextArea.value = text;
-                            document.body.appendChild(tempTextArea);
-                            tempTextArea.select();
-                            document.execCommand('copy');
-                            document.body.removeChild(tempTextArea);
-                            alert("Copied results to clipboard!");
+                                    .map(cell => cell.innerText).join("\\t"))
+                                .join("\\n");
+                            navigator.clipboard.writeText(text).then(function() {
+                                alert("Copied results to clipboard!");
+                            }, function(err) {
+                                alert("Failed to copy results: " + err);
+                            });
                         } else {
                             alert("No results to copy!");
                         }
@@ -220,9 +235,10 @@ with gr.Blocks(css="""
 
             # Connect JavaScript function to button
             btn_copy_results.click(
-                None, None, None, _js="copyToClipboard"
+                _js="copyToClipboard"
             )
 
+    # Dataset Schema Tab
     with gr.Tab("📋 Dataset Schema", elem_classes="schema-tab"):
         gr.Markdown("### Dataset Schema")
         schema_display = gr.JSON(label="Schema", value=get_schema())
@@ -233,19 +249,19 @@ with gr.Blocks(css="""
 
     def generate_sql(nl_query):
         if not nl_query.strip():
-            return "", "<p class='error-message'>Please enter a query.</p>", gr.update(visible=True)
+            return "", "<p>Please enter a query.</p>", gr.update(visible=True)
         sql_query, error = parse_query(nl_query)
         if error:
-            return sql_query, f"<p class='error-message'>{error}</p>", gr.update(visible=True)
+            return sql_query, f"<p>{error}</p>", gr.update(visible=True)
         else:
             return sql_query, "", gr.update(visible=False)
 
     def execute_query(sql_query):
         if not sql_query.strip():
-            return None, "<p class='error-message'>No SQL query to execute.</p>", gr.update(visible=True)
+            return None, "<p>No SQL query to execute.</p>", gr.update(visible=True)
         result_df, error = execute_sql_query(sql_query)
         if error:
-            return None, f"<p class='error-message'>{error}</p>", gr.update(visible=True)
+            return None, f"<p>{error}</p>", gr.update(visible=True)
         else:
             return result_df, "", gr.update(visible=False)
 
@@ -254,16 +270,16 @@ with gr.Blocks(css="""
             sql_query = example_query
             result_df, error = execute_sql_query(sql_query)
             if error:
-                return sql_query, f"<p class='error-message'>{error}</p>", None, gr.update(visible=True)
+                return sql_query, f"<p>{error}</p>", None, gr.update(visible=True)
             else:
-                return sql_query, "", result_df, gr.update(visible=False)
+                return sql_query, "", result_df, gr.update(visible(False))
         else:
             sql_query, error = parse_query(example_query)
             if error:
-                return sql_query, f"<p class='error-message'>{error}</p>", None, gr.update(visible=True)
+                return sql_query, f"<p>{error}</p>", None, gr.update(visible=True)
             result_df, exec_error = execute_sql_query(sql_query)
             if exec_error:
-                return sql_query, f"<p class='error-message'>{exec_error}</p>", None, gr.update(visible=True)
+                return sql_query, f"<p>{exec_error}</p>", None, gr.update(visible=True)
             else:
                 return sql_query, "", result_df, gr.update(visible=False)
 
@@ -287,8 +303,9 @@ with gr.Blocks(css="""
             outputs=[sql_query_out, error_out, results_out, error_out],
         )
 
-    query_input.change(fn=lambda: ("", gr.update(visible=False)), inputs=None, outputs=[error_out])
-    sql_query_out.change(fn=lambda: ("", gr.update(visible=False)), inputs=None, outputs=[error_out])
+    # Hide error message when inputs change
+    query_input.change(fn=lambda: gr.update(visible=False), inputs=None, outputs=[error_out])
+    sql_query_out.change(fn=lambda: gr.update(visible=False), inputs=None, outputs=[error_out])
 
 # =========================
 # Launch the Gradio App