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"User: {msg['content']}
", unsafe_allow_html=True)
elif msg["role"] == "assistant":
content = msg.get("content", "")
if content.strip():
st.markdown(f"Agent: {content}
", unsafe_allow_html=True)
else:
st.markdown(f"Agent: [No response generated]
", unsafe_allow_html=True)
elif msg["role"] == "function":
try:
result = json.loads(msg["content"])
st.markdown(f"Function '{msg['name']}' output:
{json.dumps(result, indent=2)}