Update app.py

app.py

@@ -1,13 +1,12 @@
-import os
 import json
 import gradio as gr
 import duckdb
 
-# Load the Parquet dataset
-dataset_path = 'sample_contract_df.parquet'  # Update with your Parquet file's path
+# Load the Parquet dataset path
+dataset_path = 'sample_contract_df.parquet'  # Update with your Parquet file path
 
-# Load the dataset with DuckDB and get schema information
-def load_dataset():
+# Load the dataset schema with DuckDB
+def load_dataset_schema():
     con = duckdb.connect()
     con.execute(f"CREATE VIEW contract_data AS SELECT * FROM '{dataset_path}'")
     schema = con.execute("DESCRIBE contract_data").fetchdf()
@@ -16,23 +15,8 @@ def load_dataset():
 
 # Generate SQL based on schema and user query
 def generate_sql_query(features, query):
-    messages = [
-        {
-            "role": "system",
-            "content": "You are a SQL query expert assistant that generates DuckDB SQL queries based on the user's natural language query and dataset schema.",
-        },
-        {
-            "role": "user",
-            "content": f"""table contract_data
-# Features
-{features}
-# Query
-{query}
-""",
-        },
-    ]
-    # Here we use DuckDB directly instead of an external API
-    sql_query = f"SELECT * FROM contract_data WHERE {query}"  # Simple example; adapt for complex queries
+    # This simple example constructs an SQL condition directly, adapt for complex queries if needed
+    sql_query = f"SELECT * FROM contract_data WHERE {query}"
     return sql_query
 
 # Execute the SQL query and display results
@@ -41,38 +25,35 @@ def execute_query(sql_query):
     con.execute(f"CREATE VIEW contract_data AS SELECT * FROM '{dataset_path}'")
     result_df = con.execute(sql_query).fetchdf()
     con.close()
-    return result_df.to_markdown()  # Convert result to markdown for display
+    return result_df.to_markdown()  # Display result as markdown
 
 # Gradio app UI
 with gr.Blocks() as demo:
     gr.Markdown("""
     # Local Parquet SQL Query App
-    Query and explore the data stored in `sample_contract_df.parquet` using DuckDB and SQL queries.
+    Query and explore data in `sample_contract_df.parquet` using DuckDB and SQL queries.
     """)
     
-    # Display schema
-    schema = load_dataset()
-    features = json.dumps(schema, indent=2)
-    gr.Markdown(f"### Dataset Schema:\n\n```json\n{features}\n```")
+    # Load dataset schema and convert it to JSON for display
+    schema = load_dataset_schema()
+    schema_json = json.dumps(schema, indent=2)
+    features = gr.Code(label="Dataset Schema", value=schema_json, language="json")
 
-    # User inputs for natural language query
+    # User input components
     query = gr.Textbox(label="Natural Language Query", placeholder="Enter a condition, e.g., 'amount > 1000'")
     sql_out = gr.Code(label="Generated SQL Query", language="sql")
     results_out = gr.Markdown(label="Query Results")
 
     # Buttons to generate and execute SQL
-    with gr.Row():
-        btn_generate = gr.Button("Generate SQL")
-        btn_execute = gr.Button("Execute Query")
+    btn_generate = gr.Button("Generate SQL")
+    btn_execute = gr.Button("Execute Query")
 
-    # Generate SQL on button click
+    # Set up click events with correct inputs and outputs
     btn_generate.click(
         fn=generate_sql_query,
-        inputs=[features, query],
+        inputs=[features, query],  # Both are Gradio components now
         outputs=sql_out,
     )
-
-    # Execute SQL on button click
     btn_execute.click(
         fn=execute_query,
         inputs=sql_out,