Update app.py

app.py

@@ -50,21 +50,42 @@ if uploaded_files:
         except Exception as e:
             st.sidebar.error(f"Error reading {f.name}: {e}")
 
+    # --- System prompt with explicit example ---
     system_message = {
         "role": "system",
         "content": (
-            "You are an AI data analyst for the uploaded JSON files. "
-            "Each file may have a different structure and set of keys. "
-            "When the user asks a question about all data (like 'How many females are there?'), "
-            "use the 'search_all_jsons' function to search through every uploaded file recursively. "
-            "Return results with the count and details, grouped by file if needed. "
-            "If the question is specific to one file, use the appropriate function for that file."
+            "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. "
+            "Example:\n"
+            "User: How many females are there?\n"
+            "Assistant: (Call search_all_jsons with key='gender', value='female')"
         )
     }
     st.session_state.messages = [system_message]
 else:
     st.session_state.json_data.clear()
 
+# --- 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():
+        def recursive_search(obj):
+            if isinstance(obj, dict):
+                if key in obj and str(obj[key]).lower() == str(value).lower():
+                    found.append({**obj, "__file__": file_name})
+                for v in obj.values():
+                    recursive_search(v)
+            elif isinstance(obj, list):
+                for item in obj:
+                    recursive_search(item)
+        recursive_search(data)
+    return found
+
+# --- Other function stubs (for completeness) ---
 def search_json(file_name, key, value):
     def recursive_search(obj, key, value, found):
         if isinstance(obj, dict):
@@ -107,22 +128,7 @@ def count_key_occurrences(file_name, key):
     except Exception as e:
         return {"error": str(e)}
 
-def search_all_jsons(key, value):
-    """Search all JSON files for dicts with key==value, recursively."""
-    found = []
-    for file_name, data in st.session_state.json_data.items():
-        def recursive_search(obj):
-            if isinstance(obj, dict):
-                if key in obj and str(obj[key]).lower() == str(value).lower():
-                    found.append({**obj, "__file__": file_name})
-                for v in obj.values():
-                    recursive_search(v)
-            elif isinstance(obj, list):
-                for item in obj:
-                    recursive_search(item)
-        recursive_search(data)
-    return found
-
+# --- Function schema including the new all-files search ---
 function_schema = [
     {
         "name": "search_json",
@@ -174,6 +180,7 @@ function_schema = [
     }
 ]
 
+# --- Conversation UI ---
 st.markdown("### Conversation")
 for i, msg in enumerate(st.session_state.messages[1:]):
     if msg["role"] == "user":
@@ -191,6 +198,7 @@ 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 ---
 def send_message():
     user_input = st.session_state.temp_input
     if user_input and user_input.strip():
@@ -205,7 +213,7 @@ def send_message():
                 "https://api.openai.com/v1/chat/completions",
                 headers=HEADERS,
                 json={
-                    "model": "gpt-4.1",
+                    "model": "gpt-4o",
                     "messages": chat_messages,
                     "functions": function_schema,
                     "function_call": "auto",
@@ -248,7 +256,7 @@ def send_message():
                     "https://api.openai.com/v1/chat/completions",
                     headers=HEADERS,
                     json={
-                        "model": "gpt-4.1",
+                        "model": "gpt-4o",
                         "messages": followup_messages,
                         "temperature": 0,
                         "max_tokens": 1500,
@@ -269,5 +277,10 @@ def send_message():
 
 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)}")
+        st.json(results)
 else:
     st.info("Please upload at least one JSON file to start chatting.")