Update app.py

app.py

@@ -3,7 +3,7 @@
 Multi-Mode AI Assistant (Voice, PDF, Image) with colorful website-like UI
 - All functionality preserved
 - Dark theme, gradient buttons, visible text
-- Gradio Blocks with columns as tabs
+- Chat bubbles for user/assistant messages
 """
 import os
 import uuid
@@ -65,14 +65,31 @@ def select_relevant_chunk(question,chunks,chunk_embeds):
     top_idx=int(scores.argmax().item())
     return chunks[top_idx]
 
-def _chat_display_to_messages(chat_display):
-    msgs=[]
-    for user,assistant in chat_display:
-        msgs.append({"role":"user","content":user})
-        msgs.append({"role":"assistant","content":assistant})
-    return msgs
+def _append_chat_display(session_id,user_text,assistant_text):
+    if session_id not in CHAT_DISPLAY: CHAT_DISPLAY[session_id]=[]
+    CHAT_DISPLAY[session_id].append((user_text,assistant_text))
 
-# ------------------ LLM & Transcription ------------------
+def format_chat_messages(chat_display):
+    """
+    Returns a list of tuples (html_text, None) for gr.Chatbot with colored bubbles.
+    """
+    formatted=[]
+    for user_msg, assistant_msg in chat_display:
+        user_html=f"""
+        <div style="background:#38bdf8;color:#0f172a;padding:8px 12px;border-radius:12px;max-width:70%;margin:5px 0;align-self:flex-end;">
+            {user_msg}
+        </div>
+        """
+        assistant_html=f"""
+        <div style="background:#facc15;color:#1e293b;padding:8px 12px;border-radius:12px;max-width:70%;margin:5px 0;align-self:flex-start;">
+            {assistant_msg}
+        </div>
+        """
+        formatted.append((user_html,None))
+        formatted.append((assistant_html,None))
+    return formatted
+
+# ------------------ Transcription & LLM ------------------
 def transcribe_audio(audio_path):
     if not audio_path or not os.path.exists(audio_path):
         return "Error: audio file missing."
@@ -185,42 +202,38 @@ def download_pdf_summary(session_id):
     if not summary: summary="No summary available."
     return generate_pdf_file(summary,"summary")
 
-# ------------------ Voice ------------------
-def _append_chat_display(session_id,user_text,assistant_text):
-    if session_id not in CHAT_DISPLAY: CHAT_DISPLAY[session_id]=[]
-    CHAT_DISPLAY[session_id].append((user_text,assistant_text))
-
+# ------------------ Voice & Text Handlers ------------------
 def handle_voice_general(audio_file,session_id,tts_lang="en",enhancer_enabled=False,enhancer_tone="Helpful"):
     path=_get_path_from_gr_file(audio_file)
-    if not path: return "No audio",None,[]
+    if not path: return "No audio", None, []
     user_text=transcribe_audio(path)
     assistant_text=generate_response(session_id,user_text,enhancer_enabled,enhancer_tone)
     _append_chat_display(session_id,user_text,assistant_text)
     audio_path=synthesize_speech(assistant_text,tts_lang)
-    return assistant_text,audio_path,_chat_display_to_messages(CHAT_DISPLAY[session_id])
+    return assistant_text, audio_path, format_chat_messages(CHAT_DISPLAY[session_id])
 
 def handle_voice_pdf(audio_file,session_id,tts_lang="en"):
     path=_get_path_from_gr_file(audio_file)
-    if not path: return "No audio",None,[]
+    if not path: return "No audio", None, []
     user_text=transcribe_audio(path)
     assistant_text=handle_pdf_question(user_text,session_id)
     _append_chat_display(session_id,user_text,assistant_text)
     audio_path=synthesize_speech(assistant_text,tts_lang)
-    return assistant_text,audio_path,_chat_display_to_messages(CHAT_DISPLAY[session_id])
+    return assistant_text, audio_path, format_chat_messages(CHAT_DISPLAY[session_id])
 
 def handle_voice_image(audio_file,session_id,tts_lang="en"):
     path=_get_path_from_gr_file(audio_file)
-    if not path: return "No audio",None,[]
+    if not path: return "No audio", None, []
     user_text=transcribe_audio(path)
     assistant_text=handle_image_question(user_text,session_id)
     _append_chat_display(session_id,user_text,assistant_text)
     audio_path=synthesize_speech(assistant_text,tts_lang)
-    return assistant_text,audio_path,_chat_display_to_messages(CHAT_DISPLAY[session_id])
+    return assistant_text, audio_path, format_chat_messages(CHAT_DISPLAY[session_id])
 
 def handle_text_general(user_text,session_id,enhancer_enabled=False,enhancer_tone="Helpful"):
     assistant=generate_response(session_id,user_text,enhancer_enabled,enhancer_tone)
     _append_chat_display(session_id,user_text,assistant)
-    return assistant,_chat_display_to_messages(CHAT_DISPLAY[session_id])
+    return assistant, format_chat_messages(CHAT_DISPLAY[session_id])
 
 def handle_text_pdf(question,session_id): return handle_pdf_question(question,session_id)
 def handle_text_image(question,session_id): return handle_image_question(question,session_id)
@@ -251,7 +264,7 @@ with gr.Blocks() as demo:
 
     # --- Voice Tab ---
     col_voice=gr.Column(visible=True)
-    chat_voice=gr.Chatbot(height=300)
+    chat_voice=gr.Chatbot(height=400)
     with gr.Row():
         mic=gr.Audio(type="filepath",label="🎤 Record Voice")
         audio_output=gr.Audio(type="filepath",label="Assistant Voice",interactive=False)