Update app.py

app.py

@@ -6,7 +6,6 @@ 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
@@ -32,17 +31,25 @@ 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 (LangChain, SQLite, LLM)")
+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):
@@ -113,6 +120,7 @@ def ingest_json_files(files):
 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)
@@ -132,24 +140,63 @@ def query_vector_db(user_query, top_k=5):
             "id": row[0],
             "batch_time": str(row[1]),
             "source_file": row[2],
-            "similarity": f"{sim:.4f}",
+            "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
 
-class SQLiteVectorRetriever(BaseRetriever):
+# --- 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 query_vector_db(query, self.top_k)
+        return hybrid_query(query, self.top_k)
 
-# --- SYSTEM PROMPT & CORRECT PROMPT FORMAT ---
+# --- 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}")
@@ -157,7 +204,7 @@ prompt = ChatPromptTemplate.from_messages([
 
 llm = ChatOpenAI(model="gpt-4.1", openai_api_key=OPENAI_API_KEY, temperature=0)
 
-retriever = SQLiteVectorRetriever(top_k=5)
+retriever = HybridRetriever(top_k=5)
 qa_chain = RetrievalQA.from_chain_type(
     llm=llm,
     retriever=retriever,
@@ -165,7 +212,7 @@ qa_chain = RetrievalQA.from_chain_type(
     return_source_documents=True,
 )
 
-# --- Conversation Area (fine-tuned style) ---
+# --- 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":
@@ -204,7 +251,7 @@ def send_message():
         return
     st.session_state.messages.append({"role": "user", "content": user_input})
     with st.spinner("Thinking..."):
-        # Correct input key: "query" (not "question")
+        # Correct key: "query"
         result = qa_chain({"query": user_input})
         answer = result['result']
         st.session_state.messages.append({"role": "assistant", "content": answer})
@@ -214,12 +261,12 @@ def send_message():
             doc_list.append({
                 "file": doc.metadata["source_file"],
                 "id": doc.metadata["id"],
+                "similarity": doc.metadata["similarity"],
                 "record": json.loads(doc.metadata["raw_json"])
             })
         st.session_state.messages.append({"role": "function", "content": json.dumps(doc_list, indent=2)})
     st.session_state.temp_input = ""
 
-
 st.text_input("Your message:", key="temp_input", on_change=send_message)
 
 if st.button("Clear chat"):