Update app.py

app.py

@@ -6,12 +6,15 @@ import sqlite3
 import pandas as pd
 import numpy as np
 import datetime
-from typing import List, Dict
+from typing import List
 import openai
 from langchain.schema import Document
 from langchain.chains import RetrievalQA
 from langchain_community.llms import OpenAI as LangOpenAI
 
+# --- FIX: Import correct BaseRetriever
+from langchain_core.retrievers import BaseRetriever
+
 # ---- CONFIG ----
 OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
 EMBEDDING_MODEL = "text-embedding-ada-002"
@@ -35,7 +38,6 @@ if "last_entity" not in st.session_state:
 
 # ---- Helper: Flatten JSON ----
 def flatten_json_obj(obj, parent_key="", sep="."):
-    """Flatten nested JSON objects/lists with dot notation."""
     items = {}
     if isinstance(obj, dict):
         for k, v in obj.items():
@@ -98,7 +100,6 @@ def ingest_json_files(files):
         if isinstance(raw, list):
             records = raw
         elif isinstance(raw, dict):
-            # If dict with a single main list, use it
             main_lists = [v for v in raw.values() if isinstance(v, list)]
             if main_lists:
                 records = main_lists[0]
@@ -108,7 +109,7 @@ def ingest_json_files(files):
             records = [raw]
         for rec in records:
             flat = flatten_json_obj(rec)
-            # Heuristic: add top-level "name"/"customer" fields for entity tracking
+            # Add entity keys if found
             if "customer" in rec and isinstance(rec["customer"], str):
                 first_name = rec["customer"].split("@")[0].replace(".", " ")
                 flat["customer_name"] = first_name
@@ -139,7 +140,6 @@ def query_vector_db(user_query, top_k=5):
             continue
         sim = float(np.dot(query_emb, db_emb) / (np.linalg.norm(query_emb) * np.linalg.norm(db_emb)))
         results.append((sim, row))
-    # Top K by similarity
     results = sorted(results, reverse=True, key=lambda x: x[0])[:top_k]
     docs = []
     for sim, row in results:
@@ -154,13 +154,17 @@ def query_vector_db(user_query, top_k=5):
     return docs
 
 # ---- LangChain Retriever Adapter ----
-class SQLiteVectorRetriever:
-    def get_relevant_documents(self, query):
-        return query_vector_db(query, top_k=5)
+class SQLiteVectorRetriever(BaseRetriever):
+    def __init__(self, top_k=5):
+        self.top_k = top_k
+        super().__init__()
+
+    def get_relevant_documents(self, query: str) -> List[Document]:
+        return query_vector_db(query, self.top_k)
 
 # ---- LangChain LLM & QA Chain ----
 llm = LangOpenAI(model_name="gpt-4.1", openai_api_key=OPENAI_API_KEY, temperature=0)
-retriever = SQLiteVectorRetriever()
+retriever = SQLiteVectorRetriever(top_k=5)
 qa_chain = RetrievalQA.from_chain_type(
     llm=llm,
     retriever=retriever,
@@ -190,7 +194,6 @@ def update_last_entity(doc):
         pass
 
 def render_json_links():
-    # Tiny inline [view JSON] links, expands in-place on click
     for key in st.session_state.json_links:
         info = st.session_state.json_link_details[key]
         label = info["label"]
@@ -206,7 +209,6 @@ def send_message():
     user_input = st.session_state.temp_input.strip()
     if not user_input:
         return
-    # Entity resolution for pronouns (he, his, etc.)
     pronoun = re.search(r"\b(he|his|him|her|she|their)\b", user_input, re.I)
     if st.session_state.last_entity and pronoun:
         q = f"For {st.session_state.last_entity}: {user_input}"
@@ -230,10 +232,9 @@ def send_message():
             link_keys.append(link_key)
         st.session_state.json_links = link_keys
         st.session_state.json_link_details = link_details
-        st.session_state.modal_link = None  # reset on every new message
+        st.session_state.modal_link = None
     st.session_state.temp_input = ""
 
-# ---- Chat Conversation Rendering ----
 for msg in st.session_state.messages:
     if msg["role"] == "user":
         st.markdown(f"<b style='color:#3575dd'>User:</b> <span style='color:#111'>{msg['content']}</span>", unsafe_allow_html=True)