LangChain之大模型处理记忆+摘要生成+关系型数据库

from langchain_community.chat_message_histories import SQLChatMessageHistory
from langchain_core.chat_history import InMemoryChatMessageHistory
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables import RunnableWithMessageHistory, RunnablePassthrough
from langchain_openai import ChatOpenAI
from dotenv import load_dotenv
import os

load_dotenv()
APIKEY = os.environ.get('DEEPSEEK_API_KEY')

# 大模型定义
llm = ChatOpenAI(
    model = "deepseek-chat",
    temperature=1.3,
    base_url="https://api.deepseek.com/",
    api_key=APIKEY
)



# 1.提示词模板

prompt = ChatPromptTemplate.from_messages(
    # 通用聊天模板
    [
        # 系统提示词
        ('system', '{system_message}'), # 改为了动态注入的系统消息提示词
        # 聊天历史记录
        # 历史消息占位符，其中的'chat_history' 可以改，但是这个值是默认的不建议改，专门构建链的对象默认占位符就是这个值
        # optional=True 就是可能没有消息记录
        MessagesPlaceholder(variable_name='chat_history', optional=True),
        # ("placeholder", "{chat_history}"), # 这种消息占位符也可
        # 用户输入的新的问题
        ('human', '{input}')
    ]
)


# 构建链
chain = prompt | llm # 基础的执行链

# 2.存储聊天记录（内存，关系型数据库或者redis数据库）

# 存到内存中
# store 是存所有会话的所有历史聊天记录列表，类型是字典
# store = {} # 用来保存历史信息，key：会话id（session id）

def get_session_history(session_id: str):
    """从关系型数据库的历史消息列表中 返回当前会话的所有历史消息"""

    # 关系型数据库解决了聊天记录永久保存的问题
    return SQLChatMessageHistory(
        session_id=session_id,
        connection_string='sqlite:///chat_history.db' #sqlite是一个轻量级关系型数据库，是个文件，按照这个模板构造
    )


# LangChain 中所有的消息类型：
# SystemMessage 系统消息
# HumanMessage 用户消息
# AIMessage  AI返回的消息
# ToolMessage 工具返回的消息

# 3.创建带历史记录功能的处理链
chain_with_message_history = RunnableWithMessageHistory(
    chain,
    get_session_history,
    input_messages_key='input',
    history_messages_key='chat_history',
)

# 4.剪辑和摘要上下文历史记录：保留最近前两条数据，把之前的所有消息形成摘要

def summarize_messages(current_input):
    """剪辑和摘要上下文历史记录"""
    session_id = current_input['config']['configurable']['session_id']
    # debug: current_input={'input': '我是Atta'}

    if not session_id:
        raise ValueError("必须通过config参数提供session_id")

    # 获取当前会话ID的所有历史聊天记录
    chat_history = get_session_history(session_id)
    stored_messages = chat_history.messages # 返回类型是数组

    if len(stored_messages) <= 2: # 不满足时直接返回原始消息
        return {
        "original_messages": stored_messages,
        "summary": None
        }

    # 剪辑消息列表
    last_two_messages = stored_messages[-2:] #  最新的两条信息，也就是需要保留的消息
    message_to_summarize = stored_messages[:-2] # 之前的信息, 这里的信息需要进行摘要

    summarization_prompt = ChatPromptTemplate.from_messages(
        [
            ('system', '请将以下对话历史压缩为一条保留关键信息的摘要信息'),
            ('placeholder', '{chat_history}'),
            ('human', '请生成包含上述对话核心内容的摘要，保留重要事实和决策')
        ]
    )

    summarization_chain = summarization_prompt | llm
    # 生成的摘要类型：（AIMessage）,所以在数据库中看到的消息为"type":"ai"
    summary_message = summarization_chain.invoke({'chat_history': message_to_summarize}) # 把需要摘要的消息列表传进去，让大模型自动生成摘要

    # 返回结构化结果
    return {
        "original_messages": last_two_messages, # 保留原始的信息
        "summary": summary_message # 生成的摘要
    }



# 最终的链
# 1.{input : 原来的, message_summarized=summarize_messages 函数执行后的返回值}
# 2.{input : 原来的, chat_history: message_summarized['original_messages'], system_message: message_summarized['original_messages']}
final_chain = (RunnablePassthrough.assign(message_summarized=summarize_messages) | RunnablePassthrough.assign(
    input=lambda x: x['input'],
    chat_history=lambda x: x['message_summarized']['original_messages'],
    system_message=lambda x: f"你是一个乐于助人的助手，尽可能回答用户的问题，你的名字叫Btta，摘要{x['message_summarized']['summary']}" if x[
        'message_summarized'].get('summary') else '无摘要'
) | chain_with_message_history)

# 其中config 是有固定模板的，session_id 不能写死，正常是由客户端随机生成的，这里只是用来测试
result1 = final_chain.invoke({'input':"我是Atta", "config":{"configurable":{"session_id": "user123"}}}, config={"configurable":{"session_id": "user123"}})
print(result1)
result2 = final_chain.invoke({'input': "我叫什么名字", "config":{"configurable":{"session_id": "user123"}}}, config={"configurable":{"session_id":"user123"}})
print(result2)
result3 = final_chain.invoke({'input': "牛顿的第三定律是什么", "config":{"configurable":{"session_id": "user123"}}}, config={"configurable":{"session_id":"user123"}})
print(result3)