Update app.py

app.py

@@ -1,157 +1,179 @@
 import gradio as gr
-import os, time, re, json, base64, asyncio, threading, uuid, io
+import asyncio
+from websockets import connect, Data, ClientConnection
+from dotenv import load_dotenv
+import json
+import os 
+import threading
 import numpy as np
+import base64
 import soundfile as sf
+import io
 from pydub import AudioSegment
-from openai import OpenAI
-from websockets import connect, Data, ClientConnection
-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"
+import time
+import uuid
+
+# =========================
+# Setup & Configuration
+# =========================
+
+class LogColors:
+    OK = '\033[94m'
+    SUCCESS = '\033[92m'
+    WARNING = '\033[93m'
+    ERROR = '\033[91m'
+    ENDC = '\033[0m'
+
+load_dotenv() 
+OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY")
+if not OPENAI_API_KEY:
+    raise ValueError("OPENAI_API_KEY environment variable must be set")
+
+WEBSOCKET_URI = "wss://api.openai.com/v1/realtime?intent=transcription"
+WEBSOCKET_HEADERS = {
+    "Authorization": "Bearer " + OPENAI_API_KEY,
+    "OpenAI-Beta": "realtime=v1"
+}
+
+css = ""
 connections = {}
 
-# ============ WebSocket Client ============
+# =========================
+# WebSocket Client Class
+# =========================
+
 class WebSocketClient:
-    def __init__(self, uri, headers, client_id):
-        self.uri, self.headers, self.client_id = uri, headers, client_id
-        self.websocket = None
+    def __init__(self, uri: str, headers: dict, client_id: str):
+        self.uri = uri
+        self.headers = headers
+        self.websocket: ClientConnection = None
         self.queue = asyncio.Queue(maxsize=10)
+        self.loop = None
+        self.client_id = client_id
         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())
+        try:
+            self.websocket = await connect(self.uri, additional_headers=self.headers)
+            print(f"{LogColors.SUCCESS}Connected to OpenAI WebSocket{LogColors.ENDC}\n")
+
+            with open("openai_transcription_settings.json", "r") as f:
+                settings = f.read()
+                await self.websocket.send(settings)
+
+            await asyncio.gather(self.receive_messages(), self.send_audio_chunks())
+        except Exception as e:
+            print(f"{LogColors.ERROR}WebSocket Connection Error: {e}{LogColors.ENDC}")
 
     def run(self):
-        loop = asyncio.new_event_loop()
-        asyncio.set_event_loop(loop)
-        loop.run_until_complete(self.connect())
+        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):
+        message_object = json.loads(message)
+        if message_object["type"] != "error":
+            print(f"{LogColors.OK}Received message: {LogColors.ENDC} {message}")
+            if message_object["type"] == "conversation.item.input_audio_transcription.delta":
+                delta = message_object["delta"]
+                self.transcript += delta
+            elif message_object["type"] == "conversation.item.input_audio_transcription.completed":
+                self.transcript += ' ' if self.transcript and self.transcript[-1] != ' ' else ''
+        else:
+            print(f"{LogColors.ERROR}Error: {message}{LogColors.ENDC}")
 
     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()
-            }))
+            audio_data = await self.queue.get()
+            sample_rate, audio_array = audio_data
+            if self.websocket:
+                if audio_array.ndim > 1:
+                    audio_array = audio_array.mean(axis=1)
+                audio_array = audio_array.astype(np.float32)
+                audio_array /= np.max(np.abs(audio_array)) if np.max(np.abs(audio_array)) > 0 else 1.0
+                audio_array_int16 = (audio_array * 32767).astype(np.int16)
+
+                audio_buffer = io.BytesIO()
+                sf.write(audio_buffer, audio_array_int16, sample_rate, format='WAV', subtype='PCM_16')
+                audio_buffer.seek(0) 
+                audio_segment = AudioSegment.from_file(audio_buffer, format="wav")
+                resampled_audio = audio_segment.set_frame_rate(24000)
+
+                output_buffer = io.BytesIO()
+                resampled_audio.export(output_buffer, format="wav")
+                output_buffer.seek(0)
+                base64_audio = base64.b64encode(output_buffer.read()).decode("utf-8")
+
+                await self.websocket.send(json.dumps({"type": "input_audio_buffer.append", "audio": base64_audio}))
+                print(f"{LogColors.OK}Sent audio chunk{LogColors.ENDC}")
 
     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"]
+        async for message in self.websocket:
+            self.process_websocket_message(message)
 
-    def enqueue_audio_chunk(self, sr, arr):
+    def enqueue_audio_chunk(self, sample_rate: int, chunk_array: np.ndarray):
         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 = ""
+            asyncio.run_coroutine_threadsafe(self.queue.put((sample_rate, chunk_array)), self.loop)
+        else:
+            print(f"{LogColors.WARNING}Queue is full, dropping audio chunk{LogColors.ENDC}")
+
+    async def close(self):
+        if self.websocket:
+            await self.websocket.close()
+            connections.pop(self.client_id)
+        print(f"{LogColors.WARNING}WebSocket connection closed{LogColors.ENDC}")
+
+
+# =========================
+# Helper Functions
+# =========================
+
+def send_audio_chunk(new_chunk: gr.Audio, client_id: str):
+    if client_id not in connections:
+        return "Connection is being established, please try again in a few seconds."
+    sr, y = new_chunk
+    connections[client_id].enqueue_audio_chunk(sr, y)
+    return connections[client_id].transcript
+
+def create_new_websocket_connection():
+    client_id = str(uuid.uuid4())
+    connections[client_id] = WebSocketClient(WEBSOCKET_URI, WEBSOCKET_HEADERS, client_id)
+    threading.Thread(target=connections[client_id].run, daemon=True).start()
+    return client_id
+
+def clear_transcript(client_id):
+    if client_id in connections:
+        connections[client_id].transcript = ""
     return ""
 
-# ============ Chat Assistant ============
-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
-
-# ============ Gradio UI ============
-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()
-    voice_enabled = gr.State(False)
-
-    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=1.4):
-            chat = gr.Chatbot(label="💬 Chat", height=460)
-
-            with gr.Row():
-                user_prompt = gr.Textbox(placeholder="Ask your question...", show_label=False, scale=6)
-                mic_toggle_btn = gr.Button("🎙️", scale=1)
-                send_btn = gr.Button("Send", variant="primary", scale=2)
-
-            with gr.Accordion("🎤 Voice Transcription", open=False) as voice_section:
-                with gr.Row():
-                    voice_input = gr.Audio(label="Mic", streaming=True)
-                    voice_transcript = gr.Textbox(label="Transcript", lines=2, interactive=False)
-                clear_btn = gr.Button("🧹 Clear Transcript")
-
-    # FUNCTIONAL CONNECTIONS
-    def toggle_voice(curr):
-        return not curr, gr.update(visible=not curr)
-
-    mic_toggle_btn.click(fn=toggle_voice, inputs=voice_enabled, outputs=[voice_enabled, voice_section])
-    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)
-    clear_btn.click(fn=clear_transcript, inputs=[client_id], outputs=voice_transcript)
-    app.load(fn=create_ws, outputs=[client_id])
-
-app.launch()
+# =========================
+# Gradio UI Sections
+# =========================
+
+with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
+
+    with gr.Tab("💬 Chat Assistant"):
+        gr.Markdown("### Chat Section (Coming Soon)")
+        gr.Textbox(label="Your question")
+        gr.Button("Send")
+
+    with gr.Tab("📄 Document Viewer"):
+        gr.Markdown("### Upload and View Documents")
+        gr.File(label="Upload Document", file_types=[".pdf", ".txt", ".docx"])
+        gr.Textbox(label="Document Preview", lines=10)
+
+    with gr.Tab("🎤 Voice Transcription"):
+        gr.Markdown("### Realtime Voice-to-Text Transcription")
+        with gr.Row():
+            output_textbox = gr.Textbox(label="Transcript", lines=7, interactive=False, autoscroll=True)
+        with gr.Row():
+            with gr.Column(scale=5):
+                audio_input = gr.Audio(streaming=True, format="wav")
+            with gr.Column():
+                clear_button = gr.Button("Clear Transcript")
+        client_id = gr.State()
+        clear_button.click(clear_transcript, inputs=[client_id], outputs=[output_textbox])
+        audio_input.stream(send_audio_chunk, [audio_input, client_id], [output_textbox], stream_every=0.5)
+        demo.load(create_new_websocket_connection, outputs=[client_id])
+
+demo.launch()