Update app.py

app.py

@@ -1,47 +1,18 @@
 import gradio as gr
+import os
+import uuid
 import asyncio
-from websockets import connect, Data, ClientConnection
-from dotenv import load_dotenv
-import json
-import os 
 import threading
-import numpy as np
+import json
 import base64
-import soundfile as sf
+import numpy as np
 import io
+import soundfile as sf
 from pydub import AudioSegment
-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 Class
-# =========================
+from websockets import connect, Data, ClientConnection
+from dotenv import load_dotenv
 
+# ========== WebSocket Client Setup ==========
 class WebSocketClient:
     def __init__(self, uri: str, headers: dict, client_id: str):
         self.uri = uri
@@ -53,17 +24,10 @@ class WebSocketClient:
         self.transcript = ""
 
     async def connect(self):
-        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}")
+        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):
         self.loop = asyncio.new_event_loop()
@@ -71,109 +35,103 @@ class WebSocketClient:
         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}")
+        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 += ' '
 
     async def send_audio_chunks(self):
         while True:
-            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}")
+            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}))
 
     async def receive_messages(self):
         async for message in self.websocket:
             self.process_websocket_message(message)
 
-    def enqueue_audio_chunk(self, sample_rate: int, chunk_array: np.ndarray):
+    def enqueue_audio_chunk(self, sr, chunk):
         if not self.queue.full():
-            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}")
+            asyncio.run_coroutine_threadsafe(self.queue.put((sr, chunk)), self.loop)
 
     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."
+# ========== 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():
+    cid = str(uuid.uuid4())
+    connections[cid] = WebSocketClient(URI, HEADERS, cid)
+    threading.Thread(target=connections[cid].run, daemon=True).start()
+    return cid
+
+def send_audio(new_chunk, cid):
+    if cid not in connections:
+        return "Starting connection..."
     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 = ""
+    connections[cid].enqueue_audio_chunk(sr, y)
+    return connections[cid].transcript
+
+def clear_transcript(cid):
+    if cid in connections:
+        connections[cid].transcript = ""
     return ""
 
-# =========================
-# 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()
+# ========== Gradio UI Layout ==========
+with gr.Blocks(theme=gr.themes.Soft()) as app:
+
+    gr.Markdown("# 🧠 Document AI Assistant with Voice & Viewer")
+
+    # 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
+
+            send_btn.click(chat_response, inputs=[msg, chatbot], outputs=[msg, chatbot])
+
+        # Smaller widgets section
+        with gr.Column(scale=1):
+            # 🟡 Image Viewer
+            viewer = gr.Image(label="📄 Document Viewer", type="filepath")
+
+            # 🔵 Voice Transcription
+            transcript = gr.Textbox(label="🎤 Transcript", lines=5, interactive=False)
+            audio = gr.Audio(label="🎙️ Audio", streaming=True)
+            clear = gr.Button("Clear Transcript")
+
+            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)
+
+app.launch()