feat: new gradio UI

gradio_app.py

@@ -29,48 +29,94 @@ def image_to_base64_markdown(image_path, alt_text="Customer Status"):
 
 # ====================================== Async-compatible wrapper
 async def run_agent(request, history):
-    # configs = load_db_configs()
-    # final_answer, last_tool_answer, = await pg_mcp_exec(request)
-    # return final_answer, last_tool_answer
-
     response, messages = await lc_mcp_exec(request, history)
 
     image_path = ""
     load_dotenv()
-    PANDAS_EXPORTS_PATH = os.getenv("PANDAS_EXPORTS_PATH")
+    PANDAS_EXPORTS_PATH = os.getenv("PANDAS_EXPORTS_PATH", "exports/charts")
+    
+    # Ensure the exports directory exists
+    os.makedirs(PANDAS_EXPORTS_PATH, exist_ok=True)
+    
     generated_files = [f for f in os.listdir(PANDAS_EXPORTS_PATH) if f.startswith("temp_chart_") and f.endswith(".png")]
     if generated_files:
         image_path = os.path.join(PANDAS_EXPORTS_PATH, generated_files[0])
-        image_markdown = image_to_base64_markdown(image_path)
-        output = image_markdown + response
-        # Remove the image file after reading it
-        os.remove(image_path)
+        try:
+            image_markdown = image_to_base64_markdown(image_path)
+            output = f"{image_markdown}\n\n{response}"
+            # Remove the image file after reading it
+            os.remove(image_path)
+        except Exception as e:
+            print(f"Error processing image: {e}")
+            output = response
     else:
-        # image_path = f"{PANDAS_EXPORTS_PATH}/blank_chart.png"
-        # image_markdown = image_to_base64_markdown(image_path)
         output = response
 
     print(f"Image path: {image_path}")
-    
-    # print(output)
+    print(f"Output length: {len(output)}")
     
     return output
 
 
 # ====================================== Gradio UI with history
-LOGO_PATH = "resources/logo.png"
-with gr.Blocks() as demo:
-    # gr.Image(value=LOGO_PATH, height=100, width=100, elem_id="logo")
-    with gr.Row():
-        with gr.Column():
-            gr.Markdown("<center><h1>Talk to Your Data</h1></center>")
-            gr.ChatInterface(
+LOGO_PATH = "resources/pialogo.png"
+
+# CSS customizations
+custom_css = """
+.container {
+    max-width: 2200px !important;
+    margin: auto;
+    padding: 20px;
+}
+.chat-container {
+    height: 1000px !important;
+    min-height: 1000px !important;
+    overflow-y: auto;
+}
+.message-container {
+    padding: 15px;
+    border-radius: 10px;
+    margin: 10px 0;
+}
+.markdown-content {
+    font-size: 18px;
+    line-height: 1.7;
+}
+"""
+
+with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
+    with gr.Row(elem_classes="container"):
+        # with gr.Column(scale=1):
+        #     gr.Image(value=LOGO_PATH, height=200, show_label=False)
+            
+        with gr.Column(scale=3):
+            gr.Markdown(
+                """
+                <h1 style='text-align: center; margin-bottom: 1rem'>Talk to Your Data</h1>
+                <p style='text-align: center'>Ask questions about your database, analyze and visualize data.</p>
+                """
+            )
+    
+    with gr.Row(elem_classes="container"):
+        with gr.Column(scale=3):
+            chat = gr.ChatInterface(
                 fn=run_agent,
-                type="messages",
-                flagging_mode="manual",
-                flagging_options=["Like", "Spam", "Inappropriate", "Other"],
+                chatbot=gr.Chatbot(
+                    height=1000,
+                    show_label=False,
+                    elem_classes="chat-container",
+                    render_markdown=True
+                ),
+                textbox=gr.Textbox(
+                    placeholder="Type your questions here...",
+                    container=False,
+                    scale=4
+                ),
+                theme="soft",
                 examples=[
+                    "Describe the database",
                     "List all tables in the database",
+                    "List all tables with columns and data types",
                     "how many customers do you have?",
                     "what are the statuses my of my customers",
                     "Visualize with different colors and show legend",
@@ -80,30 +126,29 @@ with gr.Blocks() as demo:
                     "How many users and roles have been created in 2024"
                 ],
                 save_history=True,
+                type="messages"
             )
+        
+        with gr.Column(scale=1):
+            with gr.Accordion("Example Questions", open=True):
+                gr.Markdown("""
+                - 📊 List all tables in database
+                - 👥 Total number of customers
+                - 📈 Visualize it with different colors
+                - 📋 Order statistics for last 6 years
+                - 📆 User and role counts in 2024
+                """)
 
 
-# Gradio UI
-# demo = gr.Interface(
-#     fn=run_agent,
-#     inputs=gr.Textbox(
-#         label="Natural Language Request",
-#         placeholder="e.g., Show me the table of join posts and users tables."
-#     ),
-#     outputs=gr.Markdown(
-#         label="SQL Query / Result",
-#         value="# Result \n \n",
-#         # max_height=1500,
-#         height= 1500
-#     ),
-#     title="PostgreSQL Query Agent",
-#     description="Ask your database in natural language and get results.",
-#     flagging_mode="never"
-# )
 
 
 # TODO: maybe we can add a mcp tool to validate the results (those converted to DataFrame) to make sure the valid type is passed to the visualization tool by ReAct agent
 
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        share=True,
+        debug=True
+    )

langchain_mcp_client.py

@@ -12,7 +12,7 @@ import logging
 from dotenv import load_dotenv
 from langchain.globals import set_debug
 
-# set_debug(True)
+set_debug(True)
 
 
 # Set up logging