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 environment variables
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 for Voice ----------------
class WebSocketClient:
    def __init__(self, uri, headers, client_id):
        self.uri, self.headers, self.client_id = uri, headers, client_id
        self.websocket = None
        self.queue = asyncio.Queue(maxsize=10)
        self.transcript = ""

    async def connect(self):
        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())

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

    async def send_audio_chunks(self):
        while True:
            sr, arr = await self.queue.get()
            if arr.ndim > 1: arr = arr.mean(axis=1)
            arr = (arr / np.max(np.abs(arr))) if np.max(np.abs(arr)) > 0 else 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 enqueue_audio_chunk(self, sr, arr):
        if not self.queue.full():
            asyncio.run_coroutine_threadsafe(self.queue.put((sr, arr)), asyncio.get_event_loop())

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 cid not in connections: return "Connecting..."
    sr, arr = chunk
    connections[cid].enqueue_audio_chunk(sr, arr)
    return connections[cid].transcript

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

# ---------------- Chat Assistant Logic ----------------
def handle_chat(user_input, history, thread_id, image_url):
    if not OPENAI_API_KEY or not ASSISTANT_ID:
        return "❌ Missing API key or Assistant ID.", 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

# ---------------- UI ----------------
with gr.Blocks(theme="lone17/kotaemon") as app:
    gr.Markdown("# 📄 Document AI Assistant")

    # States
    chat_state = gr.State([])
    thread_state = gr.State()
    image_state = gr.State()
    client_id = gr.State()
    mic_shown = gr.State(False)

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

        # Right: Chat + Mic
        with gr.Column(scale=1.4):
            chat = gr.Chatbot(label="💬 Chat", height=450)

            with gr.Row():
                user_input = gr.Textbox(placeholder="Ask your question...", show_label=False, scale=6)
                mic_btn = gr.Button("🎙️", scale=1)
                send_btn = gr.Button("Send", scale=2)

            # Hidden Voice Section
            with gr.Row(visible=False) as mic_row:
                with gr.Column(scale=4):
                    audio = gr.Audio(label="🎤 Speak", streaming=True)
                with gr.Column(scale=5):
                    transcript = gr.Textbox(label="Transcript", lines=2, interactive=False)
                with gr.Column(scale=2):
                    clear_btn = gr.Button("🧹 Clear")

    # Logic Wiring
    def toggle_mic(state): return not state, gr.update(visible=not state)
    mic_btn.click(toggle_mic, inputs=mic_shown, outputs=[mic_shown, mic_row])
    send_btn.click(handle_chat,
                   inputs=[user_input, chat_state, thread_state, image_state],
                   outputs=[user_input, chat, thread_state, image_state])
    image_state.change(fn=lambda x: x, inputs=image_state, outputs=image_display)
    audio.stream(fn=send_audio, inputs=[audio, client_id], outputs=transcript, stream_every=0.5)
    clear_btn.click(fn=clear_transcript, inputs=[client_id], outputs=transcript)
    app.load(fn=create_ws, outputs=[client_id])

app.launch()