如何通过Flask API实现LangChain OpenAI的响应流式传输？

当前代码的核心问题在于askQuestion方法是同步调用链的完整执行逻辑，等待所有内容生成完成后才返回结果。虽然设置了streaming=True和StreamingStdOutCallbackHandler，但该回调仅将输出打印到控制台，并未将流式生成的token传递给Flask的响应生成器，导致接口只能一次性返回全部内容。

步骤1：自定义回调捕获流式Token #

创建回调处理器，将LLM生成的每个token存入队列，方便后续生成器读取：

from langchain.callbacks.base import BaseCallbackHandler
import queue

class QueueCallback(BaseCallbackHandler):
    def __init__(self):
        self.token_queue = queue.Queue()

    def on_llm_new_token(self, token: str, **kwargs) -> None:
        self.token_queue.put(token)

    def on_llm_end(self, response, **kwargs) -> None:
        self.token_queue.put(None)  # 标记生成结束

步骤2：修改askQuestion方法，拆分检索与流式生成 #

先检索源文档，再单独启动流式LLM生成回答，同时保留对话历史：

def askQuestion(self, collection_id, question):
    collection_name = "collection-" + str(collection_id)
    
    # 1. 检索相关文档
    chroma_Vectorstore = Chroma(collection_name=collection_name, embedding_function=self.embeddingsOpenAi, client=self.chroma_client)
    retriever = chroma_Vectorstore.as_retriever(similarity_search_with_score=True)
    source_documents = retriever.get_relevant_documents(question)
    
    # 整理源文档数据
    source_data = []
    for source in source_documents:
        source_data.append({
            "page_content": source.page_content,
            "metadata": source.metadata
        })
    
    # 2. 流式生成回答
    def generate_answer_tokens():
        # 初始化回调与流式LLM
        callback = QueueCallback()
        llm = ChatOpenAI(
            model_name=self.model_name,
            temperature=self.temperature,
            openai_api_key=os.environ.get('OPENAI_API_KEY'),
            streaming=True,
            callback_manager=CallbackManager([callback])
        )
        
        # 构建提示模板，包含对话历史与文档上下文
        from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder
        prompt = ChatPromptTemplate.from_messages([
            ("system", "根据提供的文档内容回答用户问题：\n{context}"),
            MessagesPlaceholder(variable_name="chat_history"),
            ("human", "{question}")
        ])
        
        # 加载对话历史
        chat_history = self.memory.load_memory_variables({})["chat_history"]
        context = "\n".join([doc.page_content for doc in source_documents])
        
        # 触发LLM生成（异步触发回调）
        llm.predict_messages(
            prompt.format_messages(context=context, chat_history=chat_history, question=question)
        )
        
        # 从队列读取token并返回
        while True:
            token = callback.token_queue.get()
            if token is None:
                break
            yield token
    
    # 保存对话历史（生成完成后执行）
    def update_memory():
        full_answer = ''.join(generate_answer_tokens())
        self.memory.save_context({"question": question}, {"answer": full_answer})
    
    # 返回生成器与源文档数据
    return {
        "answer_tokens": generate_answer_tokens(),
        "source_documents": source_data,
        "update_memory": update_memory
    }

步骤3：修改Flask路由，实现流式响应 #

调整路由逻辑，分块返回源文档与流式回答，设置正确的响应头：

import json
from flask import stream_with_context

@app.route("/collection/<int:collection_id>/ask_question", methods=["POST"])
def ask_question(collection_id):
    question = request.form["question"]
    result = document_thread.askQuestion(collection_id, question)
    source_docs = result["source_documents"]
    answer_tokens = result["answer_tokens"]
    update_memory = result["update_memory"]
    
    def stream_response():
        # 先发送源文档信息（JSON格式）
        yield '{"source_documents": ' + json.dumps(source_docs) + ', "answer": "'
        
        # 流式发送每个token，转义特殊字符避免JSON错误
        for token in answer_tokens:
            escaped_token = token.replace('"', '\\"').replace('\n', '\\n').replace('\r', '\\r')
            yield escaped_token
        
        # 结束JSON结构
        yield '"}'
        
        # 更新对话历史
        update_memory()
    
    # 设置响应头，启用分块传输
    return app.response_class(
        stream_with_context(stream_response()),
        mimetype="application/json",
        headers={
            "Transfer-Encoding": "chunked",
            "Cache-Control": "no-cache",
            "Connection": "keep-alive"
        }
    )

关键说明 #

• 拆分检索与生成：避免ConversationalRetrievalChain的同步执行逻辑，单独控制流式生成的token输出。
• 回调与队列：通过QueueCallback将LLM生成的token传递给响应生成器，实现流式转发。
• 响应头设置：Transfer-Encoding: chunked告知客户端响应为分块传输，Cache-Control: no-cache避免客户端缓存内容。

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