RetrievalQA流式输出在Chainlit UI失效问题求助

问题背景 #

使用Langchain的RetrievalQA链整合LLM、Prompt、向量存储与ConversationBufferMemory时，LLMChain在CLI和Chainlit UI中均可正常流式输出，但RetrievalQA仅在CLI生效，无法在Chainlit UI实现Token流式输出。

核心问题分析 #

1. LLM未开启流式配置：代码中load_llm()里的streaming=True被注释，导致LLM本身不支持流式输出。
2. 回调处理器冲突：同时使用了AsyncFinalIteratorCallbackHandler和Chainlit的回调，导致回调信号被拦截，无法传递到UI。
3. RetrievalQA的回调传递不足：未确保自定义回调被正确传递到内部的LLMChain层级。
4. 内存与Prompt格式不匹配：ConversationBufferMemory的return_messages=True返回消息对象，但Prompt中的{chat_history}需要字符串格式，导致渲染异常。

解决方案步骤 #

• 开启LLM流式模式：在CTransformers配置中添加streaming=True，移除冲突的AsyncFinalIteratorCallbackHandler。
• 使用Chainlit专属回调处理器：自定义StreamHandler监听on_llm_new_token事件，确保Token能被实时推送到UI。
• 确保回调传递到LLM：调用chain.acall()时传入回调处理器，让RetrievalQA内部的LLM触发流式回调。
• 调整内存配置匹配Prompt：将ConversationBufferMemory的return_messages设为False，确保chat_history以字符串格式渲染。

修改后的完整代码 #

from langchain.chains import LLMChain, RetrievalQA
from langchain.prompts import PromptTemplate
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.vectorstores import FAISS
from langchain.llms import CTransformers
import chainlit as cl
from langchain.memory import ConversationBufferMemory
from langchain.callbacks.base import BaseCallbackHandler

DB_FAISS_PATH = 'vectorstores/db_faiss'

# RetrievalQA使用的Prompt模板
custom_prompt_template="""
Below is an instruction that describes a task. Write a response that appropriately completes the request from the given context.
### Chat History:
{chat_history}
## Context:
{context}
### Instruction:
{question}
### Response:
"""

def set_custom_prompt():
    prompt = PromptTemplate(
        template=custom_prompt_template,
        input_variables=["chat_history","context","question"]
    )
    return prompt

# 加载LLM，开启流式模式
def load_llm():
    llm = CTransformers(
        model="H:\WPCS_OPS_MERGED\OPS_WPCS_Q4.gguf", 
        model_type="llama",
        streaming=True,  # 必须开启流式
        config={
            "temperature":0.15,
            "context_length":1800,
            "top_k":30,
            "repetition_penalty":1.2
        }
    )
    return llm

# 创建RetrievalQA链
def create_chain_qa():
    embeddings = HuggingFaceEmbeddings(
        model_name="sentence-transformers/all-MiniLM-L6-v2",
        model_kwargs={'device': 'cpu'}
    )
    db = FAISS.load_local(DB_FAISS_PATH, embeddings)
    # 内存配置调整：return_messages=False，确保chat_history是字符串格式
    memory = ConversationBufferMemory(
        memory_key="chat_history",
        input_key="question",
        max_token_limit=150,
        return_messages=False
    )
    qa_chain = RetrievalQA.from_chain_type(
        llm=load_llm(),
        chain_type='stuff',
        retriever=db.as_retriever(search_kwargs={'k': 3}),
        verbose=True,
        return_source_documents=False,
        chain_type_kwargs={
            'prompt': set_custom_prompt(),
            "verbose":True,
            "memory": memory
        }
    )
    return qa_chain

# 自定义流式回调处理器
class StreamHandler(BaseCallbackHandler):
    def __init__(self):
        self.msg = cl.Message(content="")

    async def on_llm_new_token(self, token: str, **kwargs):
        await self.msg.stream_token(token)

    async def on_llm_end(self, response, **kwargs):
        await self.msg.send()

@cl.on_chat_start
async def start():
    chain = create_chain_qa()
    msg = cl.Message(content="Starting the bot...")
    await msg.send()
    msg.content = "Hi, Welcome to AHS Bot. What is your query?"
    await msg.update()
    cl.user_session.set("chain", chain)

@cl.on_message
async def main(message):
    chain = cl.user_session.get("chain") 
    # 初始化自定义流式处理器
    stream_handler = StreamHandler()
    # 调用链时传入回调
    res = await chain.acall(
        message, 
        callbacks=[stream_handler]
    )

关键修改说明 #

1. LLM流式开启：load_llm()中添加streaming=True，移除了AsyncFinalIteratorCallbackHandler，避免回调冲突。
2. 内存配置调整：ConversationBufferMemory设置return_messages=False，确保{chat_history}能正确渲染为字符串，匹配Prompt模板。
3. 自定义回调生效：使用StreamHandler监听on_llm_new_token事件，每收到一个Token就推送到Chainlit UI，实现流式输出。
4. 回调传递正确：调用chain.acall()时传入stream_handler，确保RetrievalQA内部的LLMChain能触发流式回调。

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