基于Langchain的流式摘要问题：仅需最终结果却输出各阶段内容

我用以下Python代码通过LangChain的自定义map-reduce方法实现文档摘要功能。我手里有一份大文档的分块列表，只需要流式输出最终的摘要结果，但现在map操作和reduce操作阶段都会流式输出每个分块的摘要内容，怎么解决？

model_llm = ChatOpenAI(temperature=0.0, model='gpt-4')
model_llm.streaming = True
docs = [Document(page_content=txt) for txt in list_chunks]
chunk_attrib = 'text'
# replacing attributes in map prompt other than document chunk
map_prompt = "Summarize document in bullet points in {text}"
combine_prompt = "Remove redundancy in summary in {text}"

human_message_prompt = HumanMessagePromptTemplate.from_template(combine_prompt)
combine_prompt_template = ChatPromptTemplate.from_messages(
    [system_message_prompt, human_message_prompt],
)
human_message_prompt_reduce = HumanMessagePromptTemplate.from_template(map_prompt)
map_prompt_template = ChatPromptTemplate.from_messages(
    [system_message_prompt, human_message_prompt],
)
chain = load_summarize_chain(self.model_llm,
                             chain_type="map_reduce",
                             map_prompt=map_prompt_template,
                             combine_prompt=map_prompt_template,
                             combine_document_variable_name=chunk_attrib,
                             map_reduce_document_variable_name=chunk_attrib,
                             input_key=chunk_attrib,
                             output_key=chunk_attrib
                             )
output = chain.run(docs)

解决方案 #

核心问题 #

你给单个LLM实例设置了streaming=True，导致map和combine（reduce）两个阶段的LLM调用都会触发流式输出。要实现只输出最终摘要，需要让只有combine阶段启用流式，map阶段关闭。

方法一：为不同阶段分配不同LLM实例 #

创建两个LLM实例，map阶段用非流式版本，combine阶段用流式版本，从根源上分开输出逻辑：

from langchain.chat_models import ChatOpenAI
from langchain.prompts import ChatPromptTemplate, HumanMessagePromptTemplate
from langchain.schema import SystemMessage
from langchain.chains.summarize import load_summarize_chain
from langchain.docstore.document import Document

# 定义系统提示词（匹配你原代码的system_message_prompt）
system_message_prompt = SystemMessage(content="You are a helpful assistant that summarizes documents accurately.")

# 分别创建map和combine阶段的LLM
map_llm = ChatOpenAI(temperature=0.0, model='gpt-4', streaming=False)
combine_llm = ChatOpenAI(temperature=0.0, model='gpt-4', streaming=True)

docs = [Document(page_content=txt) for txt in list_chunks]
chunk_attrib = 'text'

# 修正提示词模板，避免变量混淆
map_prompt = "Summarize the following document content in bullet points: {text}"
combine_prompt = "Remove redundancy from the following summaries and generate a coherent final summary: {text}"

# 构建map阶段提示模板
map_human_prompt = HumanMessagePromptTemplate.from_template(map_prompt)
map_prompt_template = ChatPromptTemplate.from_messages(
    [system_message_prompt, map_human_prompt]
)

# 构建combine阶段提示模板
combine_human_prompt = HumanMessagePromptTemplate.from_template(combine_prompt)
combine_prompt_template = ChatPromptTemplate.from_messages(
    [system_message_prompt, combine_human_prompt]
)

# 加载汇总链，指定不同阶段的LLM
chain = load_summarize_chain(
    llm=map_llm,
    chain_type="map_reduce",
    map_prompt=map_prompt_template,
    combine_prompt=combine_prompt_template,
    combine_llm=combine_llm,  # 为combine阶段单独配置流式LLM
    combine_document_variable_name=chunk_attrib,
    map_reduce_document_variable_name=chunk_attrib,
    input_key="input_documents",  # 修正输入key，load_summarize_chain默认用这个
    return_intermediate_steps=False  # 不返回中间分块摘要
)

# 流式获取最终结果
for chunk in chain.stream({"input_documents": docs}):
    print(chunk[chunk_attrib], end="", flush=True)

方法二：自定义回调过滤输出 #

如果不想创建多个LLM实例，可通过自定义回调函数，只捕获combine阶段的流式token：

from langchain.chat_models import ChatOpenAI
from langchain.prompts import ChatPromptTemplate, HumanMessagePromptTemplate
from langchain.schema import SystemMessage
from langchain.chains.summarize import load_summarize_chain
from langchain.docstore.document import Document
from langchain.callbacks.base import BaseCallbackHandler

# 自定义回调处理器，只输出combine阶段的token
class FinalSummaryStreamHandler(BaseCallbackHandler):
    def __init__(self, combine_keyword):
        self.is_combine_stage = False
        self.combine_keyword = combine_keyword

    def on_llm_start(self, serialized, prompts, **kwargs):
        # 通过提示词判断是否进入combine阶段
        self.is_combine_stage = any(self.combine_keyword in prompt for prompt in prompts)

    def on_llm_new_token(self, token, **kwargs):
        # 仅在combine阶段输出token
        if self.is_combine_stage:
            print(token, end="", flush=True)

# 初始化系统提示词
system_message_prompt = SystemMessage(content="You are a helpful assistant that summarizes documents accurately.")

# 初始化带自定义回调的流式LLM
combine_keyword = "Remove redundancy"
model_llm = ChatOpenAI(
    temperature=0.0,
    model='gpt-4',
    streaming=True,
    callbacks=[FinalSummaryStreamHandler(combine_keyword)]
)

docs = [Document(page_content=txt) for txt in list_chunks]
chunk_attrib = 'text'

map_prompt = "Summarize the following document content in bullet points: {text}"
combine_prompt = "Remove redundancy from the following summaries and generate a coherent final summary: {text}"

# 修正之前的模板错误：map阶段用对应提示词
map_human_prompt = HumanMessagePromptTemplate.from_template(map_prompt)
map_prompt_template = ChatPromptTemplate.from_messages(
    [system_message_prompt, map_human_prompt]
)

combine_human_prompt = HumanMessagePromptTemplate.from_template(combine_prompt)
combine_prompt_template = ChatPromptTemplate.from_messages(
    [system_message_prompt, combine_human_prompt]
)

chain = load_summarize_chain(
    model_llm,
    chain_type="map_reduce",
    map_prompt=map_prompt_template,
    combine_prompt=combine_prompt_template,
    combine_document_variable_name=chunk_attrib,
    map_reduce_document_variable_name=chunk_attrib,
    input_key="input_documents"
)

# 执行链，仅最终摘要会流式输出
chain.run(docs)

额外修正点 #

你原代码里有两处错误：

1. map_prompt_template误用了human_message_prompt（combine阶段的提示模板），导致map阶段用了错误的提示词，需要修正为对应map阶段的模板。
2. load_summarize_chain的默认输入key是input_documents，不是你设置的chunk_attrib，之前的配置会导致参数不匹配，需要修正。

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