Update app.py

app.py

@@ -37,7 +37,7 @@ elif not uploaded_files:
     st.session_state.json_data.clear()
     st.session_state.files_loaded = False
 
-# --- NORMALIZE ---
+# --- NORMALIZE & FUZZY ---
 def normalize(s):
     return ' '.join(str(s).lower().replace("_", " ").replace("-", " ").replace(".", " ").split())
 
@@ -120,6 +120,32 @@ def count_key_occurrences(file_name, key):
     except Exception as e:
         return {"error": str(e)}
 
+# --- NEW: FIND/COUNT IN ARRAYS ---
+def find_in_arrays(key, value, return_count=True):
+    # Searches ALL arrays for dicts where key == value
+    matches = []
+    count = 0
+    for file_name, data in st.session_state.json_data.items():
+        def recursive(obj):
+            nonlocal count
+            if isinstance(obj, list):
+                for item in obj:
+                    if isinstance(item, dict):
+                        for k, v in item.items():
+                            if normalize(k) == normalize(key) and normalize(str(v)) == normalize(str(value)):
+                                matches.append({
+                                    "file": file_name,
+                                    "item": item,
+                                    "array_path": key
+                                })
+                                count += 1
+                    recursive(item)
+            elif isinstance(obj, dict):
+                for v in obj.values():
+                    recursive(v)
+        recursive(data)
+    return count if return_count else matches
+
 # --- FUNCTION SCHEMA for OpenAI ---
 function_schema = [
     {
@@ -167,14 +193,28 @@ function_schema = [
             },
             "required": ["file_name", "key"]
         }
+    },
+    {
+        "name": "find_in_arrays",
+        "description": "Find/count all objects in any arrays/lists where key equals value (e.g. done:true for completed tasks).",
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "key": {"type": "string", "description": "The key to search for, e.g., 'done'"},
+                "value": {"type": "string", "description": "The value to match, e.g., 'true' or 'false'"},
+                "return_count": {"type": "boolean", "description": "Return the count (true) or matching records (false)."}
+            },
+            "required": ["key", "value"]
+        }
     }
 ]
 
-# --- SYSTEM PROMPT: Set expectations for OpenAI! ---
+# --- SYSTEM PROMPT ---
 system_message = {
     "role": "system",
     "content": (
         "You are a JSON data assistant. Use the functions provided to answer the user's question. "
+        "If the user asks for the number or details of items in a list/array (e.g., completed tasks), use 'find_in_arrays'. "
         "If the user's query does not mention a key, use 'fuzzy_value_search' to match on any value. "
         "If a key is mentioned (like 'apps_installed'), use 'search_all_jsons' for that key and the value. "
         "You may use 'list_keys' to help discover the file structure if needed. "
@@ -209,7 +249,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",
@@ -236,6 +276,12 @@ def send_message():
                     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 == "find_in_arrays":
+                    result = find_in_arrays(
+                        args.get("key"),
+                        args.get("value"),
+                        args.get("return_count", True)
+                    )
                 else:
                     result = {"error": f"Unknown function: {func_name}"}
 
@@ -253,7 +299,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": 1200,