Update app.py

app.py

@@ -26,12 +26,44 @@ HEADERS = {
 st.sidebar.header("Upload CSV File")
 uploaded_file = st.sidebar.file_uploader("Choose a CSV file", type="csv")
 
-if uploaded_file:
+# --- Conversation memory: Use Streamlit session state
+if "messages" not in st.session_state:
+    st.session_state.messages = []
+if "temp_input" not in st.session_state:
+    st.session_state.temp_input = ""
+
+# --- Only load df and reset chat on new file upload
+if uploaded_file is not None:
     try:
         df = pd.read_csv(uploaded_file)
         st.sidebar.success("File uploaded successfully!")
         st.sidebar.write("Preview of the uploaded file:")
         st.sidebar.dataframe(df.head())
+        columns = ", ".join(df.columns)
+        system_message = {
+            "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."
+            ),
+        }
+        # 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
@@ -45,8 +77,7 @@ if df is not None:
 def search_csv(query: str):
     try:
         result_df = df.query(query)
-        # Limit output to 50 rows for large results
-        return result_df.head(50).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)}"}
 
@@ -89,45 +120,13 @@ function_schema = [
     }
 ]
 
-# --- Map function names to Python functions
 function_map = {
     "search_csv": search_csv,
     "count_unique": count_unique,
 }
 
-# --- Conversation memory: Use Streamlit session state
-if "messages" not in st.session_state:
-    st.session_state.messages = []
-
-if "temp_input" not in st.session_state:
-    st.session_state.temp_input = ""
-
-# If CSV is loaded, update the system prompt with current columns
-if df is not None:
-    columns = ", ".join(df.columns)
-    system_message = {
-    "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."
-    ),
-}
-
-    # Ensure the system message is always at the start and up-to-date
-    if not st.session_state.messages or st.session_state.messages[0]["role"] != "system":
-        st.session_state.messages.insert(0, system_message)
-    else:
-        st.session_state.messages[0] = system_message
-
 # --- Chat interface
 st.markdown("### Conversation")
-
-# Display chat history (like ChatGPT)
 for i, msg in enumerate(st.session_state.messages[1:]):  # Skip system message for display
     if msg["role"] == "user":
         st.markdown(f"<div style='color: #4F8BF9;'><b>User:</b> {msg['content']}</div>", unsafe_allow_html=True)
@@ -145,10 +144,12 @@ def send_message():
     user_input = st.session_state.temp_input
     if user_input and user_input.strip():
         st.session_state.messages.append({"role": "user", "content": user_input})
-
-        # Compose messages for OpenAI (entire chat history)
-        chat_messages = st.session_state.messages.copy()
-
+        # Limit history for context size (keep system + last 8)
+        chat_messages = st.session_state.messages
+        if len(chat_messages) > 10:
+            chat_messages = [chat_messages[0]] + chat_messages[-9:]
+        else:
+            chat_messages = chat_messages.copy()
         # First OpenAI call: Check for function call
         chat_resp = requests.post(
             "https://api.openai.com/v1/chat/completions",
@@ -177,15 +178,17 @@ def send_message():
                 function_result = function_map[func_name](**args)
             else:
                 function_result = {"error": f"Unknown function: {func_name}"}
-            # Append function call and output to history
             st.session_state.messages.append({
                 "role": "function",
                 "name": func_name,
                 "content": json.dumps(function_result),
             })
-
-            # Second OpenAI call: Get final answer with function result
-            followup_messages = st.session_state.messages.copy()
+            # Limit history again for second call
+            followup_messages = st.session_state.messages
+            if len(followup_messages) > 12:
+                followup_messages = [followup_messages[0]] + followup_messages[-11:]
+            else:
+                followup_messages = followup_messages.copy()
             final_resp = requests.post(
                 "https://api.openai.com/v1/chat/completions",
                 headers=HEADERS,
@@ -199,15 +202,11 @@ def send_message():
             )
             final_resp.raise_for_status()
             answer = final_resp.json()["choices"][0]["message"]["content"]
-            # Add assistant's reply to chat
             st.session_state.messages.append({"role": "assistant", "content": answer})
         else:
-            # No function call: Just add model's reply
             st.session_state.messages.append({"role": "assistant", "content": msg["content"]})
 
-        # Clear input after sending (now legal and safe)
         st.session_state.temp_input = ""
 
-# --- User input box at bottom (like ChatGPT)
 if df is not None:
     st.text_input("Your message:", key="temp_input", on_change=send_message)