Update app.py

app.py

@@ -2,6 +2,7 @@ import streamlit as st
 import httpx
 import json
 import os
+import re
 
 # ============================================================
 # 配置
@@ -17,94 +18,35 @@ st.set_page_config(
     layout="centered",
 )
 
-# 自定义样式
+# 简化样式 - 只定义思考区域
 st.markdown("""
 <style>
-/* 聊天容器 */
-.chat-container {
-    max-width: 800px;
-    margin: 0 auto;
-    padding-bottom: 100px;
-}
-
-/* 用户消息 - 右侧 */
-.user-message {
-    display: flex;
-    justify-content: flex-end;
-    margin-bottom: 20px;
-}
-.user-bubble {
-    background: linear-gradient(135deg, #3b82f6 0%, #2563eb 100%);
-    color: white;
-    padding: 12px 16px;
-    border-radius: 18px 18px 4px 18px;
-    max-width: 70%;
-    font-size: 15px;
-    line-height: 1.5;
-    word-break: break-word;
-}
-
-/* 助手消息 - 左侧 */
-.assistant-message {
-    display: flex;
-    flex-direction: column;
-    align-items: flex-start;
-    margin-bottom: 20px;
-}
-.assistant-label {
-    font-size: 12px;
-    color: #64748b;
-    margin-bottom: 6px;
-    font-weight: 500;
-}
+#MainMenu {visibility: hidden;}
+footer {visibility: hidden;}
 
-/* 思考气泡 - 可滚动 */
-.thinking-bubble {
-    background: linear-gradient(135deg, #f1f5f9 0%, #e2e8f0 100%);
-    border: 1px solid #cbd5e1;
-    border-radius: 12px;
-    padding: 12px 16px;
-    margin-bottom: 12px;
-    max-width: 85%;
+.thinking-container {
+    background: #f8fafc;
+    border: 1px solid #e2e8f0;
+    border-radius: 8px;
+    padding: 10px 14px;
+    margin-bottom: 10px;
     max-height: 150px;
     overflow-y: auto;
     font-size: 13px;
-    line-height: 1.6;
+    line-height: 1.5;
     color: #64748b;
-    white-space: pre-wrap;
-    word-break: break-word;
 }
-.thinking-bubble::-webkit-scrollbar {
+.thinking-container::-webkit-scrollbar {
     width: 4px;
 }
-.thinking-bubble::-webkit-scrollbar-thumb {
-    background: #94a3b8;
+.thinking-container::-webkit-scrollbar-thumb {
+    background: #cbd5e1;
     border-radius: 2px;
 }
-.thinking-header {
-    font-size: 11px;
+.thinking-label {
+    font-size: 12px;
     color: #94a3b8;
     margin-bottom: 4px;
-    font-weight: 500;
-}
-
-/* 回答文本 - 直接显示 */
-.answer-text {
-    max-width: 85%;
-    font-size: 15px;
-    line-height: 1.7;
-    color: #1e293b;
-    white-space: pre-wrap;
-    word-break: break-word;
-}
-
-/* 隐藏 Streamlit 默认元素 */
-#MainMenu {visibility: hidden;}
-footer {visibility: hidden;}
-
-/* 自动滚动脚本的锚点 */
-.scroll-anchor {
-    height: 1px;
 }
 </style>
 """, unsafe_allow_html=True)
@@ -182,57 +124,21 @@ def chat_stream(message: str, history: list, system_prompt: str, max_tokens: int
         yield reasoning, f"❌ 错误: {str(e)}"
 
 
-def render_user_message(content: str):
-    """渲染用户消息（右侧）"""
-    escaped = content.replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;")
-    st.markdown(f"""
-    <div class="user-message">
-        <div class="user-bubble">{escaped}</div>
-    </div>
-    """, unsafe_allow_html=True)
-
-
-def render_assistant_message(thinking: str, answer: str, is_streaming: bool = False):
-    """渲染助手消息（左侧）"""
-    thinking_html = ""
-    if thinking:
-        # 去掉 <think> 标签
-        cleaned_thinking = thinking.replace("<think>", "").replace("</think>", "")
-        escaped_thinking = cleaned_thinking.replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;")
-        scroll_script = """<script>
-            (function() {
-                var el = document.getElementById('thinking-stream');
-                if (el) el.scrollTop = el.scrollHeight;
-            })();
-        </script>""" if is_streaming else ""
-        thinking_html = f"""
-        <div class="thinking-header">💭 思考过程</div>
-        <div class="thinking-bubble" id="thinking-stream">{escaped_thinking}</div>
-        {scroll_script}
-        """
-
-    answer_html = ""
-    if answer:
-        escaped_answer = answer.replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;")
-        answer_html = f'<div class="answer-text">{escaped_answer}</div>'
-
-    st.markdown(f"""
-    <div class="assistant-message">
-        <div class="assistant-label">🚀 Step-3.5-Flash</div>
-        {thinking_html}
-        {answer_html}
-    </div>
-    """, unsafe_allow_html=True)
-
-
-def scroll_to_bottom():
-    """滚动到页面底部"""
-    st.markdown("""
-    <script>
-        window.scrollTo(0, document.body.scrollHeight);
-    </script>
-    <div class="scroll-anchor"></div>
-    """, unsafe_allow_html=True)
+def clean_thinking(text: str) -> str:
+    """清理思考内容中的标签"""
+    if not text:
+        return ""
+    # 移除 <think> 标签
+    text = re.sub(r'</?think>', '', text)
+    return text.strip()
+
+
+def render_thinking_expander(thinking_text: str, is_streaming: bool = False):
+    """使用 expander 渲染思考内容"""
+    if thinking_text:
+        cleaned = clean_thinking(thinking_text)
+        with st.expander("💭 思考过程", expanded=is_streaming):
+            st.text(cleaned)
 
 
 def main():
@@ -267,9 +173,15 @@ def main():
     # 显示历史消息
     for msg in st.session_state.messages:
         if msg["role"] == "user":
-            render_user_message(msg["content"])
+            with st.chat_message("user"):
+                st.markdown(msg["content"])
         elif msg["role"] == "assistant":
-            render_assistant_message(msg.get("thinking", ""), msg.get("content", ""))
+            with st.chat_message("assistant", avatar="🚀"):
+                # 思考内容用 expander
+                if msg.get("thinking"):
+                    render_thinking_expander(msg["thinking"], is_streaming=False)
+                # 回答内容用 markdown
+                st.markdown(msg.get("content", ""))
 
     # 示例问题（无消息时显示）
     if not st.session_state.messages:
@@ -296,39 +208,45 @@ def main():
     if prompt:
         # 添加并显示用户消息
         st.session_state.messages.append({"role": "user", "content": prompt})
-        render_user_message(prompt)
-
-        # 创建流式输出占位符
-        response_placeholder = st.empty()
-
-        full_response = ""
-        full_thinking = ""
-
-        for thinking, response in chat_stream(
-            prompt,
-            st.session_state.messages[:-1],
-            system_prompt,
-            max_tokens,
-            temperature,
-            top_p,
-        ):
-            full_thinking = thinking
-            full_response = response if response else "▌"
-
-            # 更新显示
-            with response_placeholder:
-                render_assistant_message(full_thinking, full_response, is_streaming=True)
-
-            # 滚动到底部
-            scroll_to_bottom()
-
-        # 保存消息
-        st.session_state.messages.append({
-            "role": "assistant",
-            "content": full_response,
-            "thinking": full_thinking,
-        })
-        st.rerun()
+        with st.chat_message("user"):
+            st.markdown(prompt)
+
+        # 助手回复
+        with st.chat_message("assistant", avatar="🚀"):
+            # 思考内容占位符
+            thinking_placeholder = st.empty()
+            # 回答内容占位符
+            answer_placeholder = st.empty()
+
+            full_response = ""
+            full_thinking = ""
+
+            for thinking, response in chat_stream(
+                prompt,
+                st.session_state.messages[:-1],
+                system_prompt,
+                max_tokens,
+                temperature,
+                top_p,
+            ):
+                full_thinking = thinking
+                full_response = response if response else "▌"
+
+                # 更新思考内容
+                if full_thinking:
+                    with thinking_placeholder.container():
+                        render_thinking_expander(full_thinking, is_streaming=True)
+
+                # 更新回答内容
+                answer_placeholder.markdown(full_response)
+
+            # 保存消息
+            st.session_state.messages.append({
+                "role": "assistant",
+                "content": full_response,
+                "thinking": full_thinking,
+            })
+            st.rerun()
 
 
 if __name__ == "__main__":