Update app.py

app.py

@@ -1,103 +1,105 @@
 import gradio as gr
 import os
+import json
 import uuid
+import threading
+import time
+import re
+from dotenv import load_dotenv
 from openai import OpenAI
 from realtime_transcriber import WebSocketClient, connections, WEBSOCKET_URI, WEBSOCKET_HEADERS
 
-# Load OpenAI API key
-OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY")
-if not OPENAI_API_KEY:
-    raise ValueError("OPENAI_API_KEY environment variable must be set")
+# ------------------ 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")
+
 client = OpenAI(api_key=OPENAI_API_KEY)
+session_threads = {}
 
-# Session state
-session_id = str(uuid.uuid4())
-if session_id not in connections:
-    connections[session_id] = WebSocketClient(WEBSOCKET_URI, WEBSOCKET_HEADERS, session_id)
-    import threading
-    threading.Thread(target=connections[session_id].run, daemon=True).start()
-
-# Functions for Document Assistant
-def process_user_input(message, history):
-    if not message:
-        return "Please enter a message.", history
-
-    try:
-        thread = client.beta.threads.create()
-        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=os.environ.get("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
-        messages = client.beta.threads.messages.list(thread_id=thread.id)
-        assistant_reply = next((m.content[0].text.value for m in reversed(messages.data) if m.role == "assistant"), "No response.")
-        history.append((message, assistant_reply))
-        return "", history
-    except Exception as e:
-        return f"❌ Error: {str(e)}", history
-
-# Functions for Realtime Voice Transcription
-def send_audio_chunk_realtime(mic_chunk):
-    if session_id not in connections:
-        return "Initializing voice session..."
-    if mic_chunk is not None:
-        sr, y = mic_chunk
-        connections[session_id].enqueue_audio_chunk(sr, y)
-    return connections[session_id].transcript
-
-def clear_transcript():
-    if session_id in connections:
-        connections[session_id].transcript = ""
-    return ""
+# ------------------ Chat Logic ------------------
+def reset_session():
+    session_id = str(uuid.uuid4())
+    session_threads[session_id] = client.beta.threads.create().id
+    return session_id
+
+def process_chat(message, history, session_id):
+    thread_id = session_threads.get(session_id)
+    if not thread_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 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)
+    for msg in reversed(messages.data):
+        if msg.role == "assistant":
+            return msg.content[0].text.value
+    return "⚠️ Assistant did not respond."
 
-# Gradio UI Components
-doc_image = gr.Image(label="📘 Extracted Document Image", show_label=True, elem_id="docimg", height=500, width=360)
-chatbot = gr.Chatbot(label="🧠 Document Assistant", elem_id="chatbox", bubble_full_width=False)
-prompt = gr.Textbox(placeholder="Ask about the document...", label="Ask about the document")
-send_btn = gr.Button("Send")
+def extract_image_url(text):
+    match = re.search(r'https://raw\.githubusercontent\.com/[^\s"]+\.png', text)
+    return match.group(0) if match else None
 
-# Voice Section
-audio_in = gr.Audio(label="🎵 Audio", type="numpy", streaming=True)
-live_transcript = gr.Textbox(label="Live Transcript", lines=6)
-clear_btn = gr.Button("Clear Transcript")
+def handle_chat(message, history, session_id):
+    response = process_chat(message, history, session_id)
+    history.append((message, response))
+    image = extract_image_url(response)
+    return history, image
 
-with gr.Blocks(theme=gr.themes.Base(), css="""
-    #docimg img { object-fit: contain !important; }
-    #chatbox { height: 500px; }
-    .gr-box { border-radius: 12px; }
-""") as demo:
+# ------------------ Voice Logic ------------------
+def create_websocket_client():
+    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 = ""
+    return ""
+
+def send_audio_chunk(audio, client_id):
+    if client_id not in connections:
+        return "Initializing connection..."
+    sr, y = audio
+    connections[client_id].enqueue_audio_chunk(sr, y)
+    return connections[client_id].transcript
+
+# ------------------ UI ------------------
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.Markdown("# 🧠 Document AI + 🎙️ Voice Assistant")
-    with gr.Row():
-        with gr.Column(scale=1):
-            doc_image.render()
-        with gr.Column(scale=2):
-            chatbot.render()
+
+    session_id = gr.State(value=reset_session())
+    client_id = gr.State()
 
     with gr.Row():
-        prompt.render()
-        send_btn.render()
+        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(fn=process_user_input, inputs=[prompt, chatbot], outputs=[prompt, chatbot])
+    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])
 
-    with gr.Accordion("🎙️ Or Use Voice Instead", open=False):
-        live_transcript.render()
+    # 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():
-            audio_in.render()
-            clear_btn.render()
-        audio_in.stream(fn=send_audio_chunk_realtime, inputs=audio_in, outputs=live_transcript)
-        clear_btn.click(fn=clear_transcript, outputs=live_transcript)
+            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)
 
-# LAUNCH WITH SHARE ENABLED FOR PUBLIC URL
-demo.launch(share=True)
+demo.launch()