Update app.py

app.py

@@ -1,26 +1,28 @@
 import gradio as gr
-import os
-import uuid
-import asyncio
-import threading
-import json
-import base64
+import os, time, re, json, base64, asyncio, threading, uuid, io
 import numpy as np
-import io
 import soundfile as sf
 from pydub import AudioSegment
+from openai import OpenAI
 from websockets import connect, Data, ClientConnection
 from dotenv import load_dotenv
 
-# ========== WebSocket Client Setup ==========
+# ---------------- Environment & Client Setup ----------------
+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: str, headers: dict, client_id: str):
-        self.uri = uri
-        self.headers = headers
-        self.websocket: ClientConnection = None
+    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.loop = None
-        self.client_id = client_id
         self.transcript = ""
 
     async def connect(self):
@@ -30,108 +32,114 @@ class WebSocketClient:
         await asyncio.gather(self.receive_messages(), self.send_audio_chunks())
 
     def run(self):
-        self.loop = asyncio.new_event_loop()
-        asyncio.set_event_loop(self.loop)
-        self.loop.run_until_complete(self.connect())
-
-    def process_websocket_message(self, message: Data):
-        msg = json.loads(message)
-        if msg["type"] == "conversation.item.input_audio_transcription.delta":
-            self.transcript += msg["delta"]
-        elif msg["type"] == "conversation.item.input_audio_transcription.completed":
-            self.transcript += ' '
+        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, audio_array = await self.queue.get()
-            if audio_array.ndim > 1:
-                audio_array = audio_array.mean(axis=1)
-            audio_array = (audio_array / np.max(np.abs(audio_array))) if np.max(np.abs(audio_array)) > 0 else audio_array
-            int16 = (audio_array * 32767).astype(np.int16)
-            buffer = io.BytesIO()
-            sf.write(buffer, int16, sr, format='WAV', subtype='PCM_16')
-            buffer.seek(0)
-            audio = AudioSegment.from_file(buffer, format="wav").set_frame_rate(24000)
-            out = io.BytesIO()
-            audio.export(out, format="wav")
-            out.seek(0)
-            encoded = base64.b64encode(out.read()).decode("utf-8")
-            await self.websocket.send(json.dumps({"type": "input_audio_buffer.append", "audio": encoded}))
+            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 message in self.websocket:
-            self.process_websocket_message(message)
+        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, chunk):
+    def enqueue_audio_chunk(self, sr, arr):
         if not self.queue.full():
-            asyncio.run_coroutine_threadsafe(self.queue.put((sr, chunk)), self.loop)
-
-    async def close(self):
-        if self.websocket:
-            await self.websocket.close()
+            asyncio.run_coroutine_threadsafe(self.queue.put((sr, arr)), asyncio.get_event_loop())
 
-
-# ========== Transcription Helpers ==========
-connections = {}
-load_dotenv()
-OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY")
-HEADERS = {"Authorization": f"Bearer {OPENAI_API_KEY}", "OpenAI-Beta": "realtime=v1"}
-URI = "wss://api.openai.com/v1/realtime?intent=transcription"
-
-def create_websocket():
+def create_ws():
     cid = str(uuid.uuid4())
-    connections[cid] = WebSocketClient(URI, HEADERS, cid)
-    threading.Thread(target=connections[cid].run, daemon=True).start()
+    client = WebSocketClient(WS_URI, HEADERS, cid)
+    threading.Thread(target=client.run, daemon=True).start()
+    connections[cid] = client
     return cid
 
-def send_audio(new_chunk, cid):
-    if cid not in connections:
-        return "Starting connection..."
-    sr, y = new_chunk
-    connections[cid].enqueue_audio_chunk(sr, y)
+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 = ""
+    if cid in connections: connections[cid].transcript = ""
     return ""
 
-# ========== Gradio UI Layout ==========
-with gr.Blocks(theme=gr.themes.Soft()) as app:
+# ---------------- Chat Functionality ----------------
+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
 
-    gr.Markdown("# 🧠 Document AI Assistant with Voice & Viewer")
+        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)
 
-    # State
+        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
+
+# ---------------- Gradio UI Layout ----------------
+with gr.Blocks(theme=gr.themes.Soft()) as app:
+    gr.Markdown("# 📄 Document AI Assistant")
+
+    # STATES
+    chat_state = gr.State([])
+    thread_state = gr.State()
+    image_state = gr.State()
     client_id = gr.State()
 
-    # Layout
     with gr.Row():
-        # 🟢 Chat Section (Main)
-        with gr.Column(scale=2):
-            chatbot = gr.Chatbot(label="Chat Assistant")
-            msg = gr.Textbox(label="Ask something", placeholder="e.g., Summarize this document...")
-            send_btn = gr.Button("Send")
-            
-            def chat_response(user_msg, history):
-                history = history or []
-                reply = f"🤖 This is a placeholder reply to: {user_msg}"
-                history.append((user_msg, reply))
-                return "", history
+        with gr.Column(scale=1):
+            # IMAGE VIEWER (left)
+            image_display = gr.Image(label="🖼️ Document", type="filepath")
 
-            send_btn.click(chat_response, inputs=[msg, chatbot], outputs=[msg, chatbot])
+            # VOICE (under)
+            voice_transcript = gr.Textbox(label="🎙️ Transcript", lines=4, interactive=False)
+            voice_input = gr.Audio(label="🔴 Record", streaming=True)
+            clear_btn = gr.Button("🧹 Clear Transcript")
 
-        # Smaller widgets section
-        with gr.Column(scale=1):
-            # 🟡 Image Viewer
-            viewer = gr.Image(label="📄 Document Viewer", type="filepath")
+        with gr.Column(scale=2):
+            # CHATBOT (right)
+            chat = gr.Chatbot(label="💬 Chat", height=450)
+            user_prompt = gr.Textbox(show_label=False, placeholder="Ask your question...")
+            send_btn = gr.Button("Send")
 
-            # 🔵 Voice Transcription
-            transcript = gr.Textbox(label="🎤 Transcript", lines=5, interactive=False)
-            audio = gr.Audio(label="🎙️ Audio", streaming=True)
-            clear = gr.Button("Clear Transcript")
+    # HANDLERS
+    send_btn.click(handle_chat,
+                   inputs=[user_prompt, chat_state, thread_state, image_state],
+                   outputs=[user_prompt, chat, thread_state, image_state])
 
-            audio.stream(fn=send_audio, inputs=[audio, client_id], outputs=transcript, stream_every=0.5)
-            clear.click(fn=clear_transcript, inputs=[client_id], outputs=transcript)
-            app.load(create_websocket, outputs=client_id)
+    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)
+    clear_btn.click(fn=clear_transcript, inputs=[client_id], outputs=voice_transcript)
+    app.load(create_ws, outputs=[client_id])
 
 app.launch()