Update app.py

app.py

@@ -1,97 +1,95 @@
-import gradio as gr
+import streamlit as st
 import os
+import time
 import json
+import re
 from openai import OpenAI
 
-# ========== Config ==========
+# Basic config
+st.set_page_config(page_title="Forrestdale Drawing Viewer", layout="wide")
+st.title("📐 Forrestdale Technical Drawing Assistant")
+
 OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY")
-ASSISTANT_ID = os.environ.get("ASSISTANT_ID")
-client = OpenAI(api_key=OPENAI_API_KEY)
+ASSISTANT_ID = "asst_DjvuWBc7tCvMbAhY7n1em4BZ"
 
-# ========== Functions ==========
-def query_assistant(message):
-    try:
-        thread = client.beta.threads.create()
-        client.beta.threads.messages.create(
-            thread_id=thread.id,
-            role="user",
-            content=message
-        )
-        run = client.beta.threads.runs.create(
-            thread_id=thread.id,
-            assistant_id=ASSISTANT_ID
-        )
+if not OPENAI_API_KEY:
+    st.error("❌ Missing OPENAI_API_KEY.")
+    st.stop()
 
-        # Poll run status
-        while True:
-            run_status = client.beta.threads.runs.retrieve(
-                thread_id=thread.id,
-                run_id=run.id
-            )
-            if run_status.status in ("completed", "failed", "cancelled"):
-                break
+client = OpenAI(api_key=OPENAI_API_KEY)
 
-        if run_status.status != "completed":
-            return f"❌ Assistant failed: {run_status.status}", []
+if "tech_thread_id" not in st.session_state:
+    thread = client.beta.threads.create()
+    st.session_state.tech_thread_id = thread.id
 
-        messages = client.beta.threads.messages.list(thread_id=thread.id)
-        for m in reversed(messages.data):
-            if m.role == "assistant":
-                try:
-                    parsed = json.loads(m.content[0].text.value)
-                    return None, parsed.get("results", [])
-                except Exception as e:
-                    return f"⚠️ Failed to parse assistant response: {e}", []
-        return "⚠️ No assistant response found", []
-    except Exception as e:
-        return f"❌ Error: {e}", []
+if "tech_messages" not in st.session_state:
+    st.session_state.tech_messages = []
 
-def render_cards(results):
-    cards = []
-    for r in results:
-        title = r.get("drawing_title", "Untitled")
-        summary = r.get("summary", "")
-        pages = r.get("pages", [])
-        image_elems = []
-        for page in pages:
-            url = page.get("public_image_url")
-            label = f"{title} – Page {page.get('page_number')}"
-            if url:
-                image_elems.append(gr.Image(value=url, label=label, show_label=True))
+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})
 
-        with gr.Column(scale=1):
-            cards.append(
-                gr.Group(
-                    [
-                        gr.Markdown(f"**{title}**\n\n**Summary:** {summary}"),
-                        gr.Accordion("📂 View Drawing Pages", open=False, children=image_elems)
-                    ]
-                )
+for msg in st.session_state.tech_messages:
+    with st.chat_message(msg["role"]):
+        st.markdown(msg["content"])
+
+# Fetch assistant reply
+if st.session_state.tech_messages and st.session_state.tech_messages[-1]["role"] == "user":
+    with st.spinner("⏳ Fetching results from assistant..."):
+        try:
+            client.beta.threads.messages.create(
+                thread_id=st.session_state.tech_thread_id,
+                role="user",
+                content=st.session_state.tech_messages[-1]["content"]
             )
-    return cards
 
-def on_query_submit(prompt):
-    status, results = query_assistant(prompt)
-    if status:
-        return status, []
-    return "✅ Found matching drawings", render_cards(results)
+            run = client.beta.threads.runs.create(
+                thread_id=st.session_state.tech_thread_id,
+                assistant_id=ASSISTANT_ID
+            )
 
-# ========== Interface ==========
-with gr.Blocks(theme=gr.themes.Base(), title="Forrestdale Technical Drawing Assistant") as app:
-    gr.Markdown("""
-    # 🏗️ Forrestdale Technical Drawing Assistant
-    Ask about plans, drawings or components (e.g. _Show me all electrical plans_)
-    """)
+            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)
 
-    with gr.Row():
-        query_input = gr.Textbox(placeholder="e.g. Show me all civil drainage plans", scale=5)
-        query_button = gr.Button("🔍 Search", scale=1)
+            if run_status.status != "completed":
+                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_content = message.content[0].text.value.strip()
+                        st.session_state.tech_messages.append({"role": "assistant", "content": reply_content})
 
-    status_output = gr.Markdown("", visible=True)
-    result_display = gr.Column()
+                        try:
+                            match = re.search(r"```json\s*(.*?)```", reply_content, re.DOTALL)
+                            json_str = match.group(1) if match else reply_content
+                            results = json.loads(json_str)
 
-    query_button.click(fn=on_query_submit, inputs=query_input, outputs=[status_output, result_display])
+                            if isinstance(results, list):
+                                cols = st.columns(4)
+                                for idx, item in enumerate(results):
+                                    with cols[idx % 4]:
+                                        with st.container(border=True):
+                                            st.markdown(f"### 📁 {item.get('drawing_number')} ({item.get('discipline')})", help=item.get("summary"))
+                                            st.caption(item.get("summary"))
 
-# ========== Launch ==========
-if __name__ == "__main__":
-    app.launch()
+                                            with st.expander("📂 View Drawing Details"):
+                                                if item.get("question"):
+                                                    st.markdown(f"**Question Match:** {item.get('question')}")
+                                                if item.get("image"):
+                                                    st.image(item.get("image"), caption=item.get("drawing_number"))
+                                                elif item.get("images"):
+                                                    for i, img in enumerate(item["images"]):
+                                                        if img.startswith("http"):
+                                                            st.image(img, caption=f"{item['drawing_number']} – Page {i+1}")
+                        except Exception as parse_error:
+                            st.warning("🟡 Could not parse assistant response as JSON.")
+                        break
+        except Exception as e:
+            st.error(f"❌ Error occurred: {e}")