Streamlit+OpenAI聊天应用点击STOP按钮后结果消失的解决方法

关键问题分析 #

原代码存在两个核心问题导致结果消失：

1. STOP按钮位置错误：放在Submit按钮的if分支内，点击STOP会触发页面重渲染，此时Submit的点击状态失效，代码跳出Submit逻辑块，无法正确维护停止状态。
2. 状态重置逻辑不当：每次点击Submit都会将result重置为空，页面重绘时若Submit条件不满足，之前的结果无法被正确恢复。

修改后的完整代码 #

import openai
import streamlit as st

# 设置OpenAI API密钥
openai.api_key = "<API-KEY>"

# 应用标题
st.subheader("AI Assistant : Streamlit + OpenAI: `stream` *argument*")

# 用户输入框
user_input = st.text_input("You: ", placeholder = "Ask me anything ...", key="input")

# 初始化会话状态变量
if 'stop_pressed' not in st.session_state:
    st.session_state['stop_pressed'] = False
if 'result' not in st.session_state:
    st.session_state['result'] = ""
if 'is_streaming' not in st.session_state:
    st.session_state['is_streaming'] = False

# 创建结果展示容器
res_box = st.empty()

# STOP按钮回调函数
def stop_streaming():
    st.session_state['stop_pressed'] = True
    st.session_state['is_streaming'] = False

# 放置STOP按钮在Submit块外，仅在流式传输时显示
if st.session_state['is_streaming']:
    st.button("STOP", key="stop", on_click=stop_streaming, type="secondary")

# Submit按钮逻辑
if st.button("Submit", type="primary"):
    # 仅当用户输入新内容时重置结果
    if user_input:
        st.session_state['result'] = ""
        st.session_state['stop_pressed'] = False
        st.session_state['is_streaming'] = True
    
    report = []
    try:
        # 发起流式聊天请求，修正消息角色为user
        response = openai.ChatCompletion.create(
            model='gpt-4',
            messages=[{"role": "user", "content": user_input}],
            temperature = 0,
            stream = True
        )
        # 遍历流式响应
        for resp in response:
            if st.session_state['stop_pressed']:
                break
            # 获取增量内容并处理
            delta_content = resp["choices"][0]["delta"].get("content", "")
            if delta_content:
                report.append(delta_content)
                current_result = "".join(report).strip().replace("\n", "")
                # 更新会话状态的结果
                st.session_state['result'] = current_result
                # 实时展示结果
                res_box.markdown(f'*{current_result}*')
    except Exception as e:
        print(e)
    finally:
        st.session_state['is_streaming'] = False
        if 'response' in locals():
            response.close()

# 始终展示会话状态中保存的结果
res_box.markdown(f'*{st.session_state["result"]}*')

# 分隔线
st.markdown("----")

核心修改点说明 #

• 调整STOP按钮位置：将STOP按钮移到Submit块外，通过is_streaming状态控制显示/隐藏，搭配回调函数修改会话状态，避免页面重绘导致的逻辑中断。
• 新增is_streaming状态：用于标记当前是否处于流式传输状态，优化按钮显示逻辑。
• 修正消息角色：原代码中messages的role设置为assistant是错误的，用户输入应对应user角色，否则OpenAI无法正确处理请求。
• 优化结果更新逻辑：每次获取增量内容后立即更新会话状态的result，确保停止后状态中保留最新的已生成内容。
• 精准重置结果：仅在用户点击Submit且有新输入时重置结果，避免页面重绘时清空已有内容。

内容的提问来源于stack exchange，提问作者henry