Update app.py

app.py

@@ -1,4 +1,3 @@
-# Final version with Perplexity-style UI, voice toggle, rich formatting, full-width input and working mic
 import gradio as gr
 import os, time, re, json, base64, asyncio, threading, uuid, io
 import numpy as np
@@ -18,6 +17,7 @@ HEADERS = {"Authorization": f"Bearer {OPENAI_API_KEY}", "OpenAI-Beta": "realtime
 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
@@ -30,12 +30,12 @@ class WebSocketClient:
 
     async def connect(self):
         try:
-            self.websocket = await connect(self.uri, extra_headers=self.headers)
+            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 failed: {e}")
+            print(f"🔴 WebSocket Connection Failed: {e}")
 
     def run(self):
         asyncio.set_event_loop(self.loop)
@@ -70,6 +70,7 @@ class WebSocketClient:
             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)
@@ -84,125 +85,140 @@ def send_audio(chunk, cid):
     connections[cid].enqueue_audio_chunk(sr, arr)
     return connections[cid].transcript.strip()
 
-def clear_transcript(cid):
+def clear_transcript_only(cid):
     if cid in connections:
         connections[cid].transcript = ""
     return ""
 
-def format_response(content, prompt):
-    md = f"""### ❓ {prompt}\n\n**🧠 In summary:**\n\n{content}\n"""
-    images = re.findall(r'https://raw\.githubusercontent\.com/[^\s)]+\.png', content)
-    if images:
-        md += "\n\n**📎 Sources:**\n" + "\n".join([f"![]({url})" for url in images])
-    return md
+def clear_chat_only():
+    return [], None, None
 
-def handle_chat(prompt, thread_id):
+# Assistant chat handler
+def handle_chat(user_input, history, thread_id, image_url):
     if not OPENAI_API_KEY or not ASSISTANT_ID:
-        return "❌ Missing credentials", thread_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=prompt)
+
+        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":
-                return format_response(msg.content[0].text.value, prompt), thread_id
-        return "⚠️ No reply", thread_id
+                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}", thread_id
+        return f"❌ {e}", history, thread_id, image_url
 
-def feed_transcript(transcript, thread_id, cid):
+# Feed transcript as assistant input
+def feed_transcript(transcript, history, thread_id, image_url, cid):
     if not transcript.strip():
-        return gr.update(), thread_id
+        return gr.update(), history, thread_id, image_url
     if cid in connections:
         connections[cid].transcript = ""
-    return handle_chat(transcript, thread_id)
-
-with gr.Blocks(css="""
-body {
-    background-color: #0f0f0f;
-    color: white;
-    font-family: 'Inter', sans-serif;
-}
-.markdown-container {
-    margin-top: 16px;
-    padding: 18px;
-    background: #1a1a1a;
-    border-radius: 12px;
-    font-size: 16px;
-    box-shadow: 0 2px 6px #000;
-}
-.input-bar {
-    position: fixed;
-    bottom: 16px;
-    left: 0;
-    right: 0;
-    max-width: 1000px;
-    margin: auto;
-    display: flex;
-    padding: 12px;
-    gap: 10px;
-    background: #1f1f1f;
-    border-radius: 16px;
-}
-#user_input {
-    flex: 1;
-    padding: 12px;
-    font-size: 16px;
-    background: #292929;
-    border: none;
-    color: white;
-    border-radius: 10px;
-}
-.btn {
-    background: #4f46e5;
-    color: white;
-    border: none;
-    border-radius: 10px;
-    padding: 10px 14px;
-    font-size: 18px;
-}
-.voice-area {
-    background: #222;
-    padding: 14px;
-    margin-top: 20px;
-    border-radius: 12px;
-    display: flex;
-    flex-direction: column;
-    gap: 12px;
-}
-""") as app:
+    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()
-    voice_visible = gr.State(False)
-
-    gr.Markdown("<h1 style='text-align:center;'>How can I help you today?</h1>")
-    response = gr.Markdown(elem_classes="markdown-container")
-
-    with gr.Row(elem_classes="input-bar"):
-        prompt = gr.Textbox(placeholder="Ask a question...", elem_id="user_input", show_label=False)
-        send = gr.Button("➤", elem_id="send", elem_classes="btn")
-        mic = gr.Button("🎙", elem_id="mic", elem_classes="btn")
-
-    with gr.Column(visible=False, elem_classes="voice-area") as voice_box:
-        voice = gr.Audio(label="Tap to Speak", streaming=True, type="numpy")
-        transcript = gr.Textbox(label="Transcript", lines=2, interactive=False)
-        send_voice = gr.Button("Send Voice")
-        clear = gr.Button("Clear Transcript")
-
-    send.click(handle_chat, [prompt, thread_state], [response, thread_state])
-    mic.click(lambda x: not x, voice_visible, voice_visible)
-    voice_visible.change(fn=None, inputs=voice_visible, outputs=voice_box, show_progress=False)
-    voice.stream(send_audio, [voice, client_id], transcript, stream_every=0.5)
-    send_voice.click(feed_transcript, [transcript, thread_state, client_id], [response, thread_state])
-    clear.click(clear_transcript, [client_id], transcript)
+
+    with gr.Row(equal_height=True):
+        with gr.Column(scale=0.8):  # thinner image column
+            image_display = gr.Image(label="🖼️ Document", type="filepath", show_download_button=False)
+
+            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.Column(scale=2):  # wider chat column
+            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.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()