Chatbot.py

# ── Imports ───────────────────────────────────────────────────────────────────
import os           # used by atexit to delete temp files
import atexit       # registers cleanup callbacks that run when the app process exits
import tempfile     # creates temporary files on disk for uploaded CSVs
import streamlit as st              # Streamlit UI framework
import csv_agent, plotter           # local modules: agent logic and output formatting

# ── Page Configuration ────────────────────────────────────────────────────────
# Must be the first Streamlit call in the script.
# layout="wide" uses the full browser width instead of the default narrow column.
st.set_page_config(layout="wide")
st.title("🤖📊 CSV Agent Chatbot ")

# ── Session State: Chat History ───────────────────────────────────────────────
# st.session_state persists values across Streamlit reruns (each user interaction
# triggers a full script rerun). Chat history is initialised once per session.
if 'chat_history' not in st.session_state:
    st.session_state.chat_history = []

# ── Render Chat History ───────────────────────────────────────────────────────
# On every rerun, replay all previous messages so the conversation stays visible.
# Assistant messages may include an HTML chart/table in a collapsible expander.
for chats in st.session_state.chat_history:
    with st.chat_message(chats["role"]):
        st.markdown(chats['content'])
        if chats["role"] == "Assistant":
            if chats['html_content'] != "":
                with st.expander("📈📝 See explanation "):
                    # st.iframe renders the HTML string in a sandboxed iframe
                    st.iframe(chats['html_content'], height=600)

# ── Sidebar: CSV File Upload ──────────────────────────────────────────────────
st.sidebar.header("📁 Upload an CSV File ")

# label_visibility="collapsed" hides the label visually while still satisfying
# Streamlit's requirement for a non-empty label string.
uploaded_file = st.sidebar.file_uploader("Upload CSV", type=["csv"], label_visibility="collapsed")

temp_file_path = ''
if uploaded_file is not None:
    try:
        # Write the uploaded file to a temporary path on disk.
        # LangChain's create_csv_agent requires a file path, not a file object,
        # so we can't pass the in-memory upload directly.
        # delete=False keeps the file alive after the with-block closes.
        with tempfile.NamedTemporaryFile(delete=False, suffix=".csv") as tmp_file:
            tmp_file.write(uploaded_file.read())
            temp_file_path = tmp_file.name

        # Register cleanup so the temp file is deleted when the Streamlit
        # process exits — prevents stale CSV files accumulating on disk.
        atexit.register(os.unlink, temp_file_path)

        st.sidebar.write("📁 File Upload status")
        st.sidebar.code("SUCCESSFUL")
    except Exception as e:
        st.error(f"❌ Error uploading .csv file: {e}")
        st.sidebar.write("📁 File Upload status")
        st.sidebar.code("FAILED")

# ── Agent Cache ───────────────────────────────────────────────────────────────
# @st.cache_resource caches the return value keyed by file_path.
# This means build_csv_agent() is called only once per uploaded file —
# subsequent queries reuse the same agent instead of rebuilding it each time,
# which avoids re-parsing the CSV and re-initialising the LangChain chain.
@st.cache_resource
def get_csv_agent(file_path):
    return csv_agent.build_csv_agent(file_path)

# ── Chat Input & Response ─────────────────────────────────────────────────────
# st.chat_input renders a persistent input bar at the bottom of the page.
# The walrus operator (:=) assigns the value and checks it's non-empty in one step.
if user_input := st.chat_input("Ask a question about your CSV data..."):

    # Display and store the user's message
    with st.chat_message("User"):
        st.markdown(user_input)
        st.session_state.chat_history.append({"role": "User", "content": user_input})

    # ── Guard: require a CSV to be uploaded first ─────────────────────────────
    if temp_file_path == '':
        with st.chat_message("Assistant"):
            st.markdown("Please upload a .csv file for Q&A")
            st.session_state.chat_history.append({
                "role": "Assistant",
                "content": "Please upload a .csv file for Q&A",
                "html_content": ""
            })

    else:
        # ── Agent Pipeline ────────────────────────────────────────────────────
        # Step 1: get (or build) the cached CSV agent for the uploaded file
        # Step 2: run the user's question through the agent → raw text answer
        # Step 3: pass the answer to plotter.output_formatter → (html, summary)
        try:
            agent = get_csv_agent(temp_file_path)
            csv_agent_response     = csv_agent.csv_agent_invoker(agent, user_input)
            html_content, response = plotter.output_formatter(user_input, csv_agent_response)
        except Exception as e:
            # Catch any agent or formatting errors and show a friendly message
            response     = f"Sorry, something went wrong: {e}"
            html_content = ""

        # ── Display and Store Assistant Response ──────────────────────────────
        with st.chat_message("Assistant"):
            st.markdown(response)
            # Only show the chart/table expander if the plotter returned HTML
            if html_content != "":
                with st.expander("📈📝 See explanation "):
                    st.iframe(html_content, height=600)
            st.session_state.chat_history.append({
                "role": "Assistant",
                "content": response,
                "html_content": html_content
            })