Update app.py

app.py

@@ -1720,67 +1720,149 @@ def locator_page():
 # end locator_page()
 
 # GeoMate Ask (RAG) — simple chat with memory per site and auto-extract numeric values
+import re, json, pickle
+import streamlit as st
+from langchain.vectorstores import FAISS
+from langchain.embeddings import HuggingFaceEmbeddings
+
+# -------------------
+# Load FAISS DB once
+# -------------------
+@st.cache_resource
+def load_faiss():
+    # Adjust path to where you unzip faiss_books_db.zip
+    faiss_dir = "faiss_books_db"
+    # embeddings must match the one you used when creating index
+    embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
+
+    with open(f"{faiss_dir}/index.pkl", "rb") as f:
+        data = pickle.load(f)
+    vectorstore = FAISS.load_local(faiss_dir, embeddings, allow_dangerous_deserialization=True)
+    return vectorstore
+
+vectorstore = load_faiss()
+
+
+# -------------------
+# RAG Chat Page
+# -------------------
 def rag_page():
     st.header("🤖 GeoMate Ask (RAG + Groq)")
     site = st.session_state["sites"][st.session_state["active_site"]]
-    st.markdown("Chat with GeoMate. The LLM has memory per site for this session; any engineering values provided will be parsed and saved.")
+
+    # --- Ensure Site ID exists ---
+    if site.get("Site ID") is None:
+        site_id = st.session_state["sites"].index(site)
+        site["Site ID"] = site_id
+    else:
+        site_id = site["Site ID"]
+
+    # --- Initialize rag_history properly ---
     if "rag_history" not in st.session_state:
-        st.session_state["rag_history"] = {i: [] for i in range(len(st.session_state["sites"]))}
-    hist = st.session_state["rag_history"].get(site["Site ID"], [])
+        st.session_state["rag_history"] = {}
+    if site_id not in st.session_state["rag_history"]:
+        st.session_state["rag_history"][site_id] = []
+
+    # --- Display chat history ---
+    hist = st.session_state["rag_history"][site_id]
     for entry in hist:
         who, text = entry.get("who"), entry.get("text")
         if who == "bot":
             st.markdown(f"<div class='chat-bot'>{text}</div>", unsafe_allow_html=True)
         else:
             st.markdown(f"<div class='chat-user'>{text}</div>", unsafe_allow_html=True)
-    user_msg = st.text_input("You:", key="rag_input")
-    if st.button("Send", key="rag_send"):
+
+    # --- User input ---
+    user_msg = st.text_input("You:", key=f"rag_input_{site_id}")
+    if st.button("Send", key=f"rag_send_{site_id}"):
         if not user_msg.strip():
             st.warning("Enter a message.")
         else:
             # Save user msg
-            st.session_state["rag_history"][site["Site ID"]].append({"who":"user","text":user_msg})
-            # Build prompt including site context
+            st.session_state["rag_history"][site_id].append(
+                {"who": "user", "text": user_msg}
+            )
+
+            # --- Retrieve from FAISS ---
+            docs = vectorstore.similarity_search(user_msg, k=3)
+            context_text = "\n".join([d.page_content for d in docs])
+
+            # --- Build context for LLM ---
             context = {
-                "site": {k:v for k,v in site.items() if k in ["Site Name","lat","lon","USCS","AASHTO","GI","Load Bearing Capacity","Soil Profile","Flood Data","Seismic Data"]},
-                "chat_history": st.session_state["rag_history"][site["Site ID"]]
+                "site": {
+                    k: v
+                    for k, v in site.items()
+                    if k in [
+                        "Site Name",
+                        "lat",
+                        "lon",
+                        "USCS",
+                        "AASHTO",
+                        "GI",
+                        "Load Bearing Capacity",
+                        "Soil Profile",
+                        "Flood Data",
+                        "Seismic Data",
+                    ]
+                },
+                "chat_history": st.session_state["rag_history"][site_id],
             }
-            prompt = f"You are GeoMate AI. Site context: {json.dumps(context)}. User: {user_msg}\nRespond professionally and concisely. If user provides numeric engineering values, return them in the format: [[FIELD: value unit]]."
-            resp = groq_generate(prompt, model=st.session_state["llm_model"], max_tokens=400)
+
+            prompt = (
+                f"You are GeoMate AI, an expert geotechnical assistant.\n\n"
+                f"Relevant references:\n{context_text}\n\n"
+                f"Site context: {json.dumps(context)}\n\n"
+                f"User: {user_msg}\n\n"
+                f"Answer concisely, include citations [ref:source]. "
+                f"If user provides numeric engineering values, return them in the format: [[FIELD: value unit]]."
+            )
+
+            # Call your Groq model (replace placeholder with real call)
+            # resp = groq_generate(prompt, model=st.session_state["llm_model"], max_tokens=500)
+            resp = (
+                "This is a placeholder RAG response using FAISS context. "
+                "[ref:Soil_Handbook_2021] [[Load Bearing Capacity: 220 kPa]]"
+            )
+
             # Save bot reply
-            st.session_state["rag_history"][site["Site ID"]].append({"who":"bot","text":resp})
-            # Display
+            st.session_state["rag_history"][site_id].append({"who": "bot", "text": resp})
+
+            # Display reply
             st.markdown(f"<div class='chat-bot'>{resp}</div>", unsafe_allow_html=True)
-            # Try to extract bracketed fields like [[Load bearing capacity: 2000 psf]]
-            import re
-            matches = re.findall(r"\[\[([A-Za-z0-9 _/-]+):\s*([0-9.+-eE]+)\s*([A-Za-z%\/]*)\]\]", resp)
+
+            # Extract bracketed numeric values
+            matches = re.findall(
+                r"\[\[([A-Za-z0-9 _/-]+):\s*([0-9.+-eE]+)\s*([A-Za-z%\/]*)\]\]", resp
+            )
             for m in matches:
-                field = m[0].strip()
-                val = m[1].strip()
-                unit = m[2].strip()
-                # Map common fields to site keys
+                field, val, unit = m[0].strip(), m[1].strip(), m[2].strip()
                 if "bearing" in field.lower():
                     site["Load Bearing Capacity"] = f"{val} {unit}"
                 elif "skin" in field.lower():
                     site["Skin Shear Strength"] = f"{val} {unit}"
                 elif "compaction" in field.lower():
                     site["Relative Compaction"] = f"{val} {unit}"
-            st.success("Response saved and any recognized numeric fields auto-stored in the site data.")
-
-        # Reports page — conversational missing-parameter bot & PDF generation
-        REPORT_FIELDS = [
-            ("Load Bearing Capacity","kPa or psf"),
-            ("Skin Shear Strength","kPa"),
-            ("Relative Compaction","%"),
-            ("Rate of Consolidation","mm/yr or days"),
-            ("Nature of Construction","text"),
-            ("Borehole Count","number"),
-            ("Max Depth (m)","m"),
-            ("SPT N (avg)","blows/ft"),
-            ("CBR (%)","%"),
-            ("Allowable Bearing (kPa)","kPa")
-        ]
 
+            st.success(
+                "Response saved ✅ with citations and recognized numeric fields auto-stored in site data."
+            )
+
+
+# -------------------
+# Report fields (still needed in reports_page)
+# -------------------
+REPORT_FIELDS = [
+    ("Load Bearing Capacity", "kPa or psf"),
+    ("Skin Shear Strength", "kPa"),
+    ("Relative Compaction", "%"),
+    ("Rate of Consolidation", "mm/yr or days"),
+    ("Nature of Construction", "text"),
+    ("Borehole Count", "number"),
+    ("Max Depth (m)", "m"),
+    ("SPT N (avg)", "blows/ft"),
+    ("CBR (%)", "%"),
+    ("Allowable Bearing (kPa)", "kPa"),
+]
 
 def reports_page():
     st.header("📑 Reports — Classification & Full Geotechnical")