Update app.py

app.py

@@ -5,6 +5,7 @@ import threading
 from datetime import datetime
 from typing import List, Dict, Any, Generator
 from session_manager import SessionManager
+import katex
 
 # Initialize session manager and get HF API key
 session_manager = SessionManager()
@@ -66,9 +67,16 @@ 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']
-        # Clean up response formatting
+        # Clean up response formatting and handle LaTeX
         result = result.split('<|')[0]  # Remove any remaining special tokens
         result = result.replace('**', '').replace('##', '')  # Remove markdown
+        
+        # Detect and format LaTeX content
+        import re
+        # Handle inline LaTeX
+        result = re.sub(r'\\\((.*?)\\\)', r'<span class="katex">\1</span>', result)
+        # Handle display LaTeX
+        result = re.sub(r'\\\[(.*?)\\\]', r'<div class="katex">\1</div>', result)
         result = result.strip()  # Remove leading/trailing whitespace
         return result  # Return complete response
     except Exception as e:
@@ -138,9 +146,22 @@ def respond(message: str, history: List[List[str]], session_id: str) -> Generato
     # 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}"
 
-# Create the Gradio interface
+# Create the Gradio interface with LaTeX support
 with gr.Blocks() as demo:
     gr.Markdown("## Multi-LLM Collaboration Chat")
+    # Add KaTeX CSS and JS
+    demo.css = """
+    .katex { font-size: 1.1em; }
+    """
+    demo.js = """
+    function renderLatex() {
+        document.querySelectorAll('.katex').forEach(el => {
+            katex.render(el.textContent, el, {
+                throwOnError: false
+            });
+        });
+    }
+    """
     
     with gr.Row():
         session_id = gr.State(session_manager.create_session)