import os
import streamlit as st
import pandas as pd
import openai
import sqlite3
import json
import numpy as np
import datetime
from langchain.chains import RetrievalQA
from langchain.schema import Document
from langchain_core.retrievers import BaseRetriever
from pydantic import Field
from langchain_openai import ChatOpenAI
from langchain.prompts import ChatPromptTemplate
# --- CONFIG ---
DB_PATH = "json_vector.db"
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
EMBEDDING_MODEL = "text-embedding-ada-002"
# --- State Initialization ---
if "ingested_batches" not in st.session_state:
st.session_state.ingested_batches = 0
if "messages" not in st.session_state:
st.session_state.messages = []
if "modal_open" not in st.session_state:
st.session_state.modal_open = False
if "modal_content" not in st.session_state:
st.session_state.modal_content = ""
if "modal_title" not in st.session_state:
st.session_state.modal_title = ""
st.set_page_config(page_title="Chat with Your JSON Vectors", layout="wide")
st.title("Chat with Your Vectorized JSON Files (Hybrid Retrieval, SQLite, LLM)")
uploaded_files = st.file_uploader(
"Upload JSON files in batches (any structure)", type="json", accept_multiple_files=True
)
# --- Improved Flattening: extracts entity from emails/user fields for better matching
def flatten_json_obj(obj, parent_key="", sep="."):
items = {}
if isinstance(obj, dict):
for k, v in obj.items():
new_key = f"{parent_key}{sep}{k}" if parent_key else k
# Entity extraction: add name from email
if (
k.lower() in {"customer", "user", "email", "username"} and
isinstance(v, str) and "@" in v
):
local = v.split("@")[0]
items[new_key + "_name"] = local
items.update(flatten_json_obj(v, new_key, sep=sep))
elif isinstance(obj, list):
for i, v in enumerate(obj):
new_key = f"{parent_key}{sep}{i}" if parent_key else str(i)
items.update(flatten_json_obj(v, new_key, sep=sep))
else:
items[parent_key] = obj
return items
def get_embedding(text):
client = openai.OpenAI(api_key=OPENAI_API_KEY)
response = client.embeddings.create(input=[text], model=EMBEDDING_MODEL)
return response.data[0].embedding
def ensure_table():
conn = sqlite3.connect(DB_PATH)
cursor = conn.cursor()
cursor.execute("""
CREATE TABLE IF NOT EXISTS json_records (
id INTEGER PRIMARY KEY AUTOINCREMENT,
batch_time TEXT,
source_file TEXT,
raw_json TEXT,
flat_text TEXT,
embedding BLOB
)
""")
conn.commit()
conn.close()
def ingest_json_files(files):
ensure_table()
rows = []
batch_time = datetime.datetime.utcnow().isoformat()
for file in files:
raw = json.load(file)
source_name = file.name
if isinstance(raw, list):
records = raw
elif isinstance(raw, dict):
main_lists = [v for v in raw.values() if isinstance(v, list)]
records = main_lists[0] if main_lists else [raw]
else:
records = [raw]
for rec in records:
flat = flatten_json_obj(rec)
flat_text = "; ".join([f"{k}: {v}" for k, v in flat.items()])
rows.append((batch_time, source_name, json.dumps(rec), flat_text))
if not rows:
st.warning("No records found in uploaded files!")
return
df = pd.DataFrame(rows, columns=["batch_time", "source_file", "raw_json", "flat_text"])
st.write(f"Flattened {len(df)} records. Generating embeddings (this may take time, please wait)...")
df["embedding"] = df["flat_text"].apply(get_embedding)
conn = sqlite3.connect(DB_PATH)
cursor = conn.cursor()
for _, row in df.iterrows():
emb_bytes = np.array(row.embedding, dtype=np.float32).tobytes()
cursor.execute("""
INSERT INTO json_records (batch_time, source_file, raw_json, flat_text, embedding)
VALUES (?, ?, ?, ?, ?)
""", (row.batch_time, row.source_file, row.raw_json, row.flat_text, emb_bytes))
conn.commit()
conn.close()
st.success(f"Ingested and indexed {len(df)} new records!")
st.session_state.ingested_batches += 1
if uploaded_files and st.button("Ingest batch to database"):
ingest_json_files(uploaded_files)
# --- VECTOR RETRIEVAL
def query_vector_db(user_query, top_k=5):
query_emb = get_embedding(user_query)
conn = sqlite3.connect(DB_PATH)
cursor = conn.cursor()
cursor.execute("SELECT id, batch_time, source_file, raw_json, flat_text, embedding FROM json_records")
results = []
for row in cursor.fetchall():
db_emb = np.frombuffer(row[5], dtype=np.float32)
if len(db_emb) != len(query_emb): continue
sim = np.dot(query_emb, db_emb) / (np.linalg.norm(query_emb) * np.linalg.norm(db_emb))
results.append((sim, row))
conn.close()
results = sorted(results, reverse=True)[:top_k]
docs = []
for sim, row in results:
meta = {
"id": row[0],
"batch_time": str(row[1]),
"source_file": row[2],
"similarity": f"{sim:.4f} (embedding)",
"raw_json": row[3],
}
docs.append(Document(page_content=row[4], metadata=meta))
return docs
# --- PYTHON FUZZY/KEYWORD SEARCH
def python_fuzzy_match(user_query, top_k=5):
query_terms = set(user_query.lower().replace("@", " ").replace(".", " ").split())
conn = sqlite3.connect(DB_PATH)
cursor = conn.cursor()
cursor.execute("SELECT id, batch_time, source_file, raw_json, flat_text FROM json_records")
results = []
for row in cursor.fetchall():
flat_text = row[4].lower()
# score = # of query terms present as substring in the flat_text
score = sum(any(term in flat_text for term in query_terms) for term in query_terms)
if score > 0:
results.append((score, row))
conn.close()
results = sorted(results, reverse=True)[:top_k]
docs = []
for score, row in results:
meta = {
"id": row[0],
"batch_time": str(row[1]),
"source_file": row[2],
"similarity": f"{score} (fuzzy)",
"raw_json": row[3],
}
docs.append(Document(page_content=row[4], metadata=meta))
return docs
# --- HYBRID RETRIEVER
def hybrid_query(user_query, top_k=5):
vector_docs = query_vector_db(user_query, top_k=top_k)
fuzzy_docs = python_fuzzy_match(user_query, top_k=top_k)
seen_ids = set()
all_docs = []
for doc in (vector_docs + fuzzy_docs):
doc_id = doc.metadata.get("id")
if doc_id not in seen_ids:
all_docs.append(doc)
seen_ids.add(doc_id)
return all_docs[:top_k]
class HybridRetriever(BaseRetriever):
top_k: int = Field(default=5)
def _get_relevant_documents(self, query, run_manager=None, **kwargs):
return hybrid_query(query, self.top_k)
# --- SYSTEM PROMPT & PROMPT TEMPLATE
system_prompt = (
"You are a JSON data assistant. Always give a direct, concise answer based only on the context provided. "
"If you do not see the answer in the context, reply: 'I don’t have that information.' "
"Never make up information. Never ask for clarification."
)
prompt = ChatPromptTemplate.from_messages([
("system", system_prompt),
("human", "Context:\n{context}\n\nQuestion: {question}")
])
llm = ChatOpenAI(model="gpt-4.1", openai_api_key=OPENAI_API_KEY, temperature=0)
retriever = HybridRetriever(top_k=5)
qa_chain = RetrievalQA.from_chain_type(
llm=llm,
retriever=retriever,
chain_type_kwargs={"prompt": prompt},
return_source_documents=True,
)
# --- Chat UI and Conversation Area ---
st.markdown("### Ask any question about your data, just like ChatGPT.")
for msg in st.session_state.messages:
if msg["role"] == "user":
st.markdown(f"User: {msg['content']}
", unsafe_allow_html=True)
elif msg["role"] == "assistant":
st.markdown(f"Agent: {msg['content']}
", unsafe_allow_html=True)
elif msg["role"] == "function":
st.markdown(f"Function Output:
{msg['content']}