Spaces:
Sleeping
Sleeping
| 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"] | |
| # 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 | |
| if len(connections[cid].transcript) > 1000: | |
| connections[cid].transcript = "" | |
| connections[cid].enqueue_audio_chunk(sr, arr) | |
| return connections[cid].transcript.strip() | |
| def clear_transcript(cid): | |
| if cid in connections: | |
| connections[cid].transcript = "" | |
| return "" | |
| 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({"role": "user", "content": user_input}) | |
| history.append({"role": "assistant", "content": 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 | |
| def send_transcript_to_assistant(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) | |
| def clear_chat_and_transcript(client_id): | |
| if client_id in connections: | |
| connections[client_id].transcript = "" | |
| return [], "", None, None | |
| # UI | |
| with gr.Blocks(theme=gr.themes.Soft()) as app: | |
| gr.Markdown("# 📄 Document AI Assistant") | |
| gr.HTML(""" | |
| <style> | |
| #ask-btn, #clear-chat-btn, #record-audio button { | |
| font-size: 16px !important; | |
| padding: 12px 28px !important; | |
| border-radius: 6px; | |
| margin-top: 10px; | |
| background-color: #f2f2f2 !important; | |
| color: #000 !important; | |
| white-space: nowrap; | |
| overflow: hidden; | |
| text-overflow: ellipsis; | |
| } | |
| button { | |
| margin-right: 8px; | |
| } | |
| /* Hide icon (optional) */ | |
| #record-audio button svg { | |
| margin-right: 6px; | |
| } | |
| /* Hide internal label if redundant */ | |
| #record-audio label { | |
| display: none; | |
| } | |
| </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=2): | |
| chat = gr.Chatbot(label="💬 Chat", height=460, type="messages") | |
| with gr.Row(): | |
| user_prompt = gr.Textbox(placeholder="Ask your question...", show_label=False, scale=6) | |
| send_btn = gr.Button("Send", variant="primary", scale=2) | |
| with gr.Accordion("🎤 Voice Transcription", open=False) as voice_section: | |
| gr.Markdown("**🎙️ Tap below to record your voice**") | |
| voice_input = gr.Audio(label="", streaming=True, elem_id="record-audio") | |
| voice_transcript = gr.Textbox(label="Transcript", lines=2, interactive=False) | |
| with gr.Row(): | |
| ask_btn = gr.Button("🟢 Ask", elem_id="ask-btn") | |
| clear_chat_btn = gr.Button("🧹 Clear Chat", elem_id="clear-chat-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]) | |
| 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) | |
| ask_btn.click(fn=send_transcript_to_assistant, | |
| inputs=[voice_transcript, chat_state, thread_state, image_state, client_id], | |
| outputs=[user_prompt, chat, thread_state, image_state]) | |
| clear_chat_btn.click(fn=clear_chat_and_transcript, | |
| inputs=[client_id], | |
| outputs=[chat, voice_transcript, thread_state, image_state]) | |
| app.load(fn=create_ws, outputs=[client_id]) | |
| app.launch() | |