Spaces:
Sleeping
Sleeping
| import streamlit as st | |
| import os | |
| import time | |
| import re | |
| import requests | |
| import json | |
| from PIL import Image | |
| from io import BytesIO | |
| 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.title("📄 Schlaeger Forrestdale Document Assistant") | |
| 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 environment variables.") | |
| st.stop() | |
| client = OpenAI(api_key=OPENAI_API_KEY) | |
| # ------------------ Tabs ------------------ | |
| tab1, tab2 = st.tabs(["📑 Contract", "📁 Technical"]) | |
| # ------------------ Contract Assistant ------------------ | |
| with tab1: | |
| ASSISTANT_ID = "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 "contract_image_url" not in st.session_state: | |
| st.session_state.contract_image_url = None | |
| st.markdown("### 🧠 Ask a Contract-Specific Question") | |
| contract_prompt = st.chat_input("Example: What is the defects liability period?", key="contract_input") | |
| if contract_prompt: | |
| st.session_state.contract_messages.append({"role": "user", "content": contract_prompt}) | |
| if st.session_state.contract_messages and st.session_state.contract_messages[-1]["role"] == "user": | |
| try: | |
| if st.session_state.contract_thread_id is None: | |
| thread = client.beta.threads.create() | |
| st.session_state.contract_thread_id = thread.id | |
| client.beta.threads.messages.create( | |
| thread_id=st.session_state.contract_thread_id, | |
| 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 | |
| ) | |
| with st.spinner("🤖 Parsing contract documents..."): | |
| while True: | |
| run_status = client.beta.threads.runs.retrieve( | |
| thread_id=st.session_state.contract_thread_id, | |
| run_id=run.id | |
| ) | |
| if run_status.status in ("completed", "failed", "cancelled"): | |
| break | |
| time.sleep(1) | |
| if run_status.status == "completed": | |
| messages = client.beta.threads.messages.list(thread_id=st.session_state.contract_thread_id) | |
| for message in reversed(messages.data): | |
| if message.role == "assistant": | |
| reply = message.content[0].text.value | |
| st.session_state.contract_messages.append({"role": "assistant", "content": reply}) | |
| match = re.search(r'Document Reference:\s*(.*?),\s*Page\s*(\\d+)', reply) | |
| if match: | |
| doc_name = match.group(1).strip() | |
| page = int(match.group(2)) | |
| page_str = f"{page:04d}" | |
| folder = quote(doc_name) | |
| url = f"https://raw.githubusercontent.com/AndrewLORTech/c2ozschlaegerforrestdale/main/{folder}/{folder}_page_{page_str}.png" | |
| st.session_state.contract_image_url = url | |
| break | |
| else: | |
| st.error(f"❌ Contract Assistant failed: {run_status.status}") | |
| st.rerun() | |
| except Exception as e: | |
| st.error(f"❌ Error: {e}") | |
| for msg in st.session_state.contract_messages: | |
| with st.chat_message(msg["role"]): | |
| st.markdown(msg["content"], unsafe_allow_html=True) | |
| if st.session_state.contract_image_url: | |
| st.image(st.session_state.contract_image_url, caption="📄 Contract Page Image", use_container_width=True) | |
| # ------------------ Technical Drawing Assistant ------------------ | |
| with tab2: | |
| ASSISTANT_ID = "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 "tech_results" not in st.session_state: | |
| st.session_state.tech_results = [] | |
| tech_prompt = st.chat_input("Ask about plans, drawings or components", key="tech_input") | |
| 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": | |
| 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", | |
| 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 | |
| ) | |
| with st.spinner("🤖 Searching technical drawings..."): | |
| 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": | |
| 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}) | |
| try: | |
| json_data = json.loads(reply.strip("`json ")) | |
| st.session_state.tech_results = json_data | |
| except: | |
| st.session_state.tech_results = [] | |
| break | |
| else: | |
| st.error(f"❌ Technical Assistant failed: {run_status.status}") | |
| st.rerun() | |
| except Exception as e: | |
| st.error(f"❌ Error: {e}") | |
| # ------------------ Filters and Display ------------------ | |
| 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.tech_results if selected_discipline == "All" or r.get("discipline") == selected_discipline] | |
| paged = filtered[(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']}**") | |
| st.markdown(f"_Discipline: {item['discipline']}_") | |
| st.caption(item.get("summary", "")) | |
| for url in item.get("images", [])[:1]: | |
| st.image(url, caption="Thumbnail", use_container_width=True) | |
| else: | |
| for msg in st.session_state.tech_messages: | |
| with st.chat_message(msg["role"]): | |
| st.markdown(msg["content"], unsafe_allow_html=True) | |