Update app.py

app.py

@@ -6,13 +6,13 @@ import requests
 import json
 import re
 
-# — Page config
-st.set_page_config(page_title="CSV-Backed AI Agent with Function Calling", layout="wide")
+# --- Page config
+st.set_page_config(page_title="CSV-Backed AI Chat Agent", layout="wide")
 
-# — Title & image
-st.title("CSV-Backed AI Agent with Function Calling")
+# --- Title & image
+st.title("CSV-Backed AI Chat Agent")
 
-# — Load API key
+# --- Load API key
 OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
 if not OPENAI_API_KEY:
     st.error("❌ OPENAI_API_KEY not set in Settings → Secrets.")
@@ -23,7 +23,7 @@ HEADERS = {
     "Content-Type": "application/json",
 }
 
-# — Sidebar: CSV upload & preview
+# --- Sidebar: CSV upload & preview
 st.sidebar.header("Upload CSV File")
 uploaded_file = st.sidebar.file_uploader("Choose a CSV file", type="csv")
 
@@ -42,14 +42,7 @@ else:
 if df is not None:
     st.markdown(f"**Loaded CSV:** {df.shape[0]} rows × {df.shape[1]} columns")
 
-# — Prompt input
-prompt = st.text_area(
-    "Enter your prompt for the agent",
-    placeholder="e.g. Which products have price > 100?",
-    height=150,
-)
-
-# — Functions for function calling
+# --- Functions for function calling
 def search_csv(query: str):
     try:
         result_df = df.query(query)
@@ -65,7 +58,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 schemas for OpenAI
 function_schema = [
     {
         "name": "search_csv",
@@ -97,94 +90,122 @@ function_schema = [
     }
 ]
 
-# — Map function names to Python functions
+# --- Map function names to Python functions
 function_map = {
     "search_csv": search_csv,
     "count_unique": count_unique,
 }
 
-# — Run Agent
-if st.button("Run Agent"):
-    if df is None:
-        st.error("Please upload a CSV file first.")
-    elif not prompt.strip():
-        st.error("Please enter a prompt.")
+# --- Conversation memory: Use Streamlit session state
+if "messages" not in st.session_state:
+    # Initial system prompt includes column info (blank at start)
+    st.session_state.messages = []
+
+# 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 agent helping users analyze a CSV file with these columns: {columns}. "
+            "If you need to search or filter the CSV, call the 'search_csv' function. "
+            "If the user wants to know how many unique values are in a column, call 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:
-        # Build dynamic system message with column info
-        columns = ", ".join(df.columns)
-        system_message = {
-            "role": "system",
-            "content": (
-                f"You are an AI agent helping users analyze a CSV file with these columns: {columns}. "
-                "If you need to search or filter the CSV, call the 'search_csv' function. "
-                "If the user wants to know how many unique values are in a column, call the 'count_unique' function. "
-                "If you use 'search_csv', use Pandas query syntax."
-            ),
-        }
-        messages = [
-            system_message,
-            {"role": "user", "content": prompt}
-        ]
-
-        # First OpenAI call: See if a function call is needed
-        chat_resp = requests.post(
+        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 the 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)
+    elif msg["role"] == "assistant":
+        st.markdown(f"<div style='color: #1C6E4C;'><b>Agent:</b> {msg['content']}</div>", unsafe_allow_html=True)
+    elif msg["role"] == "function":
+        # Optionally, display function results
+        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)
+        except Exception:
+            st.markdown(f"<b>Function '{msg['name']}' output:</b> {msg['content']}", unsafe_allow_html=True)
+
+# --- User input box at bottom (like ChatGPT)
+if df is not None:
+    user_input = st.text_input("Your message:", key="user_input")
+    send = st.button("Send", key="send_btn")
+else:
+    user_input = None
+    send = False
+
+if send and user_input and user_input.strip():
+    # Append user's message to conversation
+    st.session_state.messages.append({"role": "user", "content": user_input})
+
+    # --- Compose messages for OpenAI (entire chat history)
+    chat_messages = st.session_state.messages.copy()
+
+    # --- First OpenAI call: Check for function call
+    chat_resp = requests.post(
+        "https://api.openai.com/v1/chat/completions",
+        headers=HEADERS,
+        json={
+            "model": "gpt-3.5-turbo-1106",
+            "messages": chat_messages,
+            "functions": function_schema,
+            "function_call": "auto",
+            "temperature": 0,
+            "max_tokens": 1000,
+        },
+        timeout=60,
+    )
+    chat_resp.raise_for_status()
+    response_json = chat_resp.json()
+    msg = response_json["choices"][0]["message"]
+
+    # --- If OpenAI requests a function call
+    if msg.get("function_call"):
+        func_name = msg["function_call"]["name"]
+        args_json = msg["function_call"]["arguments"]
+        args = json.loads(args_json)
+        # Call the correct Python function
+        if func_name in function_map:
+            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()
+        final_resp = requests.post(
             "https://api.openai.com/v1/chat/completions",
             headers=HEADERS,
             json={
-                "model": "gpt-3.5-turbo-1106",  # or "gpt-4-1106-preview" if you have it
-                "messages": messages,
-                "functions": function_schema,
-                "function_call": "auto",
+                "model": "gpt-3.5-turbo-1106",
+                "messages": followup_messages,
                 "temperature": 0,
-                "max_tokens": 1000,
+                "max_tokens": 1500,
             },
             timeout=60,
         )
-        chat_resp.raise_for_status()
-        response_json = chat_resp.json()
-        msg = response_json["choices"][0]["message"]
-
-        # If OpenAI requests a function call
-        if msg.get("function_call"):
-            func_name = msg["function_call"]["name"]
-            args_json = msg["function_call"]["arguments"]
-            args = json.loads(args_json)
-            # Call the appropriate function
-            if func_name in function_map:
-                function_result = function_map[func_name](**args)
-            else:
-                function_result = {"error": f"Unknown function: {func_name}"}
-
-            # Send function result back for final answer
-            followup_messages = [
-                system_message,
-                {"role": "user", "content": prompt},
-                {
-                    "role": "function",
-                    "name": func_name,
-                    "content": json.dumps(function_result),
-                }
-            ]
-
-            final_resp = requests.post(
-                "https://api.openai.com/v1/chat/completions",
-                headers=HEADERS,
-                json={
-                    "model": "gpt-3.5-turbo-1106",
-                    "messages": followup_messages,
-                    "temperature": 0,
-                    "max_tokens": 1500,
-                },
-                timeout=60,
-            )
-            final_resp.raise_for_status()
-            answer = final_resp.json()["choices"][0]["message"]["content"]
-
-            st.subheader("✅ Agent Answer")
-            st.markdown(answer)
-            st.subheader("🔎 Function Output")
-            st.json(function_result)
-        else:
-            # No function call: model answered directly
-            st.subheader("✅ Agent Answer")
-            st.markdown(msg["content"])
+        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
+    st.experimental_rerun()
+