Update app.py

app.py

@@ -5,11 +5,13 @@ import requests
 import json
 
 # --- Page config
-st.set_page_config(page_title="CSV-Backed AI Chat Agent", layout="wide")
+st.set_page_config(page_title="CSV-Backed AI Chat Agent", layout="centered")
 
-# --- Title & image
-st.title("CSV-Backed AI Chat Agent")
-st.image("./nadi-lok-image.png")
+# --- Minimal title
+st.markdown(
+    "<h2 style='text-align:center; margin-bottom:1rem;'>CSV-Backed AI Chat Agent</h2>",
+    unsafe_allow_html=True,
+)
 
 # --- Load API key
 OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
@@ -52,18 +54,8 @@ if uploaded_file is not None:
                 "If you use 'search_csv', use Pandas query syntax."
             ),
         }
-        # 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
+        # Reset memory on new file load
+        st.session_state.messages = [system_message]
     except Exception as e:
         st.sidebar.error(f"Error reading file: {e}")
         df = None
@@ -71,13 +63,18 @@ else:
     df = None
 
 if df is not None:
-    st.markdown(f"**Loaded CSV:** {df.shape[0]} rows × {df.shape[1]} columns")
+    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,
+    )
 
 # --- Functions for function calling
 def search_csv(query: str):
     try:
         result_df = df.query(query)
-        return result_df.head(10).to_dict(orient="records")   # limit for safety
+        return result_df.head(10).to_dict(orient="records")
     except Exception as e:
         return {"error": f"Invalid query. Example: 'price > 100'. Details: {str(e)}"}
 
@@ -88,7 +85,6 @@ 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",
@@ -120,19 +116,89 @@ function_schema = [
     }
 ]
 
-# --- Chat interface
-st.markdown("### Conversation")
-for i, msg in enumerate(st.session_state.messages[1:]):  # Skip system message for display
+# --- 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
     if msg["role"] == "user":
-        st.markdown(f"<div style='color: #4F8BF9;'><b>User:</b> {msg['content']}</div>", unsafe_allow_html=True)
+        st.markdown(f"<div class='chat-row-user'>{msg['content']}</div>", unsafe_allow_html=True)
     elif msg["role"] == "assistant":
-        st.markdown(f"<div style='color: #1C6E4C;'><b>Agent:</b> {msg['content']}</div>", unsafe_allow_html=True)
+        st.markdown(f"<div class='chat-row-agent'>{msg['content']}</div>", unsafe_allow_html=True)
     elif msg["role"] == "function":
         try:
             result = json.loads(msg["content"])
-            st.markdown(f"<details><summary><b>Function '{msg['name']}' output:</b></summary><pre>{json.dumps(result, indent=2)}</pre></details>", unsafe_allow_html=True)
+            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)
         except Exception:
-            st.markdown(f"<b>Function '{msg['name']}' output:</b> {msg['content']}", unsafe_allow_html=True)
+            st.markdown(
+                f"<div class='function-bubble'><b>Function '{msg['name']}' output:</b> {msg['content']}</div>",
+                unsafe_allow_html=True)
 
 # --- Sending user input and OpenAI call logic using a callback
 def send_message():
@@ -168,7 +234,6 @@ 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":
@@ -190,7 +255,7 @@ def send_message():
                 "https://api.openai.com/v1/chat/completions",
                 headers=HEADERS,
                 json={
-                    "model": "gpt-4o",
+                    "model": "gpt-4.1",
                     "messages": followup_messages,
                     "temperature": 0,
                     "max_tokens": 1500,
@@ -205,5 +270,13 @@ def send_message():
 
         st.session_state.temp_input = ""
 
+# --- User input box at bottom (styled like ChatGPT input)
 if df is not None:
-    st.text_input("Your message:", key="temp_input", on_change=send_message)
+    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",
+    )