app.py

import gradio as gr
import openai
import os
import json

# === Load Secrets ===
openai.api_key = os.environ.get("OPENAI_API_KEY")
ASSISTANT_ID = os.environ.get("ASSISTANT_ID")
VECTOR_STORE_ID = "vs_68199dddd40c8191973e8a6d1b136cd3"

# === Chat Function ===
def search_drawings(user_query):
    try:
        thread = openai.beta.threads.create()
        openai.beta.threads.messages.create(
            thread_id=thread.id,
            role="user",
            content=user_query
        )

        run = openai.beta.threads.runs.create(
            thread_id=thread.id,
            assistant_id=ASSISTANT_ID,
            tool_choice="auto",
            vector_store_ids=[VECTOR_STORE_ID]
        )

        # Wait for completion
        while True:
            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 gr.update(visible=True), f"❌ Assistant failed: {run_status.status}", []

        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:
                    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 gr.update(visible=True), "⚠️ Could not parse assistant response as JSON.", []
        return gr.update(visible=True), "⚠️ No assistant response found.", []

    except Exception as e:
        return gr.update(visible=True), f"❌ Error: {str(e)}", []

# === UI ===
title = "\U0001F5A9 Forrestdale Technical Drawing Assistant"
description = "Ask for plans by discipline, components or tags (e.g. \"Show all architectural plans\")"

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>
        """)

        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)

        error_box = gr.Markdown("", visible=False)
        result_gallery = gr.Group([])

        submit_btn.click(
            fn=search_drawings,
            inputs=[query_input],
            outputs=[error_box, error_box, result_gallery]
        )

    return demo

app = ui()
app.launch()