Update app.py

app.py

@@ -1418,253 +1418,53 @@ def locator_ui():
 #  - Final UI glue: Reports page and main app routing
 # -------------------------------------------------------
 
-import os
-import io
-import zipfile
-import json
-import tempfile
-import base64
-import pickle
-from datetime import datetime
-import textwrap
-
-# LLM client (Groq)
-try:
-    from groq import Groq
-except Exception:
-    Groq = None
-
-# FAISS
-try:
-    import faiss
-except Exception:
-    faiss = None
-
-# PDF (ReportLab)
-from reportlab.lib import colors
-from reportlab.lib.pagesizes import A4, landscape
-from reportlab.lib.units import mm
-from reportlab.platypus import (
-    SimpleDocTemplate, Paragraph, Spacer, Table, TableStyle, Image as RLImage, PageBreak
-)
-from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
-
-# Matplotlib for embedding GSD plot images into PDFs
-import matplotlib.pyplot as plt
-
 # --------------------------
-# Basic secret checks (run early)
+# GeoMate Ask (RAG Chat with OCR) UI
 # --------------------------
-# Using st.secrets for HF Spaces; fallback to env
-GROQ_KEY = os.getenv("GROQ_API_KEY") or (st.secrets.get("GROQ_API_KEY") if "GROQ_API_KEY" in st.secrets else None)
-SERVICE_ACCOUNT = os.getenv("SERVICE_ACCOUNT") or (st.secrets.get("SERVICE_ACCOUNT") if "SERVICE_ACCOUNT" in st.secrets else None)
-EARTH_ENGINE_KEY = os.getenv("EARTH_ENGINE_KEY") or (st.secrets.get("EARTH_ENGINE_KEY") if "EARTH_ENGINE_KEY" in st.secrets else None)
-
-missing = []
-if not GROQ_KEY:
-    missing.append("GROQ_API_KEY")
-if not SERVICE_ACCOUNT:
-    missing.append("SERVICE_ACCOUNT")
-if not EARTH_ENGINE_KEY:
-    # Don't require EARTH_ENGINE_KEY strictly (locator will warn) but warn user
-    st.warning("EARTH_ENGINE_KEY not found in secrets. Locator functions will be limited or disabled.")
-
-if missing:
-    st.error(f"Missing required secrets: {', '.join(missing)}. Please add them to Hugging Face Secrets and reload the app.")
-    st.stop()
-
-# Initialize Groq client
-if Groq is None:
-    st.warning("Groq client lib not installed (pip package 'groq'). RAG will be disabled.")
-    GROQ_CLIENT = None
-else:
-    try:
-        GROQ_CLIENT = Groq(api_key=GROQ_KEY)
-    except Exception as e:
-        st.error(f"Failed to create Groq client: {e}")
-        GROQ_CLIENT = None
-
-# --------------------------
-# Session-state helpers (ensure present)
-# --------------------------
-if "sites" not in ss:
-    ss["sites"] = [{"Site Name": "Home", "chat_history": [], "classifier_inputs": {}, "classifier_state": 0,
-                    "GSD": None, "USCS": None, "AASHTO": None, "GI": None,
-                    "map_snapshot": None, "report_convo_state": 0, "soil_profile": None}]
-
-if "active_site_idx" not in ss:
-    ss["active_site_idx"] = 0
-
-def get_active_site():
-    idx = ss.get("active_site_idx", 0)
-    sites = ss.get("sites", [])
-    if idx < 0 or idx >= len(sites):
-        ss["active_site_idx"] = 0
-        idx = 0
-    return sites[idx]
-
-def save_active_site(site):
-    idx = ss.get("active_site_idx", 0)
-    ss["sites"][idx] = site
-    # ensure persistence in session_state
-    ss.modified = True
-
-def add_site(name="New Site"):
-    if len(ss["sites"]) >= 4:
-        st.warning("Maximum 4 sites allowed.")
-        return
-    ss["sites"].append({"Site Name": name, "chat_history": [], "classifier_inputs": {}, "classifier_state": 0,
-                       "GSD": None, "USCS": None, "AASHTO": None, "GI": None,
-                       "map_snapshot": None, "report_convo_state": 0, "soil_profile": None})
-    ss["active_site_idx"] = len(ss["sites"]) - 1
-
-def remove_site(idx):
-    if 0 <= idx < len(ss["sites"]):
-        ss["sites"].pop(idx)
-        ss["active_site_idx"] = max(0, ss["active_site_idx"] - 1)
-
-# --------------------------
-# FAISS DB loader (expects a zip containing index.faiss + meta.pkl)
-# --------------------------
-def load_faiss_db_from_zip(zip_path: str):
-    """
-    Accepts path to a zip file containing index.faiss and meta.pkl (or index.pkl).
-    Extracts to tempdir and loads faiss index + metadata list.
-    Returns (index, meta_list) or (None, None) on error.
-    """
-    if faiss is None:
-        st.error("faiss not installed. RAG unavailable.")
-        return None, None
-
-    if not os.path.exists(zip_path):
-        st.error("FAISS DB zip not found at provided path.")
-        return None, None
-    tmpd = tempfile.mkdtemp()
-    try:
-        with zipfile.ZipFile(zip_path, "r") as z:
-            z.extractall(tmpd)
-        # look for index.faiss or index.pkl
-        idx_file = None
-        meta_file = None
-        for fname in os.listdir(tmpd):
-            if fname.endswith(".faiss") or fname == "index.faiss":
-                idx_file = os.path.join(tmpd, fname)
-            if fname.endswith(".pkl") or fname == "meta.pkl":
-                meta_file = os.path.join(tmpd, fname)
-        if idx_file is None or meta_file is None:
-            st.error("Zip did not contain index.faiss and meta.pkl.")
-            return None, None
-        index = faiss.read_index(idx_file)
-        with open(meta_file, "rb") as f:
-            meta = pickle.load(f)
-        return index, meta
-    except Exception as e:
-        st.error(f"Failed to load FAISS DB: {e}")
-        return None, None
-
-# --------------------------
-# RAG helper: simple retriever + Groq caller
-# --------------------------
-def rag_retrieve_and_answer(query: str, topk: int = 5):
-    """
-    Retrieve from FAISS (if loaded) and call Groq with context.
-    """
-    # Load index from session if present
-    if "faiss_index" not in ss or ss.get("faiss_index") is None:
-        st.error("FAISS index not loaded. Please upload faiss_books_db.zip in GeoMate Ask page.")
-        return "FAISS DB missing."
-
-    index = ss["faiss_index"]
-    meta = ss["faiss_meta"]
-    # prepare embedding for query — simplified: use Groq's embed endpoint if available.
-    # For demo, we'll simply return top-K metadata concatenated as context
-    try:
-        # nearest neighbor search using user-provided embeddings if available
-        # Here: assume index is IndexFlatL2 and meta is list
-        # For safety, we will not compute embeddings here (requires sentence-transformers).
-        # We'll do a cheap fallback: return top K metas (or first K) as context.
-        context_texts = [m.get("text","") for m in (meta[:topk] if isinstance(meta, list) else meta)]
-        context = "\n\n".join(context_texts)
-    except Exception:
-        context = ""
-
-    # Build prompt for Groq
-    system = "You are GeoMate RAG assistant. Use the context to answer precisely and professionally."
-    user_prompt = f"Context:\n{context}\n\nQuestion: {query}\nAnswer concisely and cite context sections if relevant."
+import pytesseract
+from PIL import Image
+import fitz  # PyMuPDF for PDF OCR
 
-    if GROQ_CLIENT is None:
-        # fallback: simple echo
-        return "Groq client not available. Cannot complete RAG call."
+def extract_text_from_file(uploaded_file):
+    """OCR utility for images and PDFs"""
     try:
-        completion = GROQ_CLIENT.chat.completions.create(
-            model="meta-llama/llama-4-maverick-17b-128e-instruct",
-            messages=[{"role":"system","content":system},{"role":"user","content":user_prompt}],
-            temperature=0.2,
-            max_tokens=800
-        )
-        text = completion.choices[0].message.content
-        return text
+        if uploaded_file.type in ["image/png", "image/jpeg", "image/jpg"]:
+            img = Image.open(uploaded_file)
+            text = pytesseract.image_to_string(img)
+            return text.strip()
+        elif uploaded_file.type == "application/pdf":
+            text_pages = []
+            pdf = fitz.open(stream=uploaded_file.read(), filetype="pdf")
+            for page in pdf:
+                text_pages.append(page.get_text("text"))
+                # fallback to OCR if no text
+                if not text_pages[-1].strip():
+                    pix = page.get_pixmap()
+                    img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
+                    text_pages[-1] = pytesseract.image_to_string(img)
+            return "\n".join(text_pages).strip()
+        else:
+            return ""
     except Exception as e:
-        st.error(f"Groq call failed: {e}")
-        return "RAG call failed."
-
-# --------------------------
-# Entity extraction placeholder
-# --------------------------
-def update_site_description_from_text(site: dict, text: str) -> dict:
-    """
-    Very simple regex-based extraction for core engineering parameters from free text.
-    Called after each RAG/chat answer: extracts numbers like 'bearing capacity 2000 psf' etc.
-    This is a placeholder; can be replaced by a proper NER model.
-    """
-    import re
-    # load-bearing capacity (psf or kPa)
-    m = re.search(r"bearing capacity\s*(?:of)?\s*([0-9,.]+)\s*(kpa|psf|pa|kn/m2)?", text, re.IGNORECASE)
-    if m:
-        val = m.group(1).replace(",", "")
-        unit = m.group(2) or ""
-        site["Load Bearing Capacity"] = f"{val} {unit}".strip()
-
-    # skin shear strength
-    m2 = re.search(r"skin shear strength\s*(?:[:is]*)\s*([0-9,.]+)\s*(kpa|kn/m2|psf)?", text, re.IGNORECASE)
-    if m2:
-        site["Skin Shear Strength"] = f"{m2.group(1).replace(',','')} { (m2.group(2) or '')}".strip()
-
-    # % compaction
-    m3 = re.search(r"compaction\s*(?:[:is]*)\s*([0-9]{1,3})\s*%", text, re.IGNORECASE)
-    if m3:
-        site["Relative Compaction"] = f"{m3.group(1)}%"
-
-    # rate of consolidation
-    m4 = re.search(r"consolidation rate\s*(?:[:is]*)\s*([0-9,.]+)\s*(mm/year|mm/yr|mm per year|m/year)?", text, re.IGNORECASE)
-    if m4:
-        site["Rate of Consolidation"] = f"{m4.group(1)} {m4.group(2) or ''}".strip()
-
-    # nature of construction - look for keywords
-    if "residential" in text.lower():
-        site["Nature of Construction"] = "Residential"
-    elif "commercial" in text.lower():
-        site["Nature of Construction"] = "Commercial"
-    elif "pavement" in text.lower() or "road" in text.lower():
-        site["Nature of Construction"] = "Pavement / Road"
-
-    return site
+        st.error(f"OCR failed: {e}")
+        return ""
 
-# --------------------------
-# GeoMate Ask (RAG Chat) UI
-# --------------------------
 def rag_ui():
-    st.header("🤖 GeoMate Ask — RAG + Groq (per-site memory)")
+    st.header("🤖 GeoMate Ask — RAG + Groq (per-site memory + OCR)")
     site = get_active_site()
 
     if "chat_history" not in site:
         site["chat_history"] = []
 
-    st.markdown("**Context:** The RAG uses your FAISS knowledge base (upload .zip in this page) and Groq LLM for answers. Chat history is saved for this site during the session.")
+    st.markdown(
+        "**Context:** The RAG uses your FAISS knowledge base (upload .zip in this page), "
+        "Groq LLM for answers, and optional OCR from uploaded images/PDFs. "
+        "Chat history is saved per site."
+    )
 
     # FAISS DB upload (one-time)
-    with st.expander("FAISS DB (index.faiss + meta.pkl inside a zip)"):
-        uploaded = st.file_uploader("Upload faiss_books_db.zip", type=["zip"])
+    with st.expander("📂 Upload FAISS DB (zip with index.faiss + meta.pkl)"):
+        uploaded = st.file_uploader("Upload faiss_books_db.zip", type=["zip"], key="faiss_db_uploader")
         if uploaded:
             tmpf = tempfile.NamedTemporaryFile(delete=False, suffix=".zip")
             tmpf.write(uploaded.getvalue())
@@ -1673,46 +1473,63 @@ def rag_ui():
             if ix is not None:
                 ss["faiss_index"] = ix
                 ss["faiss_meta"] = meta
-                st.success("FAISS DB loaded.")
+                st.success("✅ FAISS DB loaded.")
 
     # Render chat history
     for turn in site.get("chat_history", []):
-        role = turn.get("role")
-        text = turn.get("text")
+        role, text = turn.get("role"), turn.get("text")
         if role == "bot":
-            st.markdown(f"<div style='background:{THEME['bubble_bg']};padding:8px;border-radius:12px;border:2px solid {THEME['accent']};'><b>🤖 GeoMate:</b> {text}</div>", unsafe_allow_html=True)
+            st.markdown(
+                f"<div style='background:{THEME['bubble_bg']};padding:8px;border-radius:12px;border:2px solid {THEME['accent']};'>"
+                f"<b>🤖 GeoMate:</b> {text}</div>", unsafe_allow_html=True
+            )
         else:
-            st.markdown(f"<div style='background:#1a2436;color:#fff;padding:8px;border-radius:12px;margin-left:40px;'><b>👤 You:</b> {text}</div>", unsafe_allow_html=True)
+            st.markdown(
+                f"<div style='background:#1a2436;color:#fff;padding:8px;border-radius:12px;margin-left:40px;'>"
+                f"<b>👤 You:</b> {text}</div>", unsafe_allow_html=True
+            )
+
+    # Input + OCR
+    user_q = st.text_input("Ask GeoMate:", key="geomate_rag_input")
+    uploaded_ocr = st.file_uploader("Optional OCR input (image/pdf)", type=["png","jpg","jpeg","pdf"], key="rag_ocr_uploader")
 
-    # Input box
-    user_q = st.text_input("Ask GeoMate (RAG + site memory):", key="geomate_rag_input")
     if st.button("Ask", key="geomate_rag_button"):
-        if not user_q.strip():
-            st.warning("Type a question first.")
+        if not user_q.strip() and not uploaded_ocr:
+            st.warning("Please type a question or upload a file.")
         else:
-            # Append to history
-            site["chat_history"].append({"role":"user","text":user_q, "time":datetime.utcnow().isoformat()})
+            query_text = user_q.strip()
+            if uploaded_ocr:
+                with st.spinner("Running OCR..."):
+                    ocr_text = extract_text_from_file(uploaded_ocr)
+                    if ocr_text:
+                        query_text += "\n\n[OCR Extracted Content]\n" + ocr_text
+
+            # Append user query
+            site["chat_history"].append({"role":"user","text":query_text, "time":datetime.utcnow().isoformat()})
             save_active_site(site)
+
             # Retrieve + call LLM
             with st.spinner("Retrieving context and calling LLM..."):
-                answer = rag_retrieve_and_answer(user_q, topk=5)
+                answer = rag_retrieve_and_answer(query_text, topk=5)
+
             # Append bot answer
             site["chat_history"].append({"role":"bot","text":answer, "time":datetime.utcnow().isoformat()})
-            # Try to extract any engineering parameters from answer or the user question
-            site = update_site_description_from_text(site, user_q + "\n" + answer)
+
+            # Extract parameters
+            site = update_site_description_from_text(site, query_text + "\n" + answer)
             save_active_site(site)
-            st.experimental_rerun()  # small rerun to display new messages
+
+            st.rerun()
 
     # Quick buttons
     colA, colB = st.columns(2)
-    if colA.button("Save Chat to Site JSON"):
+    if colA.button("💾 Save Chat"):
         save_active_site(site)
-        st.success("Saved chat into site JSON.")
-    if colB.button("Clear Site Chat"):
+        st.success("Chat saved into site JSON.")
+    if colB.button("🗑️ Clear Chat"):
         site["chat_history"] = []
         save_active_site(site)
         st.success("Cleared history for this site.")
-
 # --------------------------
 # REPORTS: PDF builders
 # --------------------------