Update app.py

app.py

@@ -10,76 +10,48 @@ from urllib.parse import quote
 from openai import OpenAI
 
 # ------------------ App Configuration ------------------
-st.set_page_config(page_title="Schlaeger Forrestdale DocAIA", layout="wide", initial_sidebar_state="collapsed")
+st.set_page_config(page_title="Schlaeger Forrestdale TechDocAIA", layout="wide", initial_sidebar_state="collapsed")
 st.title("📄 Schlaeger Forrestdale Document Assistant")
-st.caption("Explore City of Armadale construction documents using AI + OCR 🧠")
+st.caption("Explore City of Armadale construction documents using AI + OCR 🧐")
 
 # ------------------ Load API Key ------------------
 OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY")
 if not OPENAI_API_KEY:
-    st.error("❌ Missing OPENAI_API_KEY in Hugging Face Space secrets.")
+    st.error("❌ Missing OPENAI_API_KEY. Please set it in Hugging Face Space secrets.")
     st.stop()
 
 client = OpenAI(api_key=OPENAI_API_KEY)
 
-# ------------------ Tabs ------------------
-tabs = st.tabs(["📜 Contract", "📐 Technical"])
+# ------------------ Tabs Setup ------------------
+tabs = st.tabs(["📑 Contract", "📐 Technical"])
 
-# ================== CONTRACT TAB ==================
+# ------------------ Contract Tab ------------------
 with tabs[0]:
-    ASSISTANT_ID = "asst_KsQRedoJUnEeStzfox1o06lO"
+    from urllib.parse import quote
+
+    ASSISTANT_ID_CONTRACT = "asst_KsQRedoJUnEeStzfox1o06lO"
+
     if "contract_messages" not in st.session_state:
         st.session_state.contract_messages = []
     if "contract_thread_id" not in st.session_state:
         st.session_state.contract_thread_id = None
-    if "image_url" not in st.session_state:
-        st.session_state.image_url = None
-
-    st.sidebar.header("📘 Contract Tools")
-    if st.sidebar.button("🧹 Clear Chat", key="clear_contract"):
-        st.session_state.contract_messages = []
-        st.session_state.contract_thread_id = None
-        st.session_state.image_url = None
-        st.rerun()
-
-    show_image = st.sidebar.toggle("📑 Show Page Image", value=True, key="show_image_toggle")
-
-    keyword = st.sidebar.text_input("Search by Keyword", key="kw")
-    if st.sidebar.button("🔎 Search", key="kw_btn") and keyword:
-        st.session_state.contract_messages.append({"role": "user", "content": f"Find clauses or references related to: {keyword}"})
-
-    section = st.sidebar.selectbox("📄 Jump to Section", [
-        "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"
-    ])
-    if section != "Select a section...":
-        st.session_state.contract_messages.append({"role": "user", "content": f"Summarize or list key points from section: {section}"})
-
-    action = st.sidebar.selectbox("⚙️ Common Queries", [
-        "Select an action...",
-        "List all contractual obligations",
-        "Summarize payment terms",
-        "List WHS responsibilities",
-        "Find delay-related penalties",
-        "Extract dispute resolution steps"
-    ])
-    if action != "Select an action...":
-        st.session_state.contract_messages.append({"role": "user", "content": action})
-
-    chat_col, img_col = st.columns([2, 1])
-    with chat_col:
-        st.markdown("### 🧠 Ask Contract Document Question")
-        user_input = st.chat_input("Ask something about the contract")
-        if user_input:
-            st.session_state.contract_messages.append({"role": "user", "content": user_input})
+    if "contract_image_url" not in st.session_state:
+        st.session_state.contract_image_url = None
+    if "contract_image_updated" not in st.session_state:
+        st.session_state.contract_image_updated = False
+    if "contract_pending_prompt" not in st.session_state:
+        st.session_state.contract_pending_prompt = None
+
+    col1, col2 = st.columns([2, 1])
+    with col1:
+        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.contract_messages.append({"role": "user", "content": user_prompt})
+        elif st.session_state.contract_pending_prompt:
+            st.session_state.contract_messages.append({"role": "user", "content": st.session_state.contract_pending_prompt})
+            st.session_state.contract_pending_prompt = None
 
         if st.session_state.contract_messages and st.session_state.contract_messages[-1]["role"] == "user":
             try:
@@ -92,38 +64,43 @@ with tabs[0]:
                     role="user",
                     content=st.session_state.contract_messages[-1]["content"]
                 )
+
                 run = client.beta.threads.runs.create(
                     thread_id=st.session_state.contract_thread_id,
-                    assistant_id=ASSISTANT_ID
+                    assistant_id=ASSISTANT_ID_CONTRACT
                 )
 
-                with st.spinner("🤖 Analyzing contract..."):
+                with st.spinner("🤖 Parsing and responding with referenced content..."):
                     while True:
-                        status = client.beta.threads.runs.retrieve(
+                        run_status = client.beta.threads.runs.retrieve(
                             thread_id=st.session_state.contract_thread_id,
                             run_id=run.id
                         )
-                        if status.status in ("completed", "failed", "cancelled"):
+                        if run_status.status in ("completed", "failed", "cancelled"):
                             break
                         time.sleep(1)
 
-                if status.status == "completed":
+                if run_status.status == "completed":
                     messages = client.beta.threads.messages.list(thread_id=st.session_state.contract_thread_id)
-                    for m in reversed(messages.data):
-                        if m.role == "assistant":
-                            content = m.content[0].text.value
-                            st.session_state.contract_messages.append({"role": "assistant", "content": content})
+                    for message in reversed(messages.data):
+                        if message.role == "assistant":
+                            assistant_reply = message.content[0].text.value
+                            st.session_state.contract_messages.append({"role": "assistant", "content": assistant_reply})
 
-                            match = re.search(r'Document Reference:\s*(.*?),\s*Page\s*(\d+)', content)
+                            match = re.search(r'Document Reference:\s*(.*?),\s*Page\s*(\d+)', assistant_reply)
                             if match:
-                                doc, page = match.group(1).strip(), int(match.group(2))
+                                doc_name = match.group(1).strip()
+                                page = int(match.group(2))
                                 page_str = f"{page:04d}"
-                                folder = quote(doc)
-                                st.session_state.image_url = f"https://raw.githubusercontent.com/AndrewLORTech/c2ozschlaegerforrestdale/main/{folder}/{folder}_page_{page_str}.png"
+                                folder = quote(doc_name)
+                                image_url = f"https://raw.githubusercontent.com/AndrewLORTech/c2ozschlaegerforrestdale/main/{folder}/{folder}_page_{page_str}.png"
+                                st.session_state.contract_image_url = image_url
+                                st.session_state.contract_image_updated = True
                             break
                 else:
-                    st.error("❌ Assistant failed.")
+                    st.error(f"⚠️ Assistant failed: {run_status.status}")
                 st.rerun()
+
             except Exception as e:
                 st.error(f"❌ Error: {e}")
 
@@ -131,31 +108,32 @@ with tabs[0]:
             with st.chat_message(msg["role"]):
                 st.markdown(msg["content"], unsafe_allow_html=True)
 
-    with img_col:
-        if show_image and st.session_state.image_url:
-            with st.spinner("Loading image..."):
-                try:
-                    res = requests.get(st.session_state.image_url)
-                    img = Image.open(BytesIO(res.content))
-                    st.image(img, caption="📄 OCR Page Image", use_container_width=True)
-                except:
-                    st.error("⚠️ Failed to load image.")
-
-# ================== TECHNICAL TAB ==================
+    with col2:
+        if st.session_state.contract_image_url:
+            try:
+                response = requests.get(st.session_state.contract_image_url)
+                response.raise_for_status()
+                img = Image.open(BytesIO(response.content))
+                st.image(img, caption="📄 OCR Page Image", use_container_width=True)
+            except Exception as e:
+                st.error(f"❗ Failed to load image: {e}")
+
+# ------------------ Technical Tab ------------------
 with tabs[1]:
-    ASSISTANT_ID = "asst_DjvuWBc7tCvMbAhY7n1em4BZ"
+    ASSISTANT_ID_TECHNICAL = "asst_DjvuWBc7tCvMbAhY7n1em4BZ"
+
     if "tech_messages" not in st.session_state:
         st.session_state.tech_messages = []
     if "tech_thread_id" not in st.session_state:
         st.session_state.tech_thread_id = None
-    if "results" not in st.session_state:
-        st.session_state.results = []
-    if "lightbox_url" not in st.session_state:
-        st.session_state.lightbox_url = None
+    if "tech_results" not in st.session_state:
+        st.session_state.tech_results = []
+    if "tech_lightbox_url" not in st.session_state:
+        st.session_state.tech_lightbox_url = None
 
-    prompt = st.chat_input("Ask about plans, drawings or components")
-    if prompt:
-        st.session_state.tech_messages.append({"role": "user", "content": prompt})
+    user_prompt = st.chat_input("Ask about plans, drawings or components")
+    if user_prompt:
+        st.session_state.tech_messages.append({"role": "user", "content": user_prompt})
 
     if st.session_state.tech_messages and st.session_state.tech_messages[-1]["role"] == "user":
         try:
@@ -168,51 +146,54 @@ with tabs[1]:
                 role="user",
                 content=st.session_state.tech_messages[-1]["content"]
             )
+
             run = client.beta.threads.runs.create(
                 thread_id=st.session_state.tech_thread_id,
-                assistant_id=ASSISTANT_ID
+                assistant_id=ASSISTANT_ID_TECHNICAL
             )
 
-            with st.spinner("🤖 Querying technical documents..."):
+            with st.spinner("🤖 Parsing and responding..."):
                 while True:
-                    status = client.beta.threads.runs.retrieve(
+                    run_status = client.beta.threads.runs.retrieve(
                         thread_id=st.session_state.tech_thread_id,
                         run_id=run.id
                     )
-                    if status.status in ("completed", "failed", "cancelled"):
+                    if run_status.status in ("completed", "failed", "cancelled"):
                         break
                     time.sleep(1)
 
-            if status.status == "completed":
+            if run_status.status == "completed":
                 messages = client.beta.threads.messages.list(thread_id=st.session_state.tech_thread_id)
-                for m in reversed(messages.data):
-                    if m.role == "assistant":
-                        content = m.content[0].text.value
-                        st.session_state.tech_messages.append({"role": "assistant", "content": content})
+                for message in reversed(messages.data):
+                    if message.role == "assistant":
+                        assistant_reply = message.content[0].text.value
+                        st.session_state.tech_messages.append({"role": "assistant", "content": assistant_reply})
                         try:
-                            json_data = json.loads(content.strip("`json "))
-                            st.session_state.results = json_data
+                            json_data = json.loads(assistant_reply.strip("`json "))
+                            st.session_state.tech_results = json_data
                         except:
-                            st.session_state.results = []
+                            st.session_state.tech_results = []
                         break
             else:
-                st.error("⚠️ Assistant failed to complete.")
+                st.error(f"⚠️ Assistant failed: {run_status.status}")
             st.rerun()
+
         except Exception as e:
             st.error(f"❌ Error: {e}")
 
-    if st.session_state.results:
-        disciplines = sorted(set(d.get("discipline", "") for d in st.session_state.results))
-        selected = st.selectbox("🌍 Filter by discipline", ["All"] + disciplines)
+    if st.session_state.tech_results:
+        disciplines = sorted(set(d.get("discipline", "") for d in st.session_state.tech_results))
+        selected_discipline = st.selectbox("🌍 Filter by discipline", ["All"] + disciplines)
         page_size = 8
         page_num = st.number_input("Page", min_value=1, step=1, value=1)
 
-        filtered = [r for r in st.session_state.results if selected == "All" or r.get("discipline") == selected]
-        paged = filtered[(page_num - 1) * page_size : page_num * page_size]
+        filtered_results = [r for r in st.session_state.tech_results if selected_discipline == "All" or r.get("discipline") == selected_discipline]
+        paged = filtered_results[(page_num - 1) * page_size : page_num * page_size]
 
         st.markdown("---")
         st.subheader("📂 Drawing Results")
         cols = st.columns(4)
+
         for i, item in enumerate(paged):
             with cols[i % 4]:
                 st.markdown(f"**{item['drawing_number']}**")
@@ -220,15 +201,16 @@ with tabs[1]:
                 st.caption(item.get("summary", ""))
                 for url in item.get("images", [])[:1]:
                     if st.button("🖼️ View Image", key=f"view_{i}"):
-                        st.session_state.lightbox_url = url
-
-        if st.session_state.lightbox_url:
-            st.markdown("---")
-            st.image(st.session_state.lightbox_url, use_column_width=True, caption="🔍 Enlarged Preview")
-            if st.button("❌ Close Viewer"):
-                st.session_state.lightbox_url = None
-                st.rerun()
+                        st.session_state.tech_lightbox_url = url
+
+        if st.session_state.tech_lightbox_url:
+            col_a, col_b = st.columns([1, 2])
+            with col_b:
+                st.image(st.session_state.tech_lightbox_url, caption="🔍 Enlarged Preview", use_container_width=True)
+                if st.button("❌ Close Viewer"):
+                    st.session_state.tech_lightbox_url = None
+                    st.rerun()
     else:
         for msg in st.session_state.tech_messages:
             with st.chat_message(msg["role"]):
-                st.markdown(msg["content"], unsafe_allow_html=True)
+                st.markdown(msg["content"], unsafe_allow_html=True)