Update app.py

app.py

@@ -4,9 +4,11 @@ import json
 import requests
 import traceback
 
+# --- Page config
 st.set_page_config(page_title="JSON-Backed AI Chat Agent", layout="wide")
 st.title("JSON-Backed AI Chat Agent")
 
+# --- 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.")
@@ -51,31 +53,33 @@ if uploaded_files:
     system_message = {
         "role": "system",
         "content": (
-            "You are an AI data analyst for the following JSON files:\n" +
-            "\n".join(file_summaries) +
-            "\nEach file may have a different structure and set of keys. " 
-            "When the user asks a question, identify which file(s) it applies to, "
-            "then use the most relevant function to extract the answer. "
-            "If the user does not specify a file, make your best guess based on keys/fields mentioned."
-        ),
+            "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."
+        )
     }
     st.session_state.messages = [system_message]
 else:
     st.session_state.json_data.clear()
 
 def search_json(file_name, key, value):
+    def recursive_search(obj, key, value, found):
+        if isinstance(obj, dict):
+            if key in obj and str(obj[key]).lower() == str(value).lower():
+                found.append(obj)
+            for v in obj.values():
+                recursive_search(v, key, value, found)
+        elif isinstance(obj, list):
+            for item in obj:
+                recursive_search(item, key, value, found)
+        return found
     try:
         data = st.session_state.json_data[file_name]
-        if isinstance(data, list):
-            results = [item for item in data if isinstance(item, dict) and str(item.get(key)) == str(value)]
-            return results[:10]
-        elif isinstance(data, dict):
-            if key in data and str(data[key]) == str(value):
-                return [{key: value}]
-            else:
-                return []
-        else:
-            return []
+        results = recursive_search(data, key, value, [])
+        return results
     except Exception as e:
         return {"error": str(e)}
 
@@ -103,6 +107,22 @@ 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 = [
     {
         "name": "search_json",
@@ -140,6 +160,18 @@ function_schema = [
             "required": ["file_name", "key"],
         },
     },
+    {
+        "name": "search_all_jsons",
+        "description": "Search all uploaded JSON files recursively for dicts where a key matches a value.",
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "key": {"type": "string", "description": "The key to search for (e.g. 'gender')"},
+                "value": {"type": "string", "description": "The value to match (e.g. 'female')"}
+            },
+            "required": ["key", "value"]
+        }
+    }
 ]
 
 st.markdown("### Conversation")
@@ -192,13 +224,13 @@ def send_message():
                 args = json.loads(args_json)
 
                 if func_name == "search_json":
-                    result = search_json(
-                        args.get("file_name"), args.get("key"), args.get("value")
-                    )
+                    result = search_json(args.get("file_name"), args.get("key"), args.get("value"))
                 elif func_name == "list_keys":
                     result = list_keys(args.get("file_name"))
                 elif func_name == "count_key_occurrences":
                     result = count_key_occurrences(args.get("file_name"), args.get("key"))
+                elif func_name == "search_all_jsons":
+                    result = search_all_jsons(args.get("key"), args.get("value"))
                 else:
                     result = {"error": f"Unknown function: {func_name}"}