Update app.py

app.py

@@ -44,12 +44,12 @@ def get_sql_query(description):
                   {"role": "user", "content": f'{prompt}'},
                 ]
       )
+      sql_query = completion.choices[0].message.content.strip().split('```sql')[1].split('```')[0]
 
     except Exception as e:
      print(f'The following error ocurred: {e}\n')
-     pass
+     sql_query = None
 
-    sql_query = completion.choices[0].message.content.strip().split('```sql')[1].split('```')[0]
     return sql_query
 
 schema = fetch_table_schema(project_id, dataset_id, table_id)
@@ -59,16 +59,18 @@ def execute_sql_query(query):
 
     try:
      result = client.query(query).to_dataframe()
-     message = f'The query : {query}\n was successfully executed and returned the above result.\n'
+     message = f'The query:{query} was successfully executed.'
 
     except Exception as e:
-     result = 'No output returned'
-     message = f'The query : {query}\n could not be executed due to exception {e}\n'
+     result = None
+     message = f'The query:{query} could not be executed due to the following exception:\n{e}'
 
     return result, message
 
 def echo(text):
   query = get_sql_query(text)
+  if query is None:
+    return 'No query generated', 'No query generated'
   result, message = execute_sql_query(query)
   return result, message
 
@@ -79,6 +81,13 @@ def gradio_interface(text):
     else:
         return result, message
 
+def gradio_interface(text):
+    result, message = echo(text)
+    if isinstance(result, pd.DataFrame):
+        return gr.Dataframe(value=result), message
+    else:
+        return result, message
+
 demo = gr.Blocks(
         title="Text-to-SQL",
         theme='remilia/ghostly',
@@ -101,7 +110,7 @@ with demo:
   with gr.Row():
     with gr.Column(scale=1):
       text_input = gr.Textbox(label="Enter your query")
-      button = gr.Button("Submit")  
+      button = gr.Button("Submit")
       gr.Examples([
         'Find the correlation between RTT and Jitter for each Market',
         'Find the variance in Jitter for each 5G_Reliability_Category',
@@ -117,10 +126,10 @@ with demo:
 
     def update_output(text):
         result, message = gradio_interface(text)
-        if isinstance(result, pd.DataFrame):
-            return gr.update(visible=True), result, message
+        if result and isinstance(result, pd.DataFrame):
+            return result, message, gr.update(visible=True)
         else:
-            return gr.update(visible=False), result, message
+            return result, message, gr.update(visible=False)
 
     button.click(update_output, inputs=text_input, outputs=[output_df, output_text])