Update app.py

app.py

@@ -5,28 +5,25 @@ import uuid
 import threading
 import time
 import re
-
-from openai import OpenAI
 from dotenv import load_dotenv
+from openai import OpenAI
 from realtime_transcriber import WebSocketClient, connections, WEBSOCKET_URI, WEBSOCKET_HEADERS
 
-# ------------------ Load API Key ------------------
+# ------------------ Load Secrets ------------------
 load_dotenv()
 OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
 ASSISTANT_ID = os.getenv("ASSISTANT_ID")
 
 if not OPENAI_API_KEY or not ASSISTANT_ID:
-    raise ValueError("Missing OPENAI_API_KEY or ASSISTANT_ID in environment variables")
+    raise ValueError("Missing OPENAI_API_KEY or ASSISTANT_ID")
 
 client = OpenAI(api_key=OPENAI_API_KEY)
-
-# ------------------ Chat Threading ------------------
 session_threads = {}
 
+# ------------------ Chat Logic ------------------
 def reset_session():
     session_id = str(uuid.uuid4())
-    thread = client.beta.threads.create()
-    session_threads[session_id] = thread.id
+    session_threads[session_id] = client.beta.threads.create().id
     return session_id
 
 def process_chat(message, history, session_id):
@@ -35,50 +32,29 @@ def process_chat(message, history, session_id):
         thread_id = client.beta.threads.create().id
         session_threads[session_id] = thread_id
 
-    client.beta.threads.messages.create(
-        thread_id=thread_id,
-        role="user",
-        content=message
-    )
-
-    run = client.beta.threads.runs.create(
-        thread_id=thread_id,
-        assistant_id=ASSISTANT_ID
-    )
-
-    while True:
-        run_status = client.beta.threads.runs.retrieve(
-            thread_id=thread_id,
-            run_id=run.id
-        )
-        if run_status.status == "completed":
-            break
+    client.beta.threads.messages.create(thread_id=thread_id, role="user", content=message)
+    run = client.beta.threads.runs.create(thread_id=thread_id, assistant_id=ASSISTANT_ID)
+
+    while client.beta.threads.runs.retrieve(thread_id=thread_id, run_id=run.id).status != "completed":
         time.sleep(1)
 
     messages = client.beta.threads.messages.list(thread_id=thread_id)
-    assistant_response = "⚠️ Assistant did not respond."
     for msg in reversed(messages.data):
         if msg.role == "assistant":
-            assistant_response = msg.content[0].text.value
-            break
-
-    return assistant_response
+            return msg.content[0].text.value
+    return "⚠️ Assistant did not respond."
 
 def extract_image_url(text):
-    match = re.search(
-        r'https://raw\.githubusercontent\.com/AndrewLORTech/surgical-pathology-manual/main/[\w\-/]*\.png',
-        text
-    )
+    match = re.search(r'https://raw\.githubusercontent\.com/[^\s"]+\.png', text)
     return match.group(0) if match else None
 
-def chat_handler(message, history, session_id):
-    history.append(("user", message))
+def handle_chat(message, history, session_id):
     response = process_chat(message, history, session_id)
-    history.append(("assistant", response))
-    image_url = extract_image_url(response)
-    return history, image_url
+    history.append((message, response))
+    image = extract_image_url(response)
+    return history, image
 
-# ------------------ Transcription ------------------
+# ------------------ Voice Logic ------------------
 def create_websocket_client():
     client_id = str(uuid.uuid4())
     connections[client_id] = WebSocketClient(WEBSOCKET_URI, WEBSOCKET_HEADERS, client_id)
@@ -97,7 +73,7 @@ def send_audio_chunk(audio, client_id):
     connections[client_id].enqueue_audio_chunk(sr, y)
     return connections[client_id].transcript
 
-# ------------------ Gradio App ------------------
+# ------------------ UI ------------------
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.Markdown("# 🧠 Document AI + 🎙️ Voice Assistant")
 
@@ -105,33 +81,25 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
     client_id = gr.State()
 
     with gr.Row():
-        with gr.Column(scale=1):
-            image_display = gr.Image(label="📑 Extracted Document Image", show_label=True, height=400)
-
-        with gr.Column(scale=2):
-            chatbot = gr.Chatbot(label="💬 Document Assistant", height=400)
-            message_input = gr.Textbox(label="Ask about the document", placeholder="e.g. What does clause 3.2 mean?")
-            send_button = gr.Button("Send")
-
-    # Send message logic
-    def user_send(msg, history, session_id):
-        return chat_handler(msg, history, session_id)
-
-    send_button.click(user_send, inputs=[message_input, chatbot, session_id], outputs=[chatbot, image_display])
-    message_input.submit(user_send, inputs=[message_input, chatbot, session_id], outputs=[chatbot, image_display])
-
-    # ------------------ Voice Section ------------------
-    gr.Markdown("## 🎙️ Realtime Voice Transcription")
-
-    with gr.Row():
-        transcript_box = gr.Textbox(label="Live Transcript", lines=7, interactive=False, autoscroll=True)
-
-    with gr.Row():
-        mic_input = gr.Audio(streaming=True)
-        clear_button = gr.Button("Clear Transcript")
-
-    mic_input.stream(fn=send_audio_chunk, inputs=[mic_input, client_id], outputs=transcript_box)
-    clear_button.click(fn=clear_transcript, inputs=[client_id], outputs=transcript_box)
-    demo.load(fn=create_websocket_client, outputs=client_id)
+        image_display = gr.Image(label="📑 Extracted Document Image", show_label=True, height=360)
+        with gr.Column():
+            chatbot = gr.Chatbot(label="💬 Document Assistant", height=360)
+            text_input = gr.Textbox(label="Ask about the document", placeholder="e.g. What is clause 3.2?")
+            send_btn = gr.Button("Send")
+
+    send_btn.click(handle_chat, inputs=[text_input, chatbot, session_id], outputs=[chatbot, image_display])
+    text_input.submit(handle_chat, inputs=[text_input, chatbot, session_id], outputs=[chatbot, image_display])
+
+    # Toggle Section
+    with gr.Accordion("🎤 Or Use Voice Instead", open=False):
+        with gr.Row():
+            transcript_box = gr.Textbox(label="Live Transcript", lines=7, interactive=False, autoscroll=True)
+        with gr.Row():
+            mic_input = gr.Audio(streaming=True)
+            clear_button = gr.Button("Clear Transcript")
+
+        mic_input.stream(fn=send_audio_chunk, inputs=[mic_input, client_id], outputs=transcript_box)
+        clear_button.click(fn=clear_transcript, inputs=[client_id], outputs=transcript_box)
+        demo.load(fn=create_websocket_client, outputs=client_id)
 
 demo.launch()