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 = {}

# WebSocket Client
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"🔴 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"]

# Connection manager
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
    if len(connections[cid].transcript) > 1000:
        connections[cid].transcript = ""
    connections[cid].enqueue_audio_chunk(sr, arr)
    return connections[cid].transcript.strip()

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

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({"role": "user", "content": user_input})
                history.append({"role": "assistant", "content": 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 send_transcript_to_assistant(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 clear_chat_and_transcript(client_id):
    if client_id in connections:
        connections[client_id].transcript = ""
    return [], "", None, None

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

    gr.HTML("""
    <style>
    #ask-btn, #clear-chat-btn, #record-audio button {
        font-size: 16px !important;
        padding: 12px 28px !important;
        border-radius: 6px;
        margin-top: 10px;
        background-color: #f2f2f2 !important;
        color: #000 !important;
        white-space: nowrap;
        overflow: hidden;
        text-overflow: ellipsis;
    }

    button {
        margin-right: 8px;
    }

    /* Hide icon (optional) */
    #record-audio button svg {
        margin-right: 6px;
    }

    /* Hide internal label if redundant */
    #record-audio label {
        display: none;
    }
    </style>
    """)

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

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

        with gr.Column(scale=2):
            chat = gr.Chatbot(label="💬 Chat", height=460, type="messages")
            with gr.Row():
                user_prompt = gr.Textbox(placeholder="Ask your question...", show_label=False, scale=6)
                send_btn = gr.Button("Send", variant="primary", scale=2)

            with gr.Accordion("🎤 Voice Transcription", open=False) as voice_section:
                gr.Markdown("**🎙️ Tap below to record your voice**")
                voice_input = gr.Audio(label="", streaming=True, elem_id="record-audio")
                voice_transcript = gr.Textbox(label="Transcript", lines=2, interactive=False)

                with gr.Row():
                    ask_btn = gr.Button("🟢 Ask", elem_id="ask-btn")
                    clear_chat_btn = gr.Button("🧹 Clear Chat", elem_id="clear-chat-btn")

    # Bindings
    send_btn.click(fn=handle_chat,
                   inputs=[user_prompt, chat_state, thread_state, image_state],
                   outputs=[user_prompt, chat, thread_state, image_state])
    image_state.change(fn=lambda x: x, inputs=image_state, outputs=image_display)
    voice_input.stream(fn=send_audio, inputs=[voice_input, client_id], outputs=voice_transcript, stream_every=0.5)

    ask_btn.click(fn=send_transcript_to_assistant,
                  inputs=[voice_transcript, chat_state, thread_state, image_state, client_id],
                  outputs=[user_prompt, chat, thread_state, image_state])

    clear_chat_btn.click(fn=clear_chat_and_transcript,
                         inputs=[client_id],
                         outputs=[chat, voice_transcript, thread_state, image_state])

    app.load(fn=create_ws, outputs=[client_id])

app.launch()