如何在Streamlit应用中实现本地Ollama大模型的流式响应？

问题核心：st.write_stream要求输入为生成器或类流对象，但当前llm_chain.run()直接返回完整字符串，且原回调仅输出到控制台，无法为Streamlit提供流式数据。

步骤1：自定义Streamlit流式回调处理器 #

创建回调类，收集模型输出的每个token，并提供生成器接口供st.write_stream调用：

from langchain.callbacks.base import BaseCallbackHandler

class StreamlitStreamingCallbackHandler(BaseCallbackHandler):
    def __init__(self):
        self.tokens = []
        self.finished = False

    def on_llm_new_token(self, token: str, **kwargs) -> None:
        self.tokens.append(token)

    def on_llm_end(self, response, **kwargs) -> None:
        self.finished = True

    def stream(self):
        while not self.finished or self.tokens:
            if self.tokens:
                yield self.tokens.pop(0)

步骤2：替换原回调并初始化LLM #

用自定义回调替换原StreamingStdOutCallbackHandler，同时开启Ollama的流式输出模式：

# 替换原callback_manager初始化逻辑
callback_handler = StreamlitStreamingCallbackHandler()
callback_manager = CallbackManager([callback_handler])
llm = Ollama(model=st.session_state.selected_model, callbacks=callback_manager, streaming=True)

步骤3：修改聊天响应逻辑 #

在助理消息块中，通过回调生成器向st.write_stream喂数据，同时收集完整响应存入会话历史：

if prompt := st.chat_input(""):
    with st.chat_message(name="user", avatar="source/ai_user_2.png"):
        st.write(prompt)
        st.session_state.messages.append({"role": "user", "content": prompt})
    
    with st.chat_message(name="assistant", avatar="source/ai_bot_2.png"):
        # 重新初始化回调处理器，避免跨请求数据污染
        callback_handler = StreamlitStreamingCallbackHandler()
        # 更新LLM链的回调管理器
        st.session_state.llm_chain.llm.callbacks = CallbackManager([callback_handler])
        
        # 异步运行模型，避免阻塞UI
        import threading
        def run_model():
            st.session_state.llm_chain.run(input=prompt)
        
        threading.Thread(target=run_model).start()
        
        # 流式输出并收集完整响应
        full_response = ""
        for token in callback_handler.stream():
            full_response += token
            yield token
        
        # 将完整响应存入会话历史
        st.session_state.messages.append({"role": "assistant", "content": full_response})

额外修复：会话历史渲染的语法错误 #

原代码中elif缩进错误且使用is进行字符串判断，修正后：

for message in st.session_state.messages:
    if message["role"] == "user":
        with st.chat_message(name="user", avatar="source/ai_user_2.png"):
            st.write(message["content"])
    elif message["role"] == "assistant":
        with st.chat_message(name="assistant", avatar="source/ai_bot_2.png"):
            st.write(message["content"])

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