Update app.py

app.py

@@ -2,6 +2,7 @@ import gradio as gr
 import requests
 import json
 import os
+import re
 from datetime import datetime
 
 
@@ -27,10 +28,33 @@ def get_ai_response(message, history):
     try:
         response = requests.post(API_ENDPOINT, headers=headers, json=payload)
         response.raise_for_status()
-        return response.json()["choices"][0]["message"]["content"]
+        raw_response = response.json()["choices"][0]["message"]["content"]
+        
+        # Convert reasoning tags to collapsible HTML
+        html_response = convert_reasoning_to_collapsible(raw_response)
+        return html_response
     except Exception as e:
         return f"Error: {str(e)}"
 
+def convert_reasoning_to_collapsible(text):
+    """Convert reasoning tags to collapsible HTML sections."""
+    # Find all reasoning sections
+    reasoning_pattern = re.compile(r'<reasoning>(.*?)</reasoning>', re.DOTALL)
+    
+    # Function to replace each reasoning section with collapsible HTML
+    def replace_with_collapsible(match):
+        reasoning_content = match.group(1).strip()
+        return f'<details><summary><strong>See reasoning</strong></summary><div class="reasoning-content">{reasoning_content}</div></details>'
+    
+    # Replace reasoning tags with collapsible sections
+    html_response = reasoning_pattern.sub(replace_with_collapsible, text)
+    
+    # Remove <sep> tags
+    html_response = re.sub(r'<sep>.*?</sep>', '', html_response, flags=re.DOTALL)
+    html_response = html_response.replace('<sep>', '').replace('</sep>', '')
+    
+    return html_response
+
 def chat_interface(message, history, stored_history):
     """Handle chat interactions and update history."""
     if not history:
@@ -39,8 +63,12 @@ def chat_interface(message, history, stored_history):
     # Convert history to the format expected by the API
     api_history = []
     for user_msg, ai_msg in history:
+        # Remove HTML tags for API history
+        clean_ai_msg = re.sub(r'<details>.*?</details>', '', ai_msg, flags=re.DOTALL)
+        clean_ai_msg = re.sub(r'<[^>]*>', '', clean_ai_msg)
+        
         api_history.append({"role": "user", "content": user_msg})
-        api_history.append({"role": "assistant", "content": ai_msg})
+        api_history.append({"role": "assistant", "content": clean_ai_msg})
     
     ai_response = get_ai_response(message, api_history)
     
@@ -50,10 +78,15 @@ def chat_interface(message, history, stored_history):
     timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
     if stored_history is None:
         stored_history = []
+    
+    # Store a plain text version for history display
+    plain_response = re.sub(r'<details>.*?</details>', '[Reasoning available]', ai_response, flags=re.DOTALL)
+    plain_response = re.sub(r'<[^>]*>', '', plain_response)
+    
     stored_history.insert(0, {
         "timestamp": timestamp,
         "user": message,
-        "ai": ai_response
+        "ai": plain_response
     })
     return history, stored_history
 
@@ -67,8 +100,8 @@ def update_history_display(stored_history):
             html += f"""
                 <li class='history-item'>
                     <small>{item['timestamp']}</small><br>
-                    <strong>You:</strong> {item['user'][:50]}...<br>
-                    <strong>AI:</strong> {item['ai'][:50]}...
+                    <strong>You:</strong> {item['user'][:50]}{'...' if len(item['user']) > 50 else ''}<br>
+                    <strong>AI:</strong> {item['ai'][:50]}{'...' if len(item['ai']) > 50 else ''}
                 </li>
             """
         html += "</ul>"
@@ -91,10 +124,24 @@ body { background-color: #1a1a1a; color: #ffffff; font-family: 'Arial', sans-ser
 .history-item:hover { background-color: #404040; }
 input, button { background-color: #333333; color: #ffffff; border: 1px solid #404040; border-radius: 5px; }
 button:hover { background-color: #404040; }
+details { background-color: #333333; padding: 10px; margin: 5px 0; border-radius: 5px; }
+summary { cursor: pointer; color: #70a9e6; }
+.reasoning-content { padding: 10px; margin-top: 5px; background-color: #404040; border-radius: 5px; }
+"""
+
+# HTML head for rendering HTML in chatbot
+html_head = """
+<head>
+<style>
+details { background-color: #333333; padding: 10px; margin: 5px 0; border-radius: 5px; }
+summary { cursor: pointer; color: #70a9e6; }
+.reasoning-content { padding: 10px; margin-top: 5px; background-color: #404040; border-radius: 5px; }
+</style>
+</head>
 """
 
 # Build the Gradio app
-with gr.Blocks(css=custom_css, title="Reka Flash 3 Demo") as demo:
+with gr.Blocks(css=custom_css, title="AI Assistant with Collapsible Reasoning") as demo:
     with gr.Row():
         # Sidebar for history
         with gr.Column(scale=1, min_width=300, elem_id="sidebar"):
@@ -104,7 +151,9 @@ with gr.Blocks(css=custom_css, title="Reka Flash 3 Demo") as demo:
         
         # Main chat area
         with gr.Column(scale=3):
-            chatbot = gr.Chatbot(elem_id="chatbot")  # Removed type="messages" which was causing errors
+            gr.Markdown("## AI Assistant")
+            gr.Markdown("This assistant shows reasoning in collapsible sections.")
+            chatbot = gr.Chatbot(elem_id="chatbot", render_markdown=False, bubble_full_width=True)
             with gr.Row():
                 message = gr.Textbox(placeholder="Type your message...", show_label=False, container=False)
                 submit_btn = gr.Button("Send", size="sm")
@@ -121,6 +170,35 @@ with gr.Blocks(css=custom_css, title="Reka Flash 3 Demo") as demo:
         return history ? JSON.parse(history) : [];
     }
     """
+    
+    # JavaScript for enabling HTML in chatbot
+    js = """
+    function() {
+        // Add event listener for when new messages are added
+        const observer = new MutationObserver(function(mutations) {
+            mutations.forEach(function(mutation) {
+                if (mutation.addedNodes.length) {
+                    document.querySelectorAll('#chatbot .message:not(.processed)').forEach(msg => {
+                        msg.classList.add('processed');
+                        // Replace content with innerHTML to render HTML
+                        const content = msg.querySelector('.content');
+                        if (content) {
+                            content.innerHTML = content.textContent;
+                        }
+                    });
+                }
+            });
+        });
+        
+        // Start observing chatbot for changes
+        const chatbot = document.getElementById('chatbot');
+        if (chatbot) {
+            observer.observe(chatbot, { childList: true, subtree: true });
+        }
+
+        return [];
+    }
+    """
 
     # Event handlers
     submit_btn.click(
@@ -141,6 +219,25 @@ with gr.Blocks(css=custom_css, title="Reka Flash 3 Demo") as demo:
         message
     )
 
+    # Message submit via Enter key
+    message.submit(
+        chat_interface,
+        [message, chat_state, history_state],
+        [chat_state, history_state]
+    ).then(
+        lambda history: history,
+        chat_state,
+        chatbot
+    ).then(
+        update_history_display,
+        history_state,
+        history_display
+    ).then(
+        lambda: "",  # Clear the input box
+        None,
+        message
+    )
+
     clear_chat_btn.click(
         lambda: ([], None),
         None,
@@ -167,6 +264,11 @@ with gr.Blocks(css=custom_css, title="Reka Flash 3 Demo") as demo:
         update_history_display,
         history_state,
         history_display
+    ).then(
+        fn=load_history_from_storage,
+        inputs=None,
+        outputs=None,
+        js=js
     )
 
 demo.launch()