Update app.py

app.py

@@ -1,180 +1,43 @@
 import streamlit as st
 import os
 import time
-import re
-import requests
-from PIL import Image
-from io import BytesIO
-from urllib.parse import quote
+import json
 from openai import OpenAI
 
-# ------------------ App Configuration ------------------
-st.set_page_config(page_title="Schlaeger Forrestdale DocAIA", layout="wide", initial_sidebar_state="collapsed")
-st.title("📄 Schlaeger Forrestdale Document Assistant")
-st.caption("Explore City of Armadale construction documents using AI + OCR 🧠")
+# Configuration
+st.set_page_config(page_title="Forrestdale Drawing Viewer", layout="wide")
+st.title("📐 Forrestdale Technical Drawing Assistant")
 
-# ------------------ Load API Key and Assistant IDs ------------------
 OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY")
-CONTRACT_ASSISTANT_ID = os.environ.get("ASSISTANT_ID")  # Contract assistant
-TECH_ASSISTANT_ID = "asst_DjvuWBc7tCvMbAhY7n1em4BZ"       # Technical drawings assistant
+ASSISTANT_ID = "asst_DjvuWBc7tCvMbAhY7n1em4BZ"
 
-if not OPENAI_API_KEY or not CONTRACT_ASSISTANT_ID:
-    st.error("❌ Missing secrets. Please set OPENAI_API_KEY and ASSISTANT_ID in Hugging Face Space secrets.")
+if not OPENAI_API_KEY:
+    st.error("❌ Missing OPENAI_API_KEY.")
     st.stop()
 
 client = OpenAI(api_key=OPENAI_API_KEY)
 
-# ------------------ Tabs ------------------
-tab1, tab2 = st.tabs(["📄 Contract Assistant", "📐 Technical Drawings Assistant"])
+if "tech_thread_id" not in st.session_state:
+    thread = client.beta.threads.create()
+    st.session_state.tech_thread_id = thread.id
 
-# ------------------ Tab 1: Contract Assistant ------------------
-with tab1:
-    if "messages" not in st.session_state:
-        st.session_state.messages = []
-    if "thread_id" not in st.session_state:
-        st.session_state.thread_id = None
-    if "image_url" not in st.session_state:
-        st.session_state.image_url = None
-    if "image_updated" not in st.session_state:
-        st.session_state.image_updated = False
-    if "pending_prompt" not in st.session_state:
-        st.session_state.pending_prompt = None
+if "tech_messages" not in st.session_state:
+    st.session_state.tech_messages = []
 
-    show_image = st.sidebar.toggle("📑 Show Page Image", value=True)
+# Prompt input
+prompt = st.chat_input("Ask about plans, drawings or components (e.g. Show me all electrical plans)")
+if prompt:
+    st.session_state.tech_messages.append({"role": "user", "content": prompt})
 
-    st.sidebar.subheader("📘 Document Tools")
-    keyword = st.sidebar.text_input("Search by Keyword", placeholder="e.g. defects, WHS, delay")
-    if st.sidebar.button("🔎 Search Keyword") and keyword:
-        st.session_state.pending_prompt = f"Find clauses or references related to: {keyword}"
+# Display assistant messages
+for msg in st.session_state.tech_messages:
+    with st.chat_message(msg["role"]):
+        st.markdown(msg["content"])
 
-    section_options = [
-        "Select a section...",
-        "1. Formal Instrument of Contract",
-        "2. Offer and Acceptance",
-        "3. Key Personnel",
-        "4. Contract Pricing",
-        "5. Specifications",
-        "6. WHS Policies",
-        "7. Penalties and Delays",
-        "8. Dispute Resolution",
-        "9. Principal Obligations"
-    ]
-    section_select = st.sidebar.selectbox("📄 Jump to Section", section_options)
-    if section_select != section_options[0]:
-        st.session_state.pending_prompt = f"Summarize or list key points from section: {section_select}"
-
-    actions = [
-        "Select an action...",
-        "List all contractual obligations",
-        "Summarize payment terms",
-        "List WHS responsibilities",
-        "Find delay-related penalties",
-        "Extract dispute resolution steps"
-    ]
-    action_select = st.sidebar.selectbox("⚙️ Common Contract Queries", actions)
-    if action_select != actions[0]:
-        st.session_state.pending_prompt = action_select
-
-    chat_col, image_col = st.columns([2, 1])
-
-    with chat_col:
-        st.markdown("### 🧠 Ask a Document-Specific Question")
-        user_prompt = st.chat_input("Example: What is the defects liability period?")
-
-        if user_prompt:
-            st.session_state.messages.append({"role": "user", "content": user_prompt})
-        elif st.session_state.pending_prompt:
-            st.session_state.messages.append({"role": "user", "content": st.session_state.pending_prompt})
-            st.session_state.pending_prompt = None
-
-        if st.session_state.messages and st.session_state.messages[-1]["role"] == "user":
-            try:
-                if st.session_state.thread_id is None:
-                    thread = client.beta.threads.create()
-                    st.session_state.thread_id = thread.id
-
-                client.beta.threads.messages.create(
-                    thread_id=st.session_state.thread_id,
-                    role="user",
-                    content=st.session_state.messages[-1]["content"]
-                )
-
-                run = client.beta.threads.runs.create(
-                    thread_id=st.session_state.thread_id,
-                    assistant_id=CONTRACT_ASSISTANT_ID
-                )
-
-                with st.spinner("🤖 Parsing and responding with referenced content..."):
-                    while True:
-                        run_status = client.beta.threads.runs.retrieve(
-                            thread_id=st.session_state.thread_id,
-                            run_id=run.id
-                        )
-                        if run_status.status in ("completed", "failed", "cancelled"):
-                            break
-                        time.sleep(1)
-
-                if run_status.status != "completed":
-                    st.error(f"⚠️ Assistant failed: {run_status.status}")
-                else:
-                    messages = client.beta.threads.messages.list(thread_id=st.session_state.thread_id)
-                    for message in reversed(messages.data):
-                        if message.role == "assistant":
-                            assistant_reply = message.content[0].text.value
-                            st.session_state.messages.append({"role": "assistant", "content": assistant_reply})
-
-                            match = re.search(r'Document Reference:\s*(.*?),\s*Page\s*(\d+)', assistant_reply)
-                            if match:
-                                doc_name = match.group(1).strip()
-                                page = int(match.group(2))
-                                page_str = f"{page:04d}"
-                                folder = quote(doc_name)
-                                image_url = (
-                                    f"https://raw.githubusercontent.com/AndrewLORTech/c2ozschlaegerforrestdale/main/"
-                                    f"{folder}/{folder}_page_{page_str}.png"
-                                )
-                                st.session_state.image_url = image_url
-                                st.session_state.image_updated = True
-                            break
-
-                st.rerun()
-            except Exception as e:
-                st.error(f"❌ Error: {e}")
-
-        for msg in st.session_state.messages:
-            with st.chat_message(msg["role"]):
-                st.markdown(msg["content"], unsafe_allow_html=True)
-
-    with image_col:
-        if show_image and st.session_state.image_url:
-            with st.spinner("Loading document preview..."):
-                try:
-                    response = requests.get(st.session_state.image_url)
-                    response.raise_for_status()
-                    img = Image.open(BytesIO(response.content))
-                    st.image(img, caption="📄 OCR Page Image", use_container_width=True)
-                    st.session_state.image_updated = False
-                except Exception as e:
-                    st.error(f"❗ Failed to load image: {e}")
-
-# ------------------ Tab 2: Technical Drawing Assistant ------------------
-with tab2:
-    if "tech_messages" not in st.session_state:
-        st.session_state.tech_messages = []
-        st.session_state.tech_thread_id = None
-
-    st.markdown("### 🏗️ Ask a Question About Drawings or Diagrams")
-    tech_prompt = st.chat_input("Example: Show me all architectural drawings")
-
-    if tech_prompt:
-        st.session_state.tech_messages.append({"role": "user", "content": tech_prompt})
-
-    if st.session_state.tech_messages and st.session_state.tech_messages[-1]["role"] == "user":
+# Trigger assistant
+if st.session_state.tech_messages and st.session_state.tech_messages[-1]["role"] == "user":
+    with st.spinner("⏳ Fetching results from assistant..."):
         try:
-            if st.session_state.tech_thread_id is None:
-                thread = client.beta.threads.create()
-                st.session_state.tech_thread_id = thread.id
-
             client.beta.threads.messages.create(
                 thread_id=st.session_state.tech_thread_id,
                 role="user",
@@ -183,33 +46,45 @@ with tab2:
 
             run = client.beta.threads.runs.create(
                 thread_id=st.session_state.tech_thread_id,
-                assistant_id=TECH_ASSISTANT_ID
+                assistant_id=ASSISTANT_ID
             )
 
-            with st.spinner("🔍 Querying Technical Drawing Assistant..."):
-                while True:
-                    run_status = client.beta.threads.runs.retrieve(
-                        thread_id=st.session_state.tech_thread_id,
-                        run_id=run.id
-                    )
-                    if run_status.status in ("completed", "failed", "cancelled"):
-                        break
-                    time.sleep(1)
+            while True:
+                run_status = client.beta.threads.runs.retrieve(
+                    thread_id=st.session_state.tech_thread_id,
+                    run_id=run.id
+                )
+                if run_status.status in ["completed", "failed", "cancelled"]:
+                    break
+                time.sleep(1)
 
             if run_status.status != "completed":
-                st.error(f"⚠️ Assistant failed: {run_status.status}")
+                st.error("⚠️ Assistant run failed.")
             else:
                 messages = client.beta.threads.messages.list(thread_id=st.session_state.tech_thread_id)
                 for message in reversed(messages.data):
                     if message.role == "assistant":
-                        reply = message.content[0].text.value
-                        st.session_state.tech_messages.append({"role": "assistant", "content": reply})
+                        reply_content = message.content[0].text.value.strip()
+                        st.session_state.tech_messages.append({"role": "assistant", "content": reply_content})
+
+                        # Try to parse response as JSON
+                        try:
+                            results = json.loads(reply_content)
+                            if isinstance(results, list):
+                                for item in results:
+                                    st.markdown(f"### 📄 {item.get('drawing_number')} ({item.get('discipline')})")
+                                    st.markdown(f"**Summary:** {item.get('summary')}")
+
+                                    # Single image (keyword match)
+                                    if "image" in item:
+                                        st.image(item["image"], use_column_width=True)
+
+                                    # Multiple image pages
+                                    elif "images" in item:
+                                        for img in item["images"]:
+                                            st.image(img, use_column_width=True)
+                        except Exception as parse_error:
+                            st.warning("🟡 Could not parse assistant response as JSON.")
                         break
-
-            st.rerun()
         except Exception as e:
-            st.error(f"❌ Error: {e}")
-
-    for msg in st.session_state.tech_messages:
-        with st.chat_message(msg["role"]):
-            st.markdown(msg["content"], unsafe_allow_html=True)
+            st.error(f"❌ Error occurred: {e}")