Update app.py

app.py

@@ -1,4 +1,3 @@
-# app.py
 import gradio as gr
 import os
 import requests
@@ -19,7 +18,7 @@ MODEL_ENDPOINTS = {
 }
 
 def query_model(model_name: str, messages: List[Dict[str, str]]) -> str:
-    """Query a single model with the chat history"""
+    """Query a single model with the chat history."""
     endpoint = MODEL_ENDPOINTS[model_name]
     headers = {
         "Authorization": f"Bearer {HF_API_KEY}",
@@ -29,7 +28,7 @@ def query_model(model_name: str, messages: List[Dict[str, str]]) -> str:
     # Build full conversation history for context
     conversation = "\n".join([f"{msg['role']}: {msg['content']}" for msg in messages])
     
-    # Model-specific prompt formatting with full history
+    # Model-specific prompt formatting
     model_prompts = {
         "Qwen2.5-72B-Instruct": (
             f"<|im_start|>system\nCollaborate with other experts. Previous discussion:\n{conversation}<|im_end|>\n"
@@ -67,14 +66,15 @@ def query_model(model_name: str, messages: List[Dict[str, str]]) -> str:
         response = requests.post(endpoint, json=payload, headers=headers)
         response.raise_for_status()
         result = response.json()[0]['generated_text']
-        # Do NOT remove ** or ## so that Markdown (including LaTeX) is preserved
-        result = result.split('<|')[0]  # Remove any remaining special tokens
-        return result.strip()  # Return the cleaned response
+        # Basic cleanup
+        result = result.split('<|')[0]  # Remove special tokens
+        result = result.replace('**', '').replace('##', '')
+        return result.strip()
     except Exception as e:
         return f"{model_name} error: {str(e)}"
 
 def respond(message: str, history: List[List[str]], session_id: str) -> Generator[str, None, None]:
-    """Handle sequential model responses with context preservation"""
+    """Handle sequential model responses with context preservation."""
     # Load or initialize session
     session = session_manager.load_session(session_id)
     if not isinstance(session, dict) or "history" not in session:
@@ -131,34 +131,51 @@ def respond(message: str, history: List[List[str]], session_id: str) -> Generato
     })
     messages.append({"role": "assistant", "content": f"Llama3.3-70B-Instruct: {response3}"})
     
-    # Save final session state
+    # Save session
     session_manager.save_session(session_id, session)
     
     # Return final combined response
-    yield f"🔵 **Qwen2.5-Coder-32B-Instruct**\n{response1}\n\n🟣 **Qwen2.5-72B-Instruct**\n{response2}\n\n🟡 **Llama3.3-70B-Instruct**\n{response3}"
+    yield (
+        f"🔵 **Qwen2.5-Coder-32B-Instruct**\n{response1}\n\n"
+        f"🟣 **Qwen2.5-72B-Instruct**\n{response2}\n\n"
+        f"🟡 **Llama3.3-70B-Instruct**\n{response3}"
+    )
 
-# Create the Gradio interface
 with gr.Blocks() as demo:
-    gr.Markdown("## Multi-LLM Collaboration Chat")
-    gr.Markdown(
-        """
-        **LaTeX Support:** You can include inline math like `$E = mc^2$` or block math like:
-        ```
-        $$
-        \\int_0^1 x^2 \\, dx
-        $$
-        ```
-        and it will be rendered correctly below.
-        """
-    )
-    
+    # -- Include KaTeX for LaTeX rendering --
+    gr.HTML("""
+    <link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/[email protected]/dist/katex.min.css" />
+    <script defer src="https://cdn.jsdelivr.net/npm/[email protected]/dist/katex.min.js"></script>
+    <script defer src="https://cdn.jsdelivr.net/npm/[email protected]/dist/contrib/auto-render.min.js"></script>
+    <script>
+    // Re-render math whenever new content is added
+    document.addEventListener("DOMContentLoaded", function() {
+      const observer = new MutationObserver(function(mutations) {
+        for (const mutation of mutations) {
+          if (mutation.type === 'childList') {
+            renderMathInElement(document.body, {
+              delimiters: [
+                {left: "$$", right: "$$", display: true},
+                {left: "$", right: "$", display: false},
+              ]
+            });
+          }
+        }
+      });
+      observer.observe(document.body, { subtree: true, childList: true });
+    });
+    </script>
+    """)
+
+    gr.Markdown("## Multi-LLM Collaboration Chat (with LaTeX support)")
+
     with gr.Row():
         session_id = gr.State(session_manager.create_session)
         new_session = gr.Button("🔄 New Session")
     
     chatbot = gr.Chatbot(height=600)
-    msg = gr.Textbox(label="Message")
-    
+    msg = gr.Textbox(label="Message (Use $...$ or $$...$$ for LaTeX)")
+
     def on_new_session():
         new_id = session_manager.create_session()
         return new_id, []
@@ -172,7 +189,7 @@ with gr.Blocks() as demo:
             for response in respond(message, history[:-1], session_id):
                 history[-1][1] = response
                 yield history
-    
+
     msg.submit(user, [msg, chatbot, session_id], [msg, chatbot]).then(
         bot, [chatbot, session_id], [chatbot]
     )