Update app.py

app.py

@@ -1,10 +1,10 @@
-import gradio as gr
-import os
-import json
+import streamlit as st
 import httpx
+import json
+import os
 
 # ============================================================
-# API 配置
+# 配置
 # ============================================================
 STEPFUN_API_KEY = os.environ.get("STEPFUN_API_KEY", "")
 STEPFUN_BASE_URL = "https://api.stepfun.com/v1"
@@ -12,38 +12,36 @@ MODEL_NAME = "step-3.5-flash"
 HF_CONFIG_URL = "https://huggingface.co/stepfun-ai/Step-3.5-Flash/raw/main/config.json"
 STEPFUN_LOGO = "https://huggingface.co/stepfun-ai/Step-3.5-Flash/resolve/main/stepfun.svg"
 
-cached_config = None
+st.set_page_config(
+    page_title="Step-3.5-Flash",
+    page_icon="🚀",
+    layout="wide",
+)
 
 
+@st.cache_data(ttl=3600)
 def fetch_model_config():
-    global cached_config
     try:
         response = httpx.get(HF_CONFIG_URL, timeout=10.0)
         if response.status_code == 200:
-            cached_config = response.json()
-            return cached_config
+            return response.json()
     except Exception as e:
-        print(f"拉取 config.json 失败: {e}")
-    return cached_config
+        st.error(f"拉取 config.json 失败: {e}")
+    return None
 
 
 def format_messages(history, system_prompt: str, user_message: str):
-    """将 chatbot history 转换为 API 消息格式"""
     messages = []
     if system_prompt.strip():
         messages.append({"role": "system", "content": system_prompt})
-    for user_msg, bot_msg in history:
-        if user_msg:
-            messages.append({"role": "user", "content": user_msg})
-        if bot_msg:
-            messages.append({"role": "assistant", "content": bot_msg})
+    for msg in history:
+        messages.append({"role": msg["role"], "content": msg["content"]})
     messages.append({"role": "user", "content": user_message})
     return messages
 
 
-def chat_stream(message: str, history, system_prompt: str, max_tokens: int, temperature: float, top_p: float):
+def chat_stream(message: str, history: list, system_prompt: str, max_tokens: int, temperature: float, top_p: float):
     """流式聊天，返回 (reasoning, content) 生成器"""
-    fetch_model_config()
     messages = format_messages(history, system_prompt, message)
 
     reasoning = ""
@@ -90,92 +88,116 @@ def chat_stream(message: str, history, system_prompt: str, max_tokens: int, temp
         yield reasoning, f"❌ 错误: {str(e)}"
 
 
-def create_demo():
+def main():
+    st.title("🚀 Step-3.5-Flash")
+    st.caption("基于 [Step-3.5-Flash](https://huggingface.co/stepfun-ai/Step-3.5-Flash) 的智能对话助手")
+
+    # 初始化 session state
+    if "messages" not in st.session_state:
+        st.session_state.messages = []
+    if "thinking" not in st.session_state:
+        st.session_state.thinking = ""
+
+    # 侧边栏设置
+    with st.sidebar:
+        st.header("⚙️ 设置")
+        system_prompt = st.text_area("系统提示词", value="你是一个有帮助的 AI 助手。", height=80)
+        max_tokens = st.slider("最大长度", 256, 131072, 4096, step=256, help="最大 128k")
+        temperature = st.slider("Temperature", 0.0, 1.5, 0.7, step=0.1)
+        top_p = st.slider("Top-p", 0.1, 1.0, 0.9, step=0.05)
+
+        st.divider()
+        if st.button("🗑️ 清空对话", use_container_width=True):
+            st.session_state.messages = []
+            st.session_state.thinking = ""
+            st.rerun()
+
+        st.divider()
+        with st.expander("📋 模型配置"):
+            config = fetch_model_config()
+            if config:
+                st.json(config)
+
+    # 主界面：左右布局
+    col_thinking, col_chat = st.columns([1, 4])
+
+    # 左侧：思考过程
+    with col_thinking:
+        st.subheader("💭 思考过程")
+        thinking_container = st.container(height=500)
+        with thinking_container:
+            if st.session_state.thinking:
+                st.markdown(st.session_state.thinking)
+            else:
+                st.caption("等待输入...")
+
+    # 右侧：对话
+    with col_chat:
+        st.subheader("💬 对话")
+
+        # 显示历史消息
+        chat_container = st.container(height=450)
+        with chat_container:
+            for msg in st.session_state.messages:
+                with st.chat_message(msg["role"], avatar=STEPFUN_LOGO if msg["role"] == "assistant" else None):
+                    st.markdown(msg["content"])
+
+        # 输入框
+        if prompt := st.chat_input("输入消息..."):
+            # 添加用户消息
+            st.session_state.messages.append({"role": "user", "content": prompt})
+
+            # 显示用户消息
+            with chat_container:
+                with st.chat_message("user"):
+                    st.markdown(prompt)
+
+            # 生成回复
+            with chat_container:
+                with st.chat_message("assistant", avatar=STEPFUN_LOGO):
+                    response_placeholder = st.empty()
+                    thinking_placeholder = col_thinking.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 "▌"
+                        response_placeholder.markdown(full_response)
+
+                        # 更新思考过程
+                        with thinking_placeholder.container(height=500):
+                            if full_thinking:
+                                st.markdown(full_thinking)
+
+                    # 保存消息
+                    st.session_state.messages.append({"role": "assistant", "content": full_response})
+                    st.session_state.thinking = full_thinking
+                    st.rerun()
+
+    # 示例问题
+    st.divider()
+    st.subheader("💡 试试这些")
     examples = [
         "请解释一下什么是机器学习？",
         "帮我写一个 Python 快速排序算法",
         "1000以内有多少个质数？",
         "一个农夫需要把狼、羊和白菜都带过河，请问农夫该怎么办？",
     ]
-
-    with gr.Blocks(title="Step-3.5-Flash", theme=gr.themes.Soft()) as demo:
-        gr.Markdown("# 🚀 Step-3.5-Flash")
-
-        with gr.Row():
-            # 左侧：思考过程 (1)
-            with gr.Column(scale=1, min_width=250):
-                gr.Markdown("### 💭 思考过程")
-                thinking_display = gr.Textbox(
-                    value="等待输入...",
-                    lines=20,
-                    max_lines=20,
-                    interactive=False,
-                    show_label=False,
-                )
-
-            # 右侧：对话 (4)
-            with gr.Column(scale=4):
-                gr.Markdown("### 💬 对话")
-                chatbot = gr.Chatbot(
-                    height=450,
-                    show_label=False,
-                    avatar_images=(None, STEPFUN_LOGO),
-                )
-
-                with gr.Row():
-                    msg = gr.Textbox(
-                        placeholder="输入消息...",
-                        show_label=False,
-                        scale=8,
-                        container=False,
-                    )
-                    submit_btn = gr.Button("发送", variant="primary", scale=1)
-                    clear_btn = gr.Button("🗑️", scale=0, min_width=50)
-
-        # 设置（折叠）
-        with gr.Accordion("⚙️ 设置", open=False):
-            with gr.Row():
-                system_prompt = gr.Textbox(label="系统提示词", value="你是一个有帮助的 AI 助手。", scale=2)
-                max_tokens = gr.Slider(256, 131072, value=4096, step=256, label="最大长度", scale=1)
-                temperature = gr.Slider(0.0, 1.5, value=0.7, step=0.1, label="Temperature", scale=1)
-                top_p = gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="Top-p", scale=1)
-
-        gr.Examples(examples, inputs=msg, label="💡 试试这些")
-
-        # 事件处理
-        def respond(message, history, system_prompt, max_tokens, temperature, top_p):
-            if not message.strip():
-                yield history, "", ""
-                return
-
-            # 添加用户消息
-            history = history + [[message, None]]
-            yield history, "", "思考中..."
-
-            reasoning = ""
-            content = ""
-
-            for r, c in chat_stream(message, history[:-1], system_prompt, max_tokens, temperature, top_p):
-                reasoning = r if r else ""
-                content = c if c else "▌"
-                history[-1][1] = content
-                yield history, "", reasoning
-
-            history[-1][1] = content
-            yield history, "", reasoning
-
-        def on_clear():
-            return [], "", "等待输入..."
-
-        msg.submit(respond, [msg, chatbot, system_prompt, max_tokens, temperature, top_p], [chatbot, msg, thinking_display])
-        submit_btn.click(respond, [msg, chatbot, system_prompt, max_tokens, temperature, top_p], [chatbot, msg, thinking_display])
-        clear_btn.click(on_clear, outputs=[chatbot, msg, thinking_display])
-        demo.load(fetch_model_config)
-
-    return demo
+    cols = st.columns(len(examples))
+    for i, example in enumerate(examples):
+        if cols[i].button(example, key=f"example_{i}", use_container_width=True):
+            st.session_state.messages.append({"role": "user", "content": example})
+            st.rerun()
 
 
 if __name__ == "__main__":
-    demo = create_demo()
-    demo.queue()
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+    main()