Update app.py

app.py

@@ -74,7 +74,7 @@ def query_model(model_name: str, messages: List[Dict[str, str]]) -> str:
     except Exception as e:
         return f"{model_name} error: {str(e)}"
 
-def respond(message: str, history: List[List[str]], session_id: str) -> tuple[str, str]:
+def respond(message: str, history: List[List[str]], session_id: str) -> Generator[tuple[str, str], None, None]:
     """Handle sequential model responses with context preservation"""
     # Load or initialize session
     session = session_manager.load_session(session_id)
@@ -97,9 +97,8 @@ def respond(message: str, history: List[List[str]], session_id: str) -> tuple[st
         "content": message
     })
     
-    responses = []
-    
-    # Get first model's response
+    # First model response
+    yield message, "💭 *Qwen2.5-Coder-32B-Instruct is thinking...*"
     response1 = query_model("Qwen2.5-Coder-32B-Instruct", messages)
     session["history"].append({
         "timestamp": datetime.now().isoformat(),
@@ -108,9 +107,10 @@ def respond(message: str, history: List[List[str]], session_id: str) -> tuple[st
         "content": response1
     })
     messages.append({"role": "assistant", "content": f"Qwen2.5-Coder-32B-Instruct: {response1}"})
-    responses.append(f"**Qwen2.5-Coder-32B-Instruct**:\n{response1}")
+    yield message, f"🔵 **Qwen2.5-Coder-32B-Instruct**\n{response1}"
     
-    # Get second model's response
+    # Second model response
+    yield message, f"🔵 **Qwen2.5-Coder-32B-Instruct**\n{response1}\n\n💭 *Qwen2.5-72B-Instruct is thinking...*"
     response2 = query_model("Qwen2.5-72B-Instruct", messages)
     session["history"].append({
         "timestamp": datetime.now().isoformat(),
@@ -119,9 +119,10 @@ def respond(message: str, history: List[List[str]], session_id: str) -> tuple[st
         "content": response2
     })
     messages.append({"role": "assistant", "content": f"Qwen2.5-72B-Instruct: {response2}"})
-    responses.append(f"**Qwen2.5-72B-Instruct**:\n{response2}")
+    yield message, f"🔵 **Qwen2.5-Coder-32B-Instruct**\n{response1}\n\n🟣 **Qwen2.5-72B-Instruct**\n{response2}"
     
-    # Get final model's response
+    # Final model response
+    yield message, f"🔵 **Qwen2.5-Coder-32B-Instruct**\n{response1}\n\n🟣 **Qwen2.5-72B-Instruct**\n{response2}\n\n💭 *Llama3.3-70B-Instruct is thinking...*"
     response3 = query_model("Llama3.3-70B-Instruct", messages)
     session["history"].append({
         "timestamp": datetime.now().isoformat(),
@@ -130,44 +131,49 @@ def respond(message: str, history: List[List[str]], session_id: str) -> tuple[st
         "content": response3
     })
     messages.append({"role": "assistant", "content": f"Llama3.3-70B-Instruct: {response3}"})
-    responses.append(f"**Llama3.3-70B-Instruct**:\n{response3}")
     
     # Save final session state
     session_manager.save_session(session_id, session)
     
-    # Return response as a single tuple for Gradio chat
-    return message, "\n\n".join(responses)
+    # Return final combined response
+    yield message, f"🔵 **Qwen2.5-Coder-32B-Instruct**\n{response1}\n\n🟣 **Qwen2.5-72B-Instruct**\n{response2}\n\n🟡 **Llama3.3-70B-Instruct**\n{response3}"
+
+# Custom CSS for styling
+css = """
+.message { padding: 15px; margin: 10px 0; border-radius: 10px; }
+.assistant { background: #f8fafc; border-left: 4px solid #3b82f6; }
+.user { background: #eff6ff; border-left: 4px solid #60a5fa; }
+.model-name { font-weight: 600; color: #1e40af; margin-bottom: 8px; }
+.thinking { color: #6b7280; font-style: italic; }
+"""
 
 # Create the Gradio interface
-with gr.Blocks() as demo:
-    session_id = gr.State(session_manager.create_session)
-    
-    gr.Markdown("## Multi-LLM Collaboration Chat")
-    gr.Markdown("A group chat with Qwen2.5-72B, Llama3.3-70B, and Qwen2.5-Coder-32B")
-    
-    chatbot = gr.Chatbot()
-    msg = gr.Textbox(label="Message")
-    clear = gr.Button("Clear")
-    
-    def user(message, history, session_id):
-        return "", history + [[message, None]]
-    
-    def bot(history, session_id):
-        if history[-1][1] is None:
-            message = history[-1][0]
-            _, response = respond(message, history[:-1], session_id)
-            history[-1][1] = response
-            return history
-        return history
-    
-    msg.submit(user, [msg, chatbot, session_id], [msg, chatbot]).then(
-        bot, [chatbot, session_id], [chatbot]
-    )
-    
-    clear.click(lambda: (session_manager.create_session(), None, []), 
-                None, 
-                [session_id, msg, chatbot],
-                queue=False)
+demo = gr.ChatInterface(
+    fn=respond,
+    title="Multi-LLM Collaboration Chat",
+    description="Experience collaborative AI thinking with three powerful language models",
+    examples=[
+        ["Explain how quantum computing works"],
+        ["Write a Python function to find prime numbers"],
+    ],
+    additional_inputs=[gr.State(session_manager.create_session)],
+    chatbot=gr.Chatbot(
+        height=600,
+        show_label=False,
+        avatar_images=("👤", "🤖"),
+        bubble_full_width=False,
+        show_copy_button=True,
+        container=True,
+        sanitize_html=False,
+    ),
+    theme=gr.themes.Soft(
+        primary_hue="blue",
+        secondary_hue="indigo",
+        neutral_hue="slate",
+        font=("Inter", "sans-serif"),
+    ),
+    css=css,
+)
 
 if __name__ == "__main__":
     demo.launch(share=True)