Spaces:
Running
Running
| 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"] | |
| # WebSocket 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 | |
| connections[cid].enqueue_audio_chunk(sr, arr) | |
| return connections[cid].transcript.strip() | |
| def clear_transcript_only(cid): | |
| if cid in connections: | |
| connections[cid].transcript = "" | |
| return "" | |
| def clear_chat_only(): | |
| return [], None, None | |
| # Assistant chat handler | |
| 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 | |
| # Feed transcript as assistant input | |
| def feed_transcript(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) | |
| # Fallback for image display | |
| def update_image_display(image_url): | |
| if image_url and isinstance(image_url, str) and image_url.startswith("http"): | |
| return image_url | |
| return None | |
| # ============ Gradio UI ============ | |
| with gr.Blocks(theme=gr.themes.Soft()) as app: | |
| gr.Markdown("# 📄 Document AI Assistant") | |
| gr.HTML(""" | |
| <style> | |
| .big-btn { | |
| font-size: 18px !important; | |
| padding: 14px 28px !important; | |
| border-radius: 8px !important; | |
| width: 100% !important; | |
| margin-top: 10px; | |
| } | |
| .voice-area { | |
| padding-top: 12px; | |
| border-top: 1px solid #444; | |
| margin-top: 12px; | |
| } | |
| </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=1.4): | |
| chat = gr.Chatbot(label="💬 Chat", height=460) | |
| with gr.Row(): | |
| user_prompt = gr.Textbox(placeholder="Ask your question...", show_label=False, scale=8) | |
| send_btn = gr.Button("Send", variant="primary", scale=2) | |
| with gr.Column(elem_classes="voice-area"): | |
| gr.Markdown("### 🎙️ Voice Input") | |
| voice_input = gr.Audio(label="Tap to Record", streaming=True, type="numpy", show_label=True) | |
| voice_transcript = gr.Textbox(label="Transcript", lines=2, interactive=False) | |
| with gr.Row(): | |
| voice_send_btn = gr.Button("🟢 Send Voice to Assistant", elem_classes="big-btn") | |
| clear_transcript_btn = gr.Button("🧹 Clear Transcript", elem_classes="big-btn") | |
| with gr.Row(): | |
| clear_chat_btn = gr.Button("🗑️ Clear Chat", elem_classes="big-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]) | |
| voice_input.stream(fn=send_audio, | |
| inputs=[voice_input, client_id], | |
| outputs=voice_transcript, | |
| stream_every=0.5) | |
| voice_send_btn.click(fn=feed_transcript, | |
| inputs=[voice_transcript, chat_state, thread_state, image_state, client_id], | |
| outputs=[user_prompt, chat, thread_state, image_state]) | |
| clear_transcript_btn.click(fn=clear_transcript_only, | |
| inputs=[client_id], | |
| outputs=voice_transcript) | |
| clear_chat_btn.click(fn=clear_chat_only, | |
| outputs=[chat, thread_state, image_state]) | |
| image_state.change(fn=update_image_display, | |
| inputs=image_state, | |
| outputs=image_display) | |
| app.load(fn=create_ws, outputs=[client_id]) | |
| app.launch() | |