Update app.py

app.py

@@ -4,6 +4,16 @@ import json
 import requests
 import traceback
 
+# --- ALWAYS INIT SESSION STATE FIRST (before any widgets)
+if "json_data" not in st.session_state:
+    st.session_state.json_data = {}
+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 "files_loaded" not in st.session_state:
+    st.session_state.files_loaded = False
+
 # --- Page config
 st.set_page_config(page_title="JSON-Backed AI Chat Agent", layout="wide")
 st.title("JSON-Backed AI Chat Agent")
@@ -24,14 +34,8 @@ uploaded_files = st.sidebar.file_uploader(
     "Choose one or more JSON files", type="json", accept_multiple_files=True
 )
 
-if "json_data" not in st.session_state:
-    st.session_state.json_data = {}
-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 uploaded_files:
+# --- Only clear/load state when files change
+if uploaded_files and not st.session_state.files_loaded:
     st.session_state.json_data.clear()
     file_summaries = []
     for f in uploaded_files:
@@ -50,34 +54,35 @@ if uploaded_files:
         except Exception as e:
             st.sidebar.error(f"Error reading {f.name}: {e}")
 
-    # --- System prompt with explicit example ---
+    # --- System prompt with explicit few-shot examples
     system_message = {
-    "role": "system",
-    "content": (
-        "You are an AI data analyst for uploaded JSON files. "
-        "Each file may have different structures and keys, including lists and nested dictionaries. "
-        "You have access to a function 'search_all_jsons' that finds all records in all JSON files where a key matches a value, recursively. "
-        "If a user asks about groups of people or wants to know counts such as 'How many females are there?', "
-        "interpret this as 'search for all records where gender equals female'. "
-        "Always use the 'search_all_jsons' function with key='gender' and value='female' for such queries, unless another key/value is clear from context. "
-        "If someone asks 'How many males?', search for gender equals male, etc. "
-        "EXAMPLES:\n"
-        "User: How many females are there?\n"
-        "Assistant: (Call search_all_jsons with key='gender', value='female')\n"
-        "User: How many males are there?\n"
-        "Assistant: (Call search_all_jsons with key='gender', value='male')\n"
-        "User: Show all females\n"
-        "Assistant: (Call search_all_jsons with key='gender', value='female')\n"
-        "User: How many people are named Emily?\n"
-        "Assistant: (Call search_all_jsons with key='firstName', value='Emily')"
-    )
-}
-
+        "role": "system",
+        "content": (
+            "You are an AI data analyst for uploaded JSON files. "
+            "Each file may have different structures and keys, including lists and nested dictionaries. "
+            "You have access to a function 'search_all_jsons' that finds all records in all JSON files where a key matches a value, recursively. "
+            "If a user asks about groups of people or wants to know counts such as 'How many females are there?', "
+            "interpret this as 'search for all records where gender equals female'. "
+            "Always use the 'search_all_jsons' function with key='gender' and value='female' for such queries, unless another key/value is clear from context. "
+            "If someone asks 'How many males?', search for gender equals male, etc. "
+            "EXAMPLES:\n"
+            "User: How many females are there?\n"
+            "Assistant: (Call search_all_jsons with key='gender', value='female')\n"
+            "User: How many males are there?\n"
+            "Assistant: (Call search_all_jsons with key='gender', value='male')\n"
+            "User: Show all females\n"
+            "Assistant: (Call search_all_jsons with key='gender', value='female')\n"
+            "User: How many people are named Emily?\n"
+            "Assistant: (Call search_all_jsons with key='firstName', value='Emily')"
+        )
+    }
     st.session_state.messages = [system_message]
-else:
+    st.session_state.files_loaded = True
+elif not uploaded_files:
     st.session_state.json_data.clear()
+    st.session_state.files_loaded = False
 
-# --- Recursive search for key/value in all files ---
+# --- Recursive search for key/value in all files
 def search_all_jsons(key, value):
     found = []
     for file_name, data in st.session_state.json_data.items():
@@ -93,7 +98,7 @@ def search_all_jsons(key, value):
         recursive_search(data)
     return found
 
-# --- Other function stubs (for completeness) ---
+# --- Other functions for LLM
 def search_json(file_name, key, value):
     def recursive_search(obj, key, value, found):
         if isinstance(obj, dict):
@@ -136,7 +141,7 @@ def count_key_occurrences(file_name, key):
     except Exception as e:
         return {"error": str(e)}
 
-# --- Function schema including the new all-files search ---
+# --- Function schema
 function_schema = [
     {
         "name": "search_json",
@@ -188,7 +193,7 @@ function_schema = [
     }
 ]
 
-# --- Conversation UI ---
+# --- Conversation UI
 st.markdown("### Conversation")
 for i, msg in enumerate(st.session_state.messages[1:]):
     if msg["role"] == "user":
@@ -206,17 +211,17 @@ for i, msg in enumerate(st.session_state.messages[1:]):
         except Exception:
             st.markdown(f"<b>Function '{msg['name']}' output:</b> {msg['content']}", unsafe_allow_html=True)
 
-# --- Chat input and OpenAI handling ---
+# --- Chat input and OpenAI handling
 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})
-        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()
-        try:
+    try:
+        user_input = st.session_state.temp_input
+        if user_input and user_input.strip():
+            st.session_state.messages.append({"role": "user", "content": user_input})
+            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()
             chat_resp = requests.post(
                 "https://api.openai.com/v1/chat/completions",
                 headers=HEADERS,
@@ -277,15 +282,13 @@ def send_message():
                 st.session_state.messages.append({"role": "assistant", "content": answer})
             else:
                 st.session_state.messages.append({"role": "assistant", "content": msg["content"]})
-        except Exception as e:
-            st.error("Exception: " + str(e))
-            st.code(traceback.format_exc())
-        finally:
-            st.session_state.temp_input = ""
+        st.session_state.temp_input = ""
+    except Exception as e:
+        st.error("Exception: " + str(e))
+        st.code(traceback.format_exc())
 
 if st.session_state.json_data:
     st.text_input("Your message:", key="temp_input", on_change=send_message)
-    # --- Manual debug/test button ---
     if st.button("Test: Count females in all JSONs"):
         results = search_all_jsons("gender", "female")
         st.write(f"Females found: {len(results)}")