Update app.py

app.py

@@ -1,6 +1,5 @@
 import gradio as gr
 import pandas as pd
-import plotly.express as px
 import os
 import io
 
@@ -10,11 +9,11 @@ from langchain.prompts import PromptTemplate
 from langchain.agents.agent_types import AgentType
 from langchain_experimental.agents import create_pandas_dataframe_agent
 
-# Load API key from environment
+# Load API key
 OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
 llm = OpenAI(temperature=0, openai_api_key=OPENAI_API_KEY)
 
-# Global storage
+# Global data stores
 dataset_dict = {}
 agent_dict = {}
 
@@ -76,8 +75,9 @@ Value counts: {value_counts}
         "value_counts": value_counts_str
     })
 
-def chat_with_csv(dataset_name, query):
+def ask_csv(dataset_name, user_input, history):
     df = dataset_dict[dataset_name]
+
     if dataset_name not in agent_dict:
         agent_dict[dataset_name] = create_pandas_dataframe_agent(
             llm=llm,
@@ -87,8 +87,9 @@ def chat_with_csv(dataset_name, query):
             handle_parsing_errors=True,
             allow_dangerous_code=True
         )
+
     agent = agent_dict[dataset_name]
-    raw_output = agent.run(query)
+    raw_output = agent.run(user_input)
 
     refine_prompt = PromptTemplate(
         input_variables=["question", "raw_output"],
@@ -105,22 +106,18 @@ Final response format:
 - Markdown-friendly format
 """
     )
+
     refine_chain = LLMChain(llm=llm, prompt=refine_prompt)
-    return refine_chain.run({
-        "question": query,
+    final_response = refine_chain.run({
+        "question": user_input,
         "raw_output": raw_output
     })
 
-def generate_plot(dataset_name, x_col, y_col, color_col):
-    df = dataset_dict[dataset_name]
-    fig = px.scatter(
-        df, x=x_col, y=y_col,
-        color=color_col if color_col and color_col != "None" else None
-    )
-    return fig
+    history.append((user_input, final_response))
+    return history, history
 
 # ─────────────────────────────────────────────
-# 🎛 Gradio UI
+# 🎛 Gradio App
 # ─────────────────────────────────────────────
 
 with gr.Blocks(title="CSV Chat Assistant") as app:
@@ -131,19 +128,18 @@ with gr.Blocks(title="CSV Chat Assistant") as app:
         upload_output = gr.Textbox(label="Upload Status")
 
     dataset_radio = gr.Radio(choices=[], label="Select Dataset", interactive=True)
-
-    # ✅ Ensures radio options are updated properly
     file_upload.change(fn=upload_csv, inputs=[file_upload], outputs=[upload_output, dataset_radio])
 
-    with gr.Accordion("🧠 Dataset Summary", open=False):
-        summary_button = gr.Button("Generate Summary")
+    with gr.Row():
+        summary_button = gr.Button("🧾 Generate Summary")
         summary_output = gr.Markdown()
         summary_button.click(fn=generate_summary, inputs=dataset_radio, outputs=summary_output)
 
-    with gr.Accordion("💬 Ask a Question", open=True):
-        query_input = gr.Textbox(lines=2, label="Ask your question")
-        query_button = gr.Button("Get Answer")
-        query_output = gr.Markdown()
-        query_button.click(fn=chat_with_csv, inputs=[dataset_radio, query_input], outputs=[query_output])
+    chatbot = gr.Chatbot(label="CSV Chat", height=400)
+    msg = gr.Textbox(label="Type your question")
+    clear = gr.Button("Clear Chat")
+
+    msg.submit(fn=ask_csv, inputs=[dataset_radio, msg, chatbot], outputs=[chatbot, chatbot])
+    clear.click(lambda: [], None, chatbot)
 
 app.launch()