| import streamlit as st |
| import os |
| import json |
| import requests |
| import traceback |
|
|
| st.set_page_config(page_title="JSON-Backed AI Chat Agent", layout="wide") |
| st.title("JSON-Backed AI Chat Agent") |
|
|
| 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.sidebar.header("Upload Multiple JSON Files") |
| uploaded_files = st.sidebar.file_uploader( |
| "Choose one or more JSON files", type="json", accept_multiple_files=True |
| ) |
|
|
| if "json_data" not in st.session_state: |
| st.session_state.json_data = {} |
| if "messages" not in st.session_state: |
| st.session_state.messages = [] |
| if "temp_input" not in st.session_state: |
| st.session_state.temp_input = "" |
|
|
| if uploaded_files: |
| 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_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." |
| ), |
| } |
| st.session_state.messages = [system_message] |
| else: |
| st.session_state.json_data.clear() |
|
|
| def search_json(file_name, key, value): |
| 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 [] |
| 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 = [ |
| { |
| "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"], |
| }, |
| }, |
| ] |
|
|
| st.markdown("### Conversation") |
| for i, msg in enumerate(st.session_state.messages[1:]): |
| 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) |
|
|
| def send_message(): |
| 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() |
| try: |
| 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")) |
| 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"]}) |
| except Exception as e: |
| st.error("Exception: " + str(e)) |
| st.code(traceback.format_exc()) |
| finally: |
| st.session_state.temp_input = "" |
|
|
| if st.session_state.json_data: |
| st.text_input("Your message:", key="temp_input", on_change=send_message) |
| else: |
| st.info("Please upload at least one JSON file to start chatting.") |
|
|
