Update app.py

app.py

@@ -1,8 +1,12 @@
 import streamlit as st
 import os
 import json
-import requests
 import traceback
+from datetime import datetime
+
+# ---- CONFIG ----
+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)
 if "json_data" not in st.session_state:
@@ -14,284 +18,118 @@ if "temp_input" not in st.session_state:
 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")
-
-# --- 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.")
-    st.stop()
-
-HEADERS = {
-    "Authorization": f"Bearer {OPENAI_API_KEY}",
-    "Content-Type": "application/json",
-}
+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
 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
 )
 
-# --- Only clear/load state when files change
+# --- Load/clear files as needed
 if uploaded_files and not st.session_state.files_loaded:
     st.session_state.json_data.clear()
-    file_summaries = []
     for f in uploaded_files:
         try:
             content = json.load(f)
             st.session_state.json_data[f.name] = content
-            if isinstance(content, dict):
-                keys = list(content.keys())
-            elif isinstance(content, list) and content and isinstance(content[0], dict):
-                keys = list(content[0].keys())
-            else:
-                keys = []
-            file_summaries.append(f"{f.name}: keys={keys[:10]}{'...' if len(keys)>10 else ''}")
             st.sidebar.success(f"Loaded: {f.name}")
-            st.sidebar.write(f"Keys: {keys[:10]}{'...' if len(keys)>10 else ''}")
         except Exception as e:
             st.sidebar.error(f"Error reading {f.name}: {e}")
-
-    # --- 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')"
-        )
-    }
-    st.session_state.messages = [system_message]
+    st.session_state.messages = []
     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
-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 functions for LLM
-def search_json(file_name, key, value):
-    def recursive_search(obj, key, value, found):
+# --- 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.
+    """
+    matches = []
+    def _search(obj, path, file_name):
         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)
+            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 item in obj:
-                recursive_search(item, key, value, found)
-        return found
-    try:
-        data = st.session_state.json_data[file_name]
-        results = recursive_search(data, key, value, [])
-        return results
-    except Exception as e:
-        return {"error": str(e)}
-
-def list_keys(file_name):
-    try:
-        data = st.session_state.json_data[file_name]
-        if isinstance(data, dict):
-            return list(data.keys())
-        elif isinstance(data, list) and data and isinstance(data[0], dict):
-            return list(data[0].keys())
-        else:
-            return []
-    except Exception as e:
-        return {"error": str(e)}
-
-def count_key_occurrences(file_name, key):
-    try:
-        data = st.session_state.json_data[file_name]
-        if isinstance(data, dict):
-            return 1 if key in data else 0
-        elif isinstance(data, list):
-            return sum(1 for item in data if isinstance(item, dict) and key in item)
-        else:
-            return 0
-    except Exception as e:
-        return {"error": str(e)}
-
-# --- Function schema
-function_schema = [
-    {
-        "name": "search_json",
-        "description": "Find records in the specified JSON file where key matches a given value.",
-        "parameters": {
-            "type": "object",
-            "properties": {
-                "file_name": {"type": "string", "description": "The uploaded JSON file to search."},
-                "key": {"type": "string", "description": "The key/field to filter by."},
-                "value": {"type": "string", "description": "The value to match."}
-            },
-            "required": ["file_name", "key", "value"],
-        },
-    },
-    {
-        "name": "list_keys",
-        "description": "List all top-level keys in a given JSON file.",
-        "parameters": {
-            "type": "object",
-            "properties": {
-                "file_name": {"type": "string", "description": "The uploaded JSON file."},
-            },
-            "required": ["file_name"],
-        },
-    },
-    {
-        "name": "count_key_occurrences",
-        "description": "Count the number of times a given key appears in a JSON file.",
-        "parameters": {
-            "type": "object",
-            "properties": {
-                "file_name": {"type": "string", "description": "The uploaded JSON file."},
-                "key": {"type": "string", "description": "The key to count."},
-            },
-            "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"]
-        }
-    }
-]
-
-# --- Conversation UI
-st.markdown("### Conversation")
-for i, msg in enumerate(st.session_state.messages[1:]):
+            for idx, item in enumerate(obj):
+                _search(item, path + [f"[{idx}]"], file_name)
+    for file_name, data in st.session_state.json_data.items():
+        _search(data, [], file_name)
+    return matches
+
+# --- MAIN QUERY LOGIC ---
+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.
+    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+(.*?)\?*$"
+    ]
+    found_name = None
+    for pat in patterns:
+        m = re.search(pat, query, re.IGNORECASE)
+        if m:
+            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)
+        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)
+        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"
+            answers.append(
+                f"**{res['matched_name']}** (in file `{res['file']}`) — {login}"
+            )
+        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!"
+
+# --- CHAT UI (NO LLM ROUNDTRIP!)
+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":
         st.markdown(f"<div style='color: #4F8BF9;'><b>User:</b> {msg['content']}</div>", unsafe_allow_html=True)
-    elif msg["role"] == "assistant":
-        content = msg.get("content", "")
-        if content.strip():
-            st.markdown(f"<div style='color: #1C6E4C;'><b>Agent:</b> {content}</div>", unsafe_allow_html=True)
-        else:
-            st.markdown(f"<div style='color: #DC143C;'><b>Agent:</b> [No response generated]</div>", unsafe_allow_html=True)
-    elif msg["role"] == "function":
-        try:
-            result = json.loads(msg["content"])
-            st.markdown(f"<details><summary><b>Function '{msg['name']}' output:</b></summary><pre>{json.dumps(result, indent=2)}</pre></details>", unsafe_allow_html=True)
-        except Exception:
-            st.markdown(f"<b>Function '{msg['name']}' output:</b> {msg['content']}", unsafe_allow_html=True)
+    else:
+        st.markdown(f"<div style='color: #1C6E4C;'><b>Agent:</b> {msg['content']}</div>", unsafe_allow_html=True)
 
-# --- Chat input and OpenAI handling
 def send_message():
-    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,
-                json={
-                    "model": "gpt-4.1",
-                    "messages": chat_messages,
-                    "functions": function_schema,
-                    "function_call": "auto",
-                    "temperature": 0,
-                    "max_tokens": 1000,
-                },
-                timeout=60,
-            )
-            chat_resp.raise_for_status()
-            response_json = chat_resp.json()
-            msg = response_json["choices"][0]["message"]
-
-            if msg.get("function_call"):
-                func_name = msg["function_call"]["name"]
-                args_json = msg["function_call"]["arguments"]
-                args = json.loads(args_json)
-
-                if func_name == "search_json":
-                    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}"}
-
-                st.session_state.messages.append({
-                    "role": "function",
-                    "name": func_name,
-                    "content": json.dumps(result),
-                })
-                followup_messages = st.session_state.messages
-                if len(followup_messages) > 12:
-                    followup_messages = [followup_messages[0]] + followup_messages[-11:]
-                else:
-                    followup_messages = followup_messages.copy()
-                final_resp = requests.post(
-                    "https://api.openai.com/v1/chat/completions",
-                    headers=HEADERS,
-                    json={
-                        "model": "gpt-4.1",
-                        "messages": followup_messages,
-                        "temperature": 0,
-                        "max_tokens": 1500,
-                    },
-                    timeout=60,
-                )
-                final_resp.raise_for_status()
-                final_json = final_resp.json()
-                answer = final_json["choices"][0]["message"]["content"]
-                st.session_state.messages.append({"role": "assistant", "content": answer})
-            else:
-                st.session_state.messages.append({"role": "assistant", "content": msg["content"]})
+    user_input = st.session_state.temp_input
+    if user_input.strip():
+        st.session_state.messages.append({"role": "user", "content": user_input})
+        answer = handle_user_query(user_input)
+        st.session_state.messages.append({"role": "assistant", "content": answer})
         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)
-    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.")