个人AI助手 - 阶段2 多轮对话记忆 & Token统计
1. 阶段定位
本节是个人 AI 助手第 2 阶段：多轮对话记忆与实时 Token 统计。
核心任务是接入 LangChain 记忆组件、打通 Token 统计、补齐前端会话管理。
本阶段不做 RAG、不做工具调用、不做持久化存储，只完成基础对话优化，为后续长对话优化提供支撑。
2. 阶段目标
- 功能目标：完成多轮连贯对话、实时 Token 统计、对话清空/管理，打通「用户输入 — 记忆拼接 — 模型调用 — Token 统计 — 前端展示」链路。
- 架构目标：落地「消息历史容器 + 会话管理器」记忆体系，统一 Token 统计流程，解耦记忆组件与核心业务链。
3. 架构与选型
3.1 架构分层
- 前端层：负责界面渲染、session_id 生成、Token 面板展示、清空对话；不负责历史拼接与 Token 计算。
- 后端层：负责会话历史存储、记忆绑定、Token 统计、结构化返回；统一 session_id 隔离。
- 模型层：负责对话生成，不感知记忆与统计逻辑。
3.2 技术选型
记忆组件
- 选型：
InMemoryChatMessageHistory+RunnableWithMessageHistory - 淘汰：手写列表、
ConversationBufferMemory+ConversationChain - 原因：手写仅适合 Demo，扩展性差；旧版组件已废弃，不支持多用户隔离与 LCEL。
- 优势：格式标准、支持多用户隔离、可无缝对接 Token 统计与持久化。
Token 统计
- 选型：dashscope 内置 Qwen 官方 tokenizer
- 淘汰：接口返回 Token、通用 tiktoken
- 原因：接口返回属于事后结果，流式场景不返回；通用 tokenizer 与 Qwen 计费口径不一致。
- 优势：与 API 计费口径一致，支持请求前预判与流式事后统计。
4. 核心设计
- 输入标准化：按 ChatML 格式拼接系统提示 + 历史 + 当前问题，插入
<|im_start|>/<|im_end|>，末尾追加<|im_start|>assistant。 - 调用标准化：用
RunnableWithMessageHistory封装记忆读写，固化input_messages_key/history_messages_key/get_session_history。 - 输出标准化：统一返回「回答内容 + 输入 Token + 输出 Token + 历史总 Token」。
5. 核心实现
5.1 核心代码
5.1.1 会话历史仓库 + 提示词模板
# config/llm_config.py
import os
from typing import List
from dotenv import load_dotenv
from langchain_community.chat_models import ChatZhipuAI
from langchain_core.chat_history import InMemoryChatMessageHistory
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables import RunnableWithMessageHistory
load_dotenv()
ZHIPU_API_KEY = os.getenv('ZHIPU_API_KEY')
# 模型参数
MODEL_NAME = 'glm-4-flash'
TIMEOUT = 60
# 全局会话仓库：session_id → 独立聊天记录
store = {}
chat_prompt_template = ChatPromptTemplate.from_messages([
("system",
"你是瑶瑶的私人助手，协助她完成工作生活学习上的事情。\n" # Python 隐式字符串拼接：多个字符串字面量相邻出现,解释器会自动把它们合并成一个字符串
"如果用户消息中提供了参考资料，请务必严格基于资料回答；"
"资料中没有的信息不要编造，可以说明“资料中未提及”。\n"
"如果没有提供参考资料，则正常回答。"),
MessagesPlaceholder("history"),
("user", "参考资料：\n{context}\n\n用户问题：{user_input}"),
])
def get_session_history(session_id: str):
if session_id not in store:
# InMemoryChatMessageHistory 是 LangChain 提供的内存存储，程序运行期间有效，重启即丢失
store[session_id] = InMemoryChatMessageHistory()
return store[session_id]
5.1.2 带记忆的链
# config/llm_config.py
import os
from typing import List
from dotenv import load_dotenv
from langchain_community.chat_models import ChatZhipuAI
from langchain_core.chat_history import InMemoryChatMessageHistory
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables import RunnableWithMessageHistory
load_dotenv()
ZHIPU_API_KEY = os.getenv('ZHIPU_API_KEY')
# 模型参数
MODEL_NAME = 'glm-4-flash'
TIMEOUT = 60
# 全局会话仓库：session_id → 独立聊天记录
store = {}
chat_prompt_template = ChatPromptTemplate.from_messages([
("system",
"你是瑶瑶的私人助手，协助她完成工作生活学习上的事情。\n" # Python 隐式字符串拼接：多个字符串字面量相邻出现,解释器会自动把它们合并成一个字符串
"如果用户消息中提供了参考资料，请务必严格基于资料回答；"
"资料中没有的信息不要编造，可以说明“资料中未提及”。\n"
"如果没有提供参考资料，则正常回答。"),
MessagesPlaceholder("history"),
("user", "参考资料：\n{context}\n\n用户问题：{user_input}"),
])
def get_session_history(session_id: str):
if session_id not in store:
# InMemoryChatMessageHistory 是 LangChain 提供的内存存储，程序运行期间有效，重启即丢失
store[session_id] = InMemoryChatMessageHistory()
return store[session_id]
class LLM:
# 瑶瑶的私人助手 Chain
# 自带：提示词 + 输出解析器 + 流式输出 + 智谱模型 + 多轮对话记忆
@staticmethod
def get_chain(
model_name: str = MODEL_NAME,
timeout: int = TIMEOUT,
):
# 预设参数组合：不同业务场景（RAG问答 / 长文深度 / 快速省钱 / 工具调用）
# plus = 日常 RAG：0.7 + 2000 + streaming = True
# max = 长文 / 深度：0.6 + 4096 + streaming = True
# turbo = 省钱快答：0.7 + 1024 + streaming = True
# 工具调用：温度 0.2、不流式
chat_model = ChatZhipuAI(
model=model_name,
api_key=ZHIPU_API_KEY,
timeout=timeout,
temperature=0.7,
max_tokens=2000,
streaming=True,
verbose=False
)
parser = StrOutputParser()
chain = chat_prompt_template | chat_model | parser
# 把记忆绑定到链（LCEL 唯一正确方式），得到带记忆的链
# 接收一个普通的 LCEL 链，给它"套上"记忆能力，返回一个新的、带记忆的链
chain_with_memory = RunnableWithMessageHistory(
chain, # 被包装的原始链
get_session_history=get_session_history, # 一个回调函数，用来根据 session_id 拿到对应的历史记录对象；按会话隔离记忆
input_messages_key="user_input", # 对应 {user_input}
history_messages_key="history", # 对应 history
)
return chain_with_memory
5.1.3 Token 统计工具
# utils/token_util.py
from dashscope import get_tokenizer
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from config.llm_config import get_session_history, chat_prompt_template
# ------------------------------
# 全局变量：存储每个会话的历史总Token
# ------------------------------
session_token_total = {}
tokenizer = get_tokenizer("qwen-plus")
def build_qwen_chatml(msg_list):
chatml = ""
for m in msg_list:
if m.type == "system":
r = "system"
elif m.type == "human":
r = "user"
elif m.type == "ai":
r = "assistant"
else:
r = "user"
chatml += f"<|im_start|>{r}\n{m.content}<|im_end|>\n"
chatml += "<|im_start|>assistant\n"
return chatml
# ===================== 工具：拼接完整prompt并计算输入token数量 =====================
def calculate_input_tokens(user_input,session_id, context):
# 获取该会话的历史聊天记录
history = get_session_history(session_id)
# 拼装系统提示词+历史消息+当前用户提问
prompt_template = ChatPromptTemplate.from_messages([
("system", "你是瑶瑶的私人助手，协助她完成工作生活学习上的事情"),
MessagesPlaceholder("history"),
("user", "{user_input}"),
])
prompt_val = chat_prompt_template.invoke({
"user_input": user_input,
"context": context,
"history": history.messages
})
# 转为通义chatml格式后统计token
chatml_text = build_qwen_chatml(prompt_val.to_messages())
return len(tokenizer.encode(chatml_text))
5.1.4 对话接口与后置统计
# api/chat_endpoint.py
import json
import time
from datetime import datetime
from fastapi import APIRouter
from fastapi.responses import StreamingResponse
from langchain_classic.retrievers import ContextualCompressionRetriever
from pydantic import BaseModel
# 项目内部模块导入：配置、LLM、向量、缓存工具、Token工具
from config.app_config import CACHE_TTL
from config.llm_config import get_session_history
from core.embeddings import embeddings
from core.llm_loader import get_chain
from utils.cache_util import session_question_answer_total, init_session_cache, clean_expired_session_cache, get_cached_answer
from utils.token_util import tokenizer, session_token_total, calculate_input_tokens
from utils.retriever import get_best_retriever
from core.reranker import RawDashScopeRerank
from core.context_compress import create_compression_retriever
router = APIRouter()
# 请求入参模型：前端请求必须携带用户提问+会话ID
class ChatRequest(BaseModel):
user_input: str
session_id: str
use_rag: bool
# ===================== 获取会话历史消息接口 =====================
@router.get("/get_history")
def get_history(session_id: str):
history = get_session_history(session_id)
return {"history": history.messages}
# ===================== 对话主接口 =====================
# 同步接口，返回流式响应，实际生成逻辑下放至 generate_stream
@router.post("/chat")
def chat(body: ChatRequest):
# body是请求体（FastAPI 按类型自动从 body 解析）
# 若要访问完整请求（headers/url/method 等），需额外加一个 Request 类型入参
print("body: ", body)
chain = get_chain()
# 把 session_id 放进 config.configurable
config = {"configurable": {"session_id": body.session_id}}
# 返回流式生成器
# generate_stream 负责"怎么产出数据"
# StreamingResponse 负责"怎么把产出的数据通过 HTTP 一块块送出去"
return StreamingResponse(generate_stream(body, chain, config), media_type="text/plain")
# ===================== 流式生成总入口函数：串联全流程 =====================
# 执行顺序：初始化缓存→清理过期→匹配缓存→命中缓存直接返回/未命中调用LLM
def generate_stream(req: ChatRequest, chain, config):
# 记录接口开始时间
start_time = datetime.now()
print("开始时间：", start_time.strftime("%Y-%m-%d %H:%M:%S"))
# 问题文本向量化
question_vector = embeddings.embed_query(req.user_input)
yield from stream_llm_result(req, chain, config, start_time, question_vector)
# ===================== 工具：调用LLM流式输出，finally中收尾统计&落地缓存 =====================
def stream_llm_result(req, chain, config, start_time, question_vector):
chunk_list = ""
try:
# 迭代流式分片，拼接完整回答
for chunk in chain.stream({"user_input": req.user_input, "context":context}, config=config):
chunk_list += chunk
yield chunk
finally:
# LLM生成结束后统一处理
yield from handle_llm_finish(req, chunk_list, start_time, question_vector, context)
# ===================== 工具：LLM生成完毕后置处理：Token统计+缓存落地+返回尾部JSON =====================
def handle_llm_finish(req, chunk_list, start_time, question_vector, context):
end = datetime.now()
cost_ms = int((end - start_time).total_seconds() * 1000)
print(f"本次输出耗时：{cost_ms} ms")
# 计算输入prompt总token
input_tokens = calculate_input_tokens(req.user_input, req.session_id, context)
# 计算模型输出回答token
output_tokens = len(tokenizer.encode(chunk_list))
session_id = req.session_id
# 累加当前会话历史总token
last_total = session_token_total.get(session_id, 0)
new_total = last_total + input_tokens + output_tokens
session_token_total[session_id] = new_total
print("本次输入 token：", input_tokens)
print("本次输出 token：", output_tokens)
print("历史总Token：", new_total)
# 将本轮问答向量+答案写入内存缓存，设置过期时间
session_question_answer_total[session_id].append({
"question_vector": question_vector,
"answer": chunk_list,
"expire_time": time.time() + CACHE_TTL
})
# 尾部返回token统计json
yield json.dumps({
"input_tokens": input_tokens,
"output_tokens": output_tokens,
"history_total": new_total,
"total_time": cost_ms
}, ensure_ascii=False)
5.2 问题修复
- Token 与接口不一致：未按 ChatML 拼接。解决：严格插入特殊标记，末尾追加
<|im_start|>assistant。 - 流式输出 Token 丢失：统计未放 finally。解决：输出 Token 统计放入 finally。
- 清空后旧历史仍被读取：未按 session_id 隔离。解决：前端重置 session_id，后端依赖变更自动隔离。
6. 交付成果
- 架构交付：记忆架构、Token 统计链路、session_id 隔离分层。
- 后端工程交付：
/chat接口、/history接口、会话仓库、Token 计数工具、历史总 Token 累加器。 - 工程规范交付：session_id 隔离规范、ChatML 拼接规范、Token 统计执行时机规范。
- 功能交付：多轮连贯对话、实时 Token 统计、对话清空/管理、刷新后历史恢复。
7. 运行效果
- 服务运行：
uvicorn正常启动，/docs可访问，无报错。 - 接口运行：
/chat返回回答 + 输入/输出/历史总 Token，参数正确。 - 对话交互：同一 session_id 下第二轮能正确回答「我是谁」，流式逐字返回、无截断。
- 整体效果：多轮记忆、Token 统计、清空对话、历史恢复四条链路均正常。
8. 后续迭代价值
本阶段搭建的记忆架构与 Token 统计底座，后续可直接迭代窗口记忆、历史摘要、会话持久化、Token 计费风控等功能，无需重构底层代码。
