Update app.py

app.py

@@ -1,14 +1,11 @@
 import streamlit as st
-import os
 import json
-import traceback
-from datetime import datetime
 
-# ---- CONFIG ----
+# --- Constants for searching ---
 COMMON_NAME_KEYS = ["user", "username", "name", "fullName", "firstName", "lastName"]
 LOGIN_KEYS = ["lastLogin", "login", "loggedIn", "lastLoggedIn", "last_login", "last_logged_in"]
 
-# --- ALWAYS INIT SESSION STATE FIRST (before any widgets)
+# --- Session state setup ---
 if "json_data" not in st.session_state:
     st.session_state.json_data = {}
 if "messages" not in st.session_state:
@@ -21,13 +18,12 @@ if "files_loaded" not in st.session_state:
 st.set_page_config(page_title="Instant JSON Q&A", layout="wide")
 st.title("Instant JSON-Backed AI Q&A (No More Clarifying Loops)")
 
-# --- Upload area
+# --- Upload area ---
 st.sidebar.header("Upload Multiple JSON Files")
 uploaded_files = st.sidebar.file_uploader(
     "Choose one or more JSON files", type="json", accept_multiple_files=True
 )
 
-# --- Load/clear files as needed
 if uploaded_files and not st.session_state.files_loaded:
     st.session_state.json_data.clear()
     for f in uploaded_files:
@@ -43,48 +39,63 @@ elif not uploaded_files:
     st.session_state.json_data.clear()
     st.session_state.files_loaded = False
 
-# --- UNIVERSAL SEARCH FUNCTION ----
-def recursive_find_user(obj, target_name):
-    """
-    Recursively search for dicts containing a value matching target_name under any COMMON_NAME_KEYS.
-    Returns list of matches: dicts with the matching name, file name, and all possible login info found.
-    """
+# --- Robust normalization function for matching ---
+def normalize(s):
+    # Lowercase, remove underscores/dashes, strip, collapse double spaces
+    return str(s).lower().replace("_", " ").replace("-", " ").replace(".", " ").strip()
+
+# --- Recursive user search with robust matching ---
+def recursive_find_user(target_name):
     matches = []
-    def _search(obj, path, file_name):
-        if isinstance(obj, dict):
-            for k, v in obj.items():
-                # Check if any "name" key matches (case-insensitive substring match)
-                if k in COMMON_NAME_KEYS and target_name.lower() in str(v).lower():
-                    login_info = {}
-                    # Find login info at this level
-                    for lk in LOGIN_KEYS:
-                        if lk in obj:
-                            login_info[lk] = obj[lk]
-                    matches.append({
-                        "match_path": path + [k],
-                        "matched_name": v,
-                        "record": obj,
-                        "file": file_name,
-                        "login_info": login_info
-                    })
-                _search(v, path + [k], file_name)
-        elif isinstance(obj, list):
-            for idx, item in enumerate(obj):
-                _search(item, path + [f"[{idx}]"], file_name)
+    norm_target = normalize(target_name)
     for file_name, data in st.session_state.json_data.items():
-        _search(data, [], file_name)
+        def _search(obj, path):
+            if isinstance(obj, dict):
+                for k, v in obj.items():
+                    if k in COMMON_NAME_KEYS and norm_target in normalize(v):
+                        login_info = {}
+                        for lk in LOGIN_KEYS:
+                            if lk in obj:
+                                login_info[lk] = obj[lk]
+                        matches.append({
+                            "match_path": path + [k],
+                            "matched_name": v,
+                            "record": obj,
+                            "file": file_name,
+                            "login_info": login_info
+                        })
+                    _search(v, path + [k])
+            elif isinstance(obj, list):
+                for idx, item in enumerate(obj):
+                    _search(item, path + [f"[{idx}]"])
+        _search(data, [])
     return matches
 
-# --- MAIN QUERY LOGIC ---
+# --- Show all user keys/values (for debug) ---
+def show_all_users():
+    found = []
+    for file_name, data in st.session_state.json_data.items():
+        def recursive(obj, path):
+            if isinstance(obj, dict):
+                for k, v in obj.items():
+                    if k in COMMON_NAME_KEYS:
+                        found.append(f"{file_name} | {'.'.join(path + [k])} = {v}")
+                    recursive(v, path + [k])
+            elif isinstance(obj, list):
+                for idx, item in enumerate(obj):
+                    recursive(item, path + [f"[{idx}]"])
+        recursive(data, [])
+    return found
+
+# --- Main Q&A logic ---
+import re
 def handle_user_query(query):
-    # 1. If question looks like "last login for <name>", extract name and search
-    import re
-    # Covers "last login for", "when did X last login", "when was X last seen", etc.
+    # Try to extract a likely user name from the query
     patterns = [
-        r"(?:last\s*login.*?for|when\s+did)\s+(.*?)\?*$", # e.g. last login for Bob, when did Bob last login
-        r"when\s+was\s+(.*?)\s+last\s+(?:login|logged\s*in)", # when was Bob last logged in
-        r"last\s*login\s*of\s+(.*?)\?*$",
-        r"(?:info|details|record) for\s+(.*?)\?*$"
+        r"(?:last\s*login.*?for|when\s+did)\s+([a-zA-Z0-9 _\-\.]+)",
+        r"when\s+was\s+([a-zA-Z0-9 _\-\.]+)\s+last\s+(?:login|logged\s*in)",
+        r"last\s*login\s*of\s+([a-zA-Z0-9 _\-\.]+)",
+        r"(?:info|details|record) for\s+([a-zA-Z0-9 _\-\.]+)"
     ]
     found_name = None
     for pat in patterns:
@@ -93,15 +104,14 @@ def handle_user_query(query):
             found_name = m.group(1).strip()
             break
     if not found_name:
-        # Fallback: just look for the first capitalized word (e.g. Bob, Bob the Builder)
+        # Fallback: look for first capitalized word or phrase
         m = re.search(r"([A-Z][a-z]+(?: [A-Z][a-z]+)*)", query)
         if m:
             found_name = m.group(1).strip()
     if found_name:
-        results = recursive_find_user(st.session_state.json_data, found_name)
+        results = recursive_find_user(found_name)
         if not results:
             return f"No records found for '{found_name}' in any file."
-        # Report each match and all login fields found
         answers = []
         for res in results:
             login = ", ".join([f"{k}: {v}" for k, v in res["login_info"].items()]) if res["login_info"] else "No login info found"
@@ -110,10 +120,9 @@ def handle_user_query(query):
             )
         return "\n\n".join(answers)
     else:
-        # For all other queries, just report "Sorry, direct name lookups only. (Add more query logic if you want!)"
-        return "Sorry, I can only answer direct user info queries (e.g., 'When did Bob the Builder last login?'). If you want more logic, let me know!"
+        return "Sorry, I can only answer direct user info queries (e.g., 'When did Bob the Builder last login?')."
 
-# --- CHAT UI (NO LLM ROUNDTRIP!)
+# --- Chat UI ---
 st.markdown("### Ask about any user directly (e.g. 'When did Bob the Builder last login?')")
 for msg in st.session_state.messages:
     if msg["role"] == "user":
@@ -131,5 +140,7 @@ def send_message():
 
 if st.session_state.json_data:
     st.text_input("Your message:", key="temp_input", on_change=send_message)
+    if st.button("Show all users in uploaded JSONs"):
+        st.write(show_all_users())
 else:
     st.info("Please upload at least one JSON file to start chatting.")