Update app.py

app.py

@@ -4,13 +4,13 @@ import numpy as np
 import os
 import requests
 import json
-import re
 
 # --- Page config
 st.set_page_config(page_title="CSV-Backed AI Chat Agent", layout="wide")
 
 # --- Title & image
 st.title("CSV-Backed AI Chat Agent")
+st.image("./nadi-lok-image.png")
 
 # --- Load API key
 OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
@@ -68,7 +68,7 @@ function_schema = [
             "properties": {
                 "query": {
                     "type": "string",
-                    "description": "A Pandas query string, e.g. 'price > 100 and category == \"A\"'"
+                    "description": "A Pandas query string, e.g. 'price > 100 and city == \"Miami\"'"
                 },
             },
             "required": ["query"],
@@ -98,7 +98,6 @@ function_map = {
 
 # --- 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
@@ -123,13 +122,12 @@ if df is not None:
 st.markdown("### Conversation")
 
 # Display chat history (like ChatGPT)
-for i, msg in enumerate(st.session_state.messages[1:]):  # Skip the system message for display
+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)
     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)
@@ -138,6 +136,8 @@ for i, msg in enumerate(st.session_state.messages[1:]):  # Skip the system messa
 
 # --- User input box at bottom (like ChatGPT)
 if df is not None:
+    if "user_input" not in st.session_state:
+        st.session_state.user_input = ""
     user_input = st.text_input("Your message:", key="user_input")
     send = st.button("Send", key="send_btn")
 else:
@@ -148,10 +148,10 @@ 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)
+    # Compose messages for OpenAI (entire chat history)
     chat_messages = st.session_state.messages.copy()
 
-    # --- First OpenAI call: Check for function call
+    # First OpenAI call: Check for function call
     chat_resp = requests.post(
         "https://api.openai.com/v1/chat/completions",
         headers=HEADERS,
@@ -169,7 +169,7 @@ if send and user_input and user_input.strip():
     response_json = chat_resp.json()
     msg = response_json["choices"][0]["message"]
 
-    # --- If OpenAI requests a function call
+    # If OpenAI requests a function call
     if msg.get("function_call"):
         func_name = msg["function_call"]["name"]
         args_json = msg["function_call"]["arguments"]
@@ -186,7 +186,7 @@ if send and user_input and user_input.strip():
             "content": json.dumps(function_result),
         })
 
-        # --- Second OpenAI call: Get final answer with 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",
@@ -206,6 +206,7 @@ if send and user_input and user_input.strip():
     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()
 
+    # Clear input after sending
+    st.session_state.user_input = ""
+    # The UI will update on the next rerender (no need for manual rerun)