Update app.py

app.py

@@ -5,13 +5,10 @@ import requests
 import json
 
 # --- Page config
-st.set_page_config(page_title="CSV-Backed AI Chat Agent", layout="centered")
+st.set_page_config(page_title="CSV-Backed AI Chat Agent", layout="wide")
 
-# --- Minimal title
-st.markdown(
-    "<h2 style='text-align:center; margin-bottom:1rem;'>CSV-Backed AI Chat Agent</h2>",
-    unsafe_allow_html=True,
-)
+# --- Title & image
+st.title("CSV-Backed AI Chat Agent")
 
 # --- Load API key
 OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
@@ -54,8 +51,18 @@ if uploaded_file is not None:
                 "If you use 'search_csv', use Pandas query syntax."
             ),
         }
-        # Reset memory on new file load
-        st.session_state.messages = [system_message]
+        # Only reset memory on new file load
+        if not st.session_state.messages or (
+            st.session_state.messages and
+            ("system" not in st.session_state.messages[0].get("role", ""))
+        ):
+            st.session_state.messages = [system_message]
+        elif (
+            st.session_state.messages and
+            st.session_state.messages[0].get("role", "") == "system" and
+            st.session_state.messages[0].get("content", "") != system_message["content"]
+        ):
+            st.session_state.messages[0] = system_message
     except Exception as e:
         st.sidebar.error(f"Error reading file: {e}")
         df = None
@@ -63,18 +70,13 @@ else:
     df = None
 
 if df is not None:
-    st.markdown(
-        f"<div style='color: #888; font-size:0.95em; text-align:center;'>"
-        f"Loaded CSV: <b>{df.shape[0]}</b> rows × <b>{df.shape[1]}</b> columns"
-        f"</div><hr style='margin-top:0.5rem; margin-bottom:1rem;'>",
-        unsafe_allow_html=True,
-    )
+    st.markdown(f"**Loaded CSV:** {df.shape[0]} rows × {df.shape[1]} columns")
 
 # --- Functions for function calling
 def search_csv(query: str):
     try:
         result_df = df.query(query)
-        return result_df.head(10).to_dict(orient="records")
+        return result_df.head(10).to_dict(orient="records")   # limit for safety
     except Exception as e:
         return {"error": f"Invalid query. Example: 'price > 100'. Details: {str(e)}"}
 
@@ -85,6 +87,7 @@ def count_unique(column: str):
     except Exception as e:
         return {"error": f"Column '{column}' not found or not countable. Details: {str(e)}"}
 
+# --- Function schemas for OpenAI
 function_schema = [
     {
         "name": "search_csv",
@@ -116,89 +119,19 @@ function_schema = [
     }
 ]
 
-# --- Chat style (CSS for chat bubbles)
-st.markdown("""
-<style>
-.chat-row-user {
-    background: #DCF8C6;
-    border-radius: 1rem 1rem 0.25rem 1rem;
-    margin: 0.25rem 0 0.75rem 25%;
-    padding: 0.6rem 1.1rem;
-    max-width: 70%;
-    text-align: right;
-    float: right;
-    clear: both;
-    box-shadow: 0 1px 2px #eee;
-}
-.chat-row-agent {
-    background: #F1F0F0;
-    border-radius: 1rem 1rem 1rem 0.25rem;
-    margin: 0.25rem 25% 0.75rem 0;
-    padding: 0.6rem 1.1rem;
-    max-width: 70%;
-    text-align: left;
-    float: left;
-    clear: both;
-    box-shadow: 0 1px 2px #eee;
-}
-.function-bubble {
-    background: #E3EDFA;
-    border-radius: 0.5rem;
-    margin: 0.1rem 15% 0.75rem 0;
-    padding: 0.4rem 1rem;
-    font-size: 0.97em;
-    color: #174075;
-    max-width: 80%;
-    box-shadow: 0 1px 2px #e2e6ea;
-}
-.reset-btn {
-    display: flex;
-    justify-content: flex-end;
-    margin-bottom: 0.8rem;
-}
-</style>
-""", unsafe_allow_html=True)
-
-# --- Chat interface header with reset
-col1, col2 = st.columns([6, 1])
-with col1:
-    st.markdown("### Conversation")
-with col2:
-    if st.button("Reset chat", key="reset_chat", help="Clear conversation history."):
-        if df is not None:
-            columns = ", ".join(df.columns)
-            st.session_state.messages = [{
-                "role": "system",
-                "content": (
-                    f"You are an AI data analyst for a CSV file with these columns: {columns}. "
-                    "When the user asks a question, always use the most relevant function to get the answer directly. "
-                    "Do not describe your plan or reasoning steps. Do not ask the user for clarification. "
-                    "Just call the function needed and give the answer, as briefly as possible. "
-                    "If you need to search or filter the CSV, use the 'search_csv' function. "
-                    "If you need to count unique values, use the 'count_unique' function. "
-                    "If you use 'search_csv', use Pandas query syntax."
-                ),
-            }]
-        else:
-            st.session_state.messages = []
-
-# --- Display chat history
-for i, msg in enumerate(st.session_state.messages[1:]):  # Skip system message
+# --- Chat interface
+st.markdown("### Conversation")
+for i, msg in enumerate(st.session_state.messages[1:]):  # Skip system message for display
     if msg["role"] == "user":
-        st.markdown(f"<div class='chat-row-user'>{msg['content']}</div>", unsafe_allow_html=True)
+        st.markdown(f"<div style='color: #4F8BF9;'><b>User:</b> {msg['content']}</div>", unsafe_allow_html=True)
     elif msg["role"] == "assistant":
-        st.markdown(f"<div class='chat-row-agent'>{msg['content']}</div>", unsafe_allow_html=True)
+        st.markdown(f"<div style='color: #1C6E4C;'><b>Agent:</b> {msg['content']}</div>", unsafe_allow_html=True)
     elif msg["role"] == "function":
         try:
             result = json.loads(msg["content"])
-            st.markdown(
-                f"<div class='function-bubble'><b>Function '{msg['name']}' output:</b>"
-                f"<pre style='background:transparent;'>{json.dumps(result, indent=2)}</pre></div>",
-                unsafe_allow_html=True)
+            st.markdown(f"<details><summary><b>Function '{msg['name']}' output:</b></summary><pre>{json.dumps(result, indent=2)}</pre></details>", unsafe_allow_html=True)
         except Exception:
-            st.markdown(
-                f"<div class='function-bubble'><b>Function '{msg['name']}' output:</b> {msg['content']}</div>",
-                unsafe_allow_html=True)
+            st.markdown(f"<b>Function '{msg['name']}' output:</b> {msg['content']}", unsafe_allow_html=True)
 
 # --- Sending user input and OpenAI call logic using a callback
 def send_message():
@@ -234,6 +167,7 @@ def send_message():
             func_name = msg["function_call"]["name"]
             args_json = msg["function_call"]["arguments"]
             args = json.loads(args_json)
+            # --- FIXED: Only pass the expected arg for each function
             if func_name == "search_csv":
                 function_result = search_csv(args.get("query", ""))
             elif func_name == "count_unique":
@@ -270,13 +204,5 @@ def send_message():
 
         st.session_state.temp_input = ""
 
-# --- User input box at bottom (styled like ChatGPT input)
 if df is not None:
-    st.markdown("<div style='height:18px;'></div>", unsafe_allow_html=True)
-    st.text_input(
-        "Type your question and press Enter...",
-        key="temp_input",
-        on_change=send_message,
-        placeholder="Ask anything about your CSV...",
-        label_visibility="collapsed",
-    )
+    st.text_input("Your message:", key="temp_input", on_change=send_message)