feat!: add visualizations to UI

.env.sample

@@ -7,6 +7,7 @@ GEMINI_API_KEY=
 GEMINI_MODEL=
 
 PANDAS_KEY=
+PANDAS_EXPORTS_PATH=
 
 OPENAI_MODEL=
 OPENAI_API_KEY=

gradio_app.py

@@ -3,6 +3,9 @@ from pathlib import Path
 import gradio as gr
 import asyncio
 from langchain_mcp_client import lc_mcp_exec
+from dotenv import load_dotenv
+import os
+import base64
 
 # ======================================= Load DB configs
 def load_db_configs():
@@ -18,23 +21,49 @@ def load_db_configs():
     return configs["db_configs"]
 
 
+def image_to_base64_markdown(image_path, alt_text="Customer Status"):
+    with open(image_path, "rb") as f:
+        encoded = base64.b64encode(f.read()).decode("utf-8")
+        return f"![{alt_text}](data:image/png;base64,{encoded})"
+
+
 # ====================================== 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, message = await lc_mcp_exec(request, history)
-    return response
+    response, messages = await lc_mcp_exec(request, history)
+
+    image_path = ""
+    load_dotenv()
+    PANDAS_EXPORTS_PATH = os.getenv("PANDAS_EXPORTS_PATH")
+    generated_files = [f for f in os.listdir(PANDAS_EXPORTS_PATH) if f.beginswith("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)
+        
+        # Remove the image file after reading it
+        os.remove(image_path)
+    else:
+        image_path = f"{PANDAS_EXPORTS_PATH}/blank_chart.png"
+        image_markdown = image_to_base64_markdown(image_path)
+
+    print(f"Image path: {image_path}")
+
+    output = image_markdown + response
+    
+    # print(output)
+    
+    return output
 
 
 # ====================================== Gradio UI with history
 demo = gr.ChatInterface(
-    run_agent,
+    fn=run_agent,
     type="messages",
     flagging_mode="manual",
     flagging_options=["Like", "Spam", "Inappropriate", "Other"],
-    # additional_outputs=[gr.Image(label="Visualization Output")],
     examples=[
         "List all tables in the database",
         "Show me the schema of the dim_customer table",
@@ -42,7 +71,6 @@ demo = gr.ChatInterface(
         "Create a pie chart showing the distribution of customer statuses in the dim_customer table.",
         "Plot a line chart showing the trend of order quantities over order dates from the dim_product_order_item table.",
         "Generate a bar chart displaying the count of products in each product class from the dim_product table. Use beautiful and vivid colors fro visualization.",
-        "Visualize the relationship between tax_rate and tax_included_amount from the dim_product table using a scatter plot.",
         "Visualize the number of product orders over time from the dim_product_order_item table, using the order_date. Show the trend monthly.",
     ],
     save_history=True,

langchain_mcp_client.py

@@ -101,7 +101,7 @@ async def lc_mcp_exec(request: str, history=None) -> Tuple[str, list]:
 
                 # Return the response and the updated history
                 return response_content, message_history.messages
-
+            
     except Exception as e:
         logger.error(f"Error in execution: {str(e)}", exc_info=True)
         return f"Error: {str(e)}", message_history.messages if 'message_history' in locals() else []

postgre_mcp_server.py

@@ -150,13 +150,9 @@ You can use the following FastMCP tools to create **read-only** queries (e.g., `
 
 Present your final answer using the following structure **exactly** in markdown language. When necessary, bold the important parts of your answer or use `` for inline code blocks:
 
-```markdown
 # Result
 {{Take the result from the execute_query tool and format it nicely using Markdown. Use a Markdown table for tabular data (rows and columns) including headers. Use bullet points or items in markdown for answers that include lists of names or descriptions. Use plain text for single values or simple messages. Ensure data alignment and clarity.}}
 
-# Visualization (if requested)
-{{If the user requested a visualization, include the result from the visualize_results tool, e.g., "Visualization saved as visualization_output.png". Otherwise, omit this section.}}
-
 # Explanation
 {{Provide a concise explanation or interpretation of the results (and visualization, if applicable) in 1-3 sentences. Explain what the data and visualization (if any) represent in the context of the user's request.}}
 
@@ -165,7 +161,9 @@ Present your final answer using the following structure **exactly** in markdown
 {{Display the exact SQL query you generated and executed here to answer the user's request.}}
 ```
 
-**Reminder:**
+==========================
+# Reminder
+==========================
 - **Every time you generate a SQL query, call `execute_query` immediately and include the result.**
 - **If the user requests a visualization (e.g., "create a chart", "visualize", "plot"), call `visualize_results` with:**
   - A visualization prompt starting with "plot" (e.g., "plot a bar chart showing sales by region").
@@ -611,37 +609,30 @@ async def visualize_results(json_data: dict, vis_prompt: str) -> str:
     """
     try:
         # Debug prints to see what's being received
-        print("\nVisualization Tool Debug:")
-        print(f"Received json_data: {json_data}")
-        print(f"Received vis_prompt: {vis_prompt}")
 
         # Convert JSON to DataFrame
         df = pd.DataFrame(json_data["data"], columns=json_data["columns"])
-        print(f"Created DataFrame:\n{df.head()}")
 
         # Initialize PandasAI
         df_ai = pai.DataFrame(df)
-        print("Initialized PandasAI DataFrame")
 
+        # Load api key
         load_dotenv()
         api_key = os.environ.get("PANDAS_KEY")
-        print(f"Using PandasAI API key: {api_key[:5]}...")
         pai.api_key.set(api_key)
 
         # Generate visualization
-        print(f"Attempting to generate visualization with prompt: '{vis_prompt}'")
         df_ai.chat(vis_prompt)
 
-        # Save plot
-        output_file = "visualization_output.png"
-        plt.savefig(output_file)
-        plt.close()
-        print(f"Saved visualization to {output_file}")
+        # Get the visualization path
+        PANDAS_EXPORTS_PATH = os.getenv("PANDAS_EXPORTS_PATH")
+        generated_files = [f for f in os.listdir(PANDAS_EXPORTS_PATH) if f.startswith("temp_chart")]
+        if generated_files:
+            visualization_path = os.path.join(PANDAS_EXPORTS_PATH, generated_files[0])
+
+        return f"Visualization saved as {visualization_path}"
 
-        return f"Visualization saved as {output_file}"
     except Exception as e:
-        print(f"Visualization error: {str(e)}")
-        print(f"Error type: {type(e)}")
         return f"Visualization error: {str(e)}"