view thinking tokens

src/streamlit_app.py

@@ -47,11 +47,26 @@ def display_message(msg):
         with st.chat_message("assistant"):
             # Remove <think> tags if present
             content = msg['content']
+            
+            # Handle <think> tags if present
             if '<think>' in content and '</think>' in content:
                 start = content.find('<think>')
                 end = content.find('</think>') + len('</think>')
-                content = content[:start] + content[end:]
-            st.markdown(content)
+
+                # Extract thinking content and response content
+                thinking_content = content[start + len('<think>'):content.find('</think>')]
+                response_content = content[end:].strip()
+
+                # Display thinking content in a collapsed expander
+                with st.expander("🧠 Internal Thinking", expanded=False):
+                    st.code(thinking_content.strip(), language="text", wrap_lines=True)
+
+                # Display the actual response
+                if response_content:
+                    st.markdown(response_content)
+            else:
+                # No thinking tags, display content normally
+                st.markdown(content)
 
 def main():
     st.title("💬 Conversation Viewer")