使用LangChain流式输出RAG响应时CPU占用率飙升至100%的优化求助

问题现象 #

使用LangChain构建RAG系统，基于OpenAI API生成响应时，非流式输出场景CPU占用率正常，但启用流式输出后CPU会骤升至100%，具体代码及运行数据如下：

非流式输出代码 #

chat_history = []

question = input("Enter")
log_cpu_usage("Before Processing")
ai_msg = rag_chain.invoke({"question": question, "chat_history": chat_history})
print(ai_msg)
log_cpu_usage("After Processing")
chat_history.extend([HumanMessage(content = question), AIMessage(content=ai_msg)])

非流式输出结果 #

[CPU Usage] Before Processing: 14.3%
User: Hi
Hello Faiq! How can I assist you today?
[CPU Usage] After Processing: 16.8%
[CPU Usage] Before Processing: 13.6%
User: How are you?
Thank you for asking, Faiq! As an AI, I don't have feelings, but I'm here and ready to help you with any questions or concerns you may have. How can I assist you today?
[CPU Usage] After Processing: 2.5%
[CPU Usage] Before Processing: 11.4%

流式输出代码 #

chat_history = []

question = input("Enter")
answer = ""
for chunk in rag_chain.stream({"question":question, "chat_history": chat_history}):
    (chunk, end="", flush=True)  # 存在语法错误，应为print语句
    log_cpu_usage(":")
    answer +=chunk
    chat_history.extend([HumanMessage(content = question), AIMessage(content=answer)])

流式输出结果 #

[CPU Usage] CPU Usage:: 13.5%
User: Hi
[CPU Usage] :: 10.4%
Hello[CPU Usage] :: 0.0%
 F[CPU Usage] :: 100.0%
ai[CPU Usage] :: 100.0%
q[CPU Usage] :: 16.7%
![CPU Usage] :: 100.0%
 How[CPU Usage] :: 100.0%
 can[CPU Usage] :: 4.8%
 I[CPU Usage] :: 100.0%
 assist[CPU Usage] :: 0.0%
 you[CPU Usage] :: 16.7%
 today[CPU Usage] :: 100.0%
?[CPU Usage] :: 100.0%
[CPU Usage] :: 0.0%

优化方案 #

1. 移除循环内无效的chat_history更新 #

流式循环中逐chunk更新chat_history是完全冗余且错误的操作，会导致chat_history快速膨胀，且每次循环都创建HumanMessage和AIMessage对象，带来额外CPU开销。正确做法是等待整个回答生成完成后，再一次性更新chat_history。

2. 修复print语句语法错误 #

原代码中(chunk, end="", flush=True)是无效表达式，应改为print(chunk, end="", flush=True)。若对实时性要求不是极端苛刻，可减少flush频率（比如每3-5个chunk执行一次flush），降低IO操作带来的CPU消耗。

3. 降低CPU采样频率 #

每个chunk都调用log_cpu_usage会触发频繁的系统调用，额外消耗CPU资源。可设置时间间隔采样，比如每1秒记录一次：

import time

last_log_time = time.time()
log_interval = 1  # 采样间隔1秒

for chunk in rag_chain.stream({"question":question, "chat_history": chat_history}):
    print(chunk, end="", flush=True)
    current_time = time.time()
    if current_time - last_log_time >= log_interval:
        log_cpu_usage(":")
        last_log_time = current_time

4. 优化字符串拼接方式 #

频繁执行answer += chunk会产生多次内存分配和拷贝开销，改用列表收集chunk后一次性join，能显著降低CPU消耗：

answer_chunks = []
for chunk in rag_chain.stream({"question":question, "chat_history": chat_history}):
    print(chunk, end="", flush=True)
    answer_chunks.append(chunk)
answer = ''.join(answer_chunks)
# 回答生成完成后再更新chat_history
chat_history.extend([HumanMessage(content=question), AIMessage(content=answer)])

5. 确认LangChain链的流式逻辑 #

检查rag_chain结构，确保检索阶段仅执行一次，而非每个chunk都重复检索文档。例如使用RetrievalQA链时，检索操作应在stream调用前完成，流式仅作用于LLM的生成阶段，可通过添加日志确认检索是否重复执行。

6. 调整LLM流式输出的chunk大小 #

若LLM支持，可设置max_tokens_per_chunk参数，减少chunk的生成数量，从而降低循环迭代的次数，减少CPU上下文切换开销。

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