Update app.py

app.py

@@ -14,12 +14,10 @@ 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:
@@ -38,13 +36,11 @@ uploaded_files = st.file_uploader(
     "Upload JSON files in batches (any structure)", type="json", accept_multiple_files=True
 )
 
-# --- Enhanced Flattening: extract names from emails/user fields for LLM context
 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(s) from email/user
             if (
                 k.lower() in {"customer", "user", "email", "username"} and
                 isinstance(v, str) and "@" in v
@@ -125,7 +121,40 @@ def ingest_json_files(files):
 if uploaded_files and st.button("Ingest batch to database"):
     ingest_json_files(uploaded_files)
 
-# --- VECTOR RETRIEVAL
+# --- Improved entity search/filter
+def extract_main_entity(question):
+    # crude: get the first capitalized word, or all words
+    tokens = re.findall(r"\b([A-Za-z0-9]+)\b", question)
+    keywords = [t.lower() for t in tokens if t.lower() not in {"how", "much", "did", "spend", "was", "the", "is", "in", "on", "for", "a", "an", "of", "to", "with"}]
+    # e.g. ["johnny", "spend"] → "johnny"
+    return keywords[0] if keywords else None
+
+def filter_records_by_entity(records, entity):
+    matches = []
+    for doc in records:
+        if entity and entity in doc.page_content.lower():
+            matches.append(doc)
+    return matches if matches else records
+
+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)
+    all_docs = []
+    seen_ids = set()
+    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)
+    # Filter for entity match if possible
+    entity = extract_main_entity(user_query)
+    entity_docs = filter_records_by_entity(all_docs, entity) if entity else all_docs
+    # Optionally, highlight the entity in the flat_text for the LLM
+    for doc in entity_docs:
+        if entity:
+            doc.page_content = re.sub(rf"({re.escape(entity)})", r"**\1**", doc.page_content, flags=re.IGNORECASE)
+    return entity_docs[:top_k]
+
 def query_vector_db(user_query, top_k=5):
     query_emb = get_embedding(user_query)
     conn = sqlite3.connect(DB_PATH)
@@ -151,7 +180,6 @@ def query_vector_db(user_query, top_k=5):
         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)
@@ -177,45 +205,23 @@ def python_fuzzy_match(user_query, top_k=5):
         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
+# --- Prompt (explicitly tells LLM what to do)
 system_prompt = (
-    "You are a JSON data assistant. Always answer using only the provided context records. "
-    "If a question mentions a person or entity (e.g. 'Johnny'), find all records where any field includes that name "
-    "(including as part of an email or username, e.g. customer_name: johnny, customer: johnny.appleseed@gmail.com). "
-    "If you find a relevant record, answer as directly as possible using its fields (e.g. 'Johnny spent $100'). "
-    "If you cannot find the answer, reply: 'I don’t have that information.'"
-    "If the user asks for the number or details of items in a list/array (e.g., completed tasks), use 'find_in_arrays'. "
-    "If the user asks about the sum/total of a field for a name or identifier, use 'sum_field_by_name'. "
-    "If the user asks about female names, use 'count_female_names'. "
-    "If the user's query does not mention a key, use 'fuzzy_value_search' to match on any value. "
-    "If a key is mentioned (like 'apps_installed'), use 'search_all_jsons' for that key and the value. "
-    "You may use 'list_keys' to help discover the file structure if needed. "
-    "Always give a direct answer from the data if possible."
+    "You are a JSON data assistant. "
+    "If a question mentions a name (like Johnny), find any record where that name appears as part of any field value (including emails or usernames). "
+    "Use the provided records to answer directly. If you can't find the answer, reply: 'I don’t have that information.' "
+    "Never make up data. Never ask for clarification."
 )
 prompt = ChatPromptTemplate.from_messages([
     ("system", system_prompt),
     ("human", "Here are the most relevant records:\n{context}\n\nQuestion: {question}")
 ])
 
-
 llm = ChatOpenAI(model="gpt-4.1", openai_api_key=OPENAI_API_KEY, temperature=0)
 
 retriever = HybridRetriever(top_k=5)
@@ -226,7 +232,6 @@ qa_chain = RetrievalQA.from_chain_type(
     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":
@@ -236,29 +241,6 @@ for msg in st.session_state.messages:
     elif msg["role"] == "function":
         st.markdown(f"<details><summary><b>Function Output:</b></summary><pre>{msg['content']}</pre></details>", unsafe_allow_html=True)
 
-def show_json_links_and_modal():
-    for msg in reversed(st.session_state.messages):
-        if msg.get("role") == "function" and msg.get("content"):
-            try:
-                docs = json.loads(msg["content"])
-                if isinstance(docs, list):
-                    for idx, doc in enumerate(docs):
-                        if isinstance(doc, dict) and "record" in doc:
-                            if st.button(f"View JSON: {doc.get('file', 'unknown')} record #{idx+1}", key=f"modal_function_{idx}"):
-                                st.session_state.modal_open = True
-                                st.session_state.modal_content = json.dumps(doc["record"], indent=2)
-                                st.session_state.modal_title = f"{doc.get('file', 'unknown')} record #{idx+1}"
-            except Exception:
-                continue
-            break
-    if st.session_state.modal_open:
-        with st.expander(f"JSON Record: {st.session_state.modal_title}", expanded=True):
-            st.code(st.session_state.modal_content, language="json")
-            if st.button("Close", key="close_modal"):
-                st.session_state.modal_open = False
-
-show_json_links_and_modal()
-
 def send_message():
     user_input = st.session_state.temp_input.strip()
     if not user_input: