Update app.py

app.py

@@ -1,95 +1,97 @@
-import streamlit as st
+import gradio as gr
 import os
-import time
 import json
-import re
 from openai import OpenAI
 
-# Basic config
-st.set_page_config(page_title="Forrestdale Drawing Viewer", layout="wide")
-st.title("📐 Forrestdale Technical Drawing Assistant")
-
+# ========== Config ==========
 OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY")
-ASSISTANT_ID = "asst_DjvuWBc7tCvMbAhY7n1em4BZ"
-
-if not OPENAI_API_KEY:
-    st.error("❌ Missing OPENAI_API_KEY.")
-    st.stop()
-
+ASSISTANT_ID = os.environ.get("ASSISTANT_ID")
 client = OpenAI(api_key=OPENAI_API_KEY)
 
-if "tech_thread_id" not in st.session_state:
-    thread = client.beta.threads.create()
-    st.session_state.tech_thread_id = thread.id
+# ========== 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
+        )
+
+        # 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
 
-if "tech_messages" not in st.session_state:
-    st.session_state.tech_messages = []
+        if run_status.status != "completed":
+            return f"❌ Assistant failed: {run_status.status}", []
 
-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})
+        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}", []
 
-for msg in st.session_state.tech_messages:
-    with st.chat_message(msg["role"]):
-        st.markdown(msg["content"])
+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))
 
-# 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"]
+        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)
+                    ]
+                )
             )
+    return cards
 
-            run = client.beta.threads.runs.create(
-                thread_id=st.session_state.tech_thread_id,
-                assistant_id=ASSISTANT_ID
-            )
+def on_query_submit(prompt):
+    status, results = query_assistant(prompt)
+    if status:
+        return status, []
+    return "✅ Found matching drawings", render_cards(results)
 
-            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)
+# ========== 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_)
+    """)
 
-            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})
+    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)
 
-                        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)
+    status_output = gr.Markdown("", visible=True)
+    result_display = gr.Column()
 
-                            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"))
+    query_button.click(fn=on_query_submit, inputs=query_input, outputs=[status_output, result_display])
 
-                                            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}")
+# ========== Launch ==========
+if __name__ == "__main__":
+    app.launch()