app.py

import gradio as gr
import os, time, re, json, base64, asyncio, threading, uuid, io
import numpy as np
import soundfile as sf
from pydub import AudioSegment
from openai import OpenAI
from websockets import connect
from dotenv import load_dotenv

# Load secrets
load_dotenv()
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
ASSISTANT_ID = os.getenv("ASSISTANT_ID")
client = OpenAI(api_key=OPENAI_API_KEY)

HEADERS = {"Authorization": f"Bearer {OPENAI_API_KEY}", "OpenAI-Beta": "realtime=v1"}
WS_URI = "wss://api.openai.com/v1/realtime?intent=transcription"
connections = {}

class WebSocketClient:
    def __init__(self, uri, headers, client_id):
        self.uri = uri
        self.headers = headers
        self.client_id = client_id
        self.websocket = None
        self.queue = asyncio.Queue(maxsize=10)
        self.transcript = ""
        self.loop = asyncio.new_event_loop()

    async def connect(self):
        try:
            self.websocket = await connect(self.uri, additional_headers=self.headers)
            with open("openai_transcription_settings.json", "r") as f:
                await self.websocket.send(f.read())
            await asyncio.gather(self.receive_messages(), self.send_audio_chunks())
        except Exception as e:
            print(f"\U0001F534 WebSocket Connection Failed: {e}")

    def run(self):
        asyncio.set_event_loop(self.loop)
        self.loop.run_until_complete(self.connect())

    def enqueue_audio_chunk(self, sr, arr):
        if not self.queue.full():
            asyncio.run_coroutine_threadsafe(self.queue.put((sr, arr)), self.loop)

    async def send_audio_chunks(self):
        while True:
            sr, arr = await self.queue.get()
            if arr.ndim > 1:
                arr = arr.mean(axis=1)
            if np.max(np.abs(arr)) > 0:
                arr = arr / np.max(np.abs(arr))
            int16 = (arr * 32767).astype(np.int16)
            buf = io.BytesIO()
            sf.write(buf, int16, sr, format='WAV', subtype='PCM_16')
            audio = AudioSegment.from_file(buf, format="wav").set_frame_rate(24000)
            out = io.BytesIO()
            audio.export(out, format="wav")
            out.seek(0)
            await self.websocket.send(json.dumps({
                "type": "input_audio_buffer.append",
                "audio": base64.b64encode(out.read()).decode()
            }))

    async def receive_messages(self):
        async for msg in self.websocket:
            data = json.loads(msg)
            if data["type"] == "conversation.item.input_audio_transcription.delta":
                self.transcript += data["delta"]

def create_ws():
    cid = str(uuid.uuid4())
    client = WebSocketClient(WS_URI, HEADERS, cid)
    threading.Thread(target=client.run, daemon=True).start()
    connections[cid] = client
    return cid

def send_audio(chunk, cid):
    if not cid or cid not in connections:
        return "Connecting..."
    sr, arr = chunk
    connections[cid].enqueue_audio_chunk(sr, arr)
    return connections[cid].transcript.strip()

def clear_transcript_only(cid):
    if cid in connections:
        connections[cid].transcript = ""
    return ""

def clear_chat_only():
    return [], None, None

def handle_chat(user_input, history, thread_id, image_url):
    if not OPENAI_API_KEY or not ASSISTANT_ID:
        return "❌ Missing secrets!", history, thread_id, image_url
    try:
        if thread_id is None:
            thread = client.beta.threads.create()
            thread_id = thread.id
        client.beta.threads.messages.create(thread_id=thread_id, role="user", content=user_input)
        run = client.beta.threads.runs.create(thread_id=thread_id, assistant_id=ASSISTANT_ID)
        while True:
            status = client.beta.threads.runs.retrieve(thread_id=thread_id, run_id=run.id)
            if status.status == "completed":
                break
            time.sleep(1)
        msgs = client.beta.threads.messages.list(thread_id=thread_id)
        for msg in reversed(msgs.data):
            if msg.role == "assistant":
                content = msg.content[0].text.value
                history.append((user_input, content))
                match = re.search(
                    r'https://raw\.githubusercontent\.com/AndrewLORTech/surgical-pathology-manual/main/[\w\-/]*\.png',
                    content)
                if match:
                    image_url = match.group(0)
                break
        return "", history, thread_id, image_url
    except Exception as e:
        return f"❌ {e}", history, thread_id, image_url

def feed_transcript(transcript, history, thread_id, image_url, cid):
    if not transcript.strip():
        return gr.update(), history, thread_id, image_url
    if cid in connections:
        connections[cid].transcript = ""
    return handle_chat(transcript, history, thread_id, image_url)

def update_image_display(image_url):
    if image_url and isinstance(image_url, str) and image_url.startswith("http"):
        return image_url
    return None

with gr.Blocks(theme=gr.themes.Soft()) as app:
    gr.Markdown("# 📄 Document AI Assistant")

    chat_state = gr.State([])
    thread_state = gr.State()
    image_state = gr.State()
    client_id = gr.State()

    # Hidden toggles as backend state
    show_left = gr.Checkbox(value=True, visible=False)
    show_right = gr.Checkbox(value=True, visible=False)

    with gr.Row():
        toggle_left_btn = gr.Button("📄 Toggle Document Panel")
        toggle_right_btn = gr.Button("🎙️ Toggle Voice Panel")

    with gr.Row(equal_height=True) as layout:
        left_col = gr.Column(visible=True, scale=1)
        with left_col:
            image_display = gr.Image(label="🖼️ Document", type="filepath", show_download_button=False, height=600)

        center_col = gr.Column(scale=2)
        with center_col:
            chat = gr.Chatbot(label="💬 Chat", height=600)
            with gr.Row():
                user_prompt = gr.Textbox(placeholder="Ask your question...", show_label=False, scale=8)
                send_btn = gr.Button("Send", variant="primary", scale=2)
            with gr.Row():
                clear_chat_btn = gr.Button("🗑️ Clear Chat")

        right_col = gr.Column(visible=True, scale=1)
        with right_col:
            gr.Markdown("### 🎙️ Voice Input")
            voice_input = gr.Audio(label="Tap to Record", streaming=True, type="numpy", show_label=True)
            voice_transcript = gr.Textbox(label="Transcript", lines=2, interactive=False)
            with gr.Row():
                voice_send_btn = gr.Button("🟢 Send Voice to Assistant")
                clear_transcript_btn = gr.Button("🧹 Clear Transcript")

    def update_layout(show_left_val, show_right_val):
        center_scale = 2
        if not show_left_val and not show_right_val:
            center_scale = 4
        elif not show_left_val or not show_right_val:
            center_scale = 3
        return (
            gr.update(visible=show_left_val),
            gr.update(scale=center_scale),
            gr.update(visible=show_right_val)
        )

    def toggle_left_fn(current):
        return not current

    def toggle_right_fn(current):
        return not current

    toggle_left_btn.click(fn=toggle_left_fn, inputs=show_left, outputs=show_left)
    toggle_right_btn.click(fn=toggle_right_fn, inputs=show_right, outputs=show_right)

    show_left.change(fn=update_layout, inputs=[show_left, show_right], outputs=[left_col, center_col, right_col])
    show_right.change(fn=update_layout, inputs=[show_left, show_right], outputs=[left_col, center_col, right_col])

    send_btn.click(fn=handle_chat, inputs=[user_prompt, chat_state, thread_state, image_state],
                   outputs=[user_prompt, chat, thread_state, image_state])
    voice_input.stream(fn=send_audio, inputs=[voice_input, client_id], outputs=voice_transcript, stream_every=0.5)
    voice_send_btn.click(fn=feed_transcript,
                         inputs=[voice_transcript, chat_state, thread_state, image_state, client_id],
                         outputs=[user_prompt, chat, thread_state, image_state])
    clear_transcript_btn.click(fn=clear_transcript_only, inputs=[client_id], outputs=voice_transcript)
    clear_chat_btn.click(fn=clear_chat_only, outputs=[chat, thread_state, image_state])
    image_state.change(fn=update_image_display, inputs=image_state, outputs=image_display)
    app.load(fn=create_ws, outputs=[client_id])

app.launch()