Update app.py

app.py

@@ -1,97 +1,92 @@
 import gradio as gr
+import openai
 import os
 import json
-from openai import OpenAI
 
-# ========== Config ==========
-OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY")
+# === Load Secrets ===
+openai.api_key = os.environ.get("OPENAI_API_KEY")
 ASSISTANT_ID = os.environ.get("ASSISTANT_ID")
-client = OpenAI(api_key=OPENAI_API_KEY)
+VECTOR_STORE_ID = "vs_68199dddd40c8191973e8a6d1b136cd3"
 
-# ========== Functions ==========
-def query_assistant(message):
+# === Chat Function ===
+def search_drawings(user_query):
     try:
-        thread = client.beta.threads.create()
-        client.beta.threads.messages.create(
+        thread = openai.beta.threads.create()
+        openai.beta.threads.messages.create(
             thread_id=thread.id,
             role="user",
-            content=message
+            content=user_query
         )
-        run = client.beta.threads.runs.create(
+
+        run = openai.beta.threads.runs.create(
             thread_id=thread.id,
-            assistant_id=ASSISTANT_ID
+            assistant_id=ASSISTANT_ID,
+            tool_choice="auto",
+            vector_store_ids=[VECTOR_STORE_ID]
         )
 
-        # Poll run status
+        # Wait for completion
         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"):
+            run_status = openai.beta.threads.runs.retrieve(thread_id=thread.id, run_id=run.id)
+            if run_status.status in ["completed", "failed", "cancelled"]:
                 break
 
         if run_status.status != "completed":
-            return f"❌ Assistant failed: {run_status.status}", []
+            return gr.update(visible=True), f"❌ Assistant failed: {run_status.status}", []
 
-        messages = client.beta.threads.messages.list(thread_id=thread.id)
+        messages = openai.beta.threads.messages.list(thread_id=thread.id)
         for m in reversed(messages.data):
             if m.role == "assistant":
+                raw = m.content[0].text.value
                 try:
-                    parsed = json.loads(m.content[0].text.value)
-                    return None, parsed.get("results", [])
+                    result = json.loads(raw)
+                    cards = []
+                    for entry in result.get("drawings", []):
+                        img_tag = f'<img src="{entry["image_url"]}" style="width:100%; border-radius:12px"/>'
+                        page_caption = f"{entry['drawing_number']} – Page {entry['page_number']}"
+                        summary = f"<b>Summary:</b> {entry['summary']}"
+
+                        with gr.Column(scale=1):
+                            cards.append(gr.HTML(f"""
+                            <div style='background:#1a1a1a; border:1px solid #333; border-radius:14px; padding:16px; height:100%'>
+                                {img_tag}<br>
+                                <p style='margin-top:8px; font-weight:bold;'>{page_caption}</p>
+                                <p style='font-size:14px'>{summary}</p>
+                            </div>
+                            """))
+                    return gr.update(visible=False), "", cards
                 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}", []
+                    return gr.update(visible=True), "⚠️ Could not parse assistant response as JSON.", []
+        return gr.update(visible=True), "⚠️ No assistant response found.", []
 
-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))
+    except Exception as e:
+        return gr.update(visible=True), f"❌ Error: {str(e)}", []
 
-        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
+# === UI ===
+title = "\U0001F5A9 Forrestdale Technical Drawing Assistant"
+description = "Ask for plans by discipline, components or tags (e.g. \"Show all architectural plans\")"
 
-def on_query_submit(prompt):
-    status, results = query_assistant(prompt)
-    if status:
-        return status, []
-    return "✅ Found matching drawings", render_cards(results)
+def ui():
+    with gr.Blocks(theme=gr.themes.Monochrome()) as demo:
+        gr.Markdown(f"""
+        <h1 style='font-size:32px'>{title}</h1>
+        <p style='color:#aaa; font-size:16px'>{description}</p>
+        """)
 
-# ========== 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_)
-    """)
+        with gr.Row():
+            query_input = gr.Textbox(placeholder="e.g. Show all electrical plans", scale=4)
+            submit_btn = gr.Button("Search", variant="primary", scale=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)
+        error_box = gr.Markdown("", visible=False)
+        result_gallery = gr.Group([])
 
-    status_output = gr.Markdown("", visible=True)
-    result_display = gr.Column()
+        submit_btn.click(
+            fn=search_drawings,
+            inputs=[query_input],
+            outputs=[error_box, error_box, result_gallery]
+        )
 
-    query_button.click(fn=on_query_submit, inputs=query_input, outputs=[status_output, result_display])
+    return demo
 
-# ========== Launch ==========
-if __name__ == "__main__":
-    app.launch()
+app = ui()
+app.launch()