使用 TiDB Serverless 存储聊天记录并构建历史数据驱动的聊天应用

技术背景介绍

TiDB 是一款分布式数据库，提供强大的 SQL 兼容性及横向扩展能力。近期，TiDB Cloud 推出了全新的 Serverless 模式，并集成了内置的向量搜索功能，进一步降低了 AI 应用开发的门槛。这意味着开发者可以在熟悉的 MySQL 环境中，直接实现复杂的 AI 应用数据存储需求。

在本篇文章中，我们将展示如何利用 TiDB Serverless 存储聊天记录数据，并基于这些记录构建一个动态对话应用。TiDB Serverless 提供免费集群，推荐访问 https://pingcap.com/ai 以创建集群并开始使用。

---

核心原理解析

以下是实现动态聊天交互的核心步骤：

1. TiDB 存储聊天历史记录： 借助 langchain-community 提供的 TiDBChatMessageHistory，我们能够将用户与 AI 的聊天记录存储到 TiDB。
2. 历史记录的动态加载： 使用 RunnableWithMessageHistory 可以在对话过程中自动加载并重用历史聊天记录。
3. LangChain 集成： 通过 LangChain 提供的 ChatPromptTemplate 和 ChatOpenAI，快速搭建基于历史记录的 AI 回答逻辑。

---

代码实现演示

以下是完整的实现流程，请按照步骤操作确保代码能够正常运行。

1. 环境依赖安装

首先，安装所需的 Python 包：

pip install --upgrade --quiet langchain langchain_openai langchain-community

2. 配置 OpenAI API Key

LangChain 需要调用 OpenAI 的模型服务，以下代码用于设置 API 密钥。

import getpass
import os

# 输入你的 OpenAI API 密钥
os.environ["OPENAI_API_KEY"] = getpass.getpass("Input your OpenAI API key:")

3. 配置 TiDB 数据库连接

使用 TiDB Cloud 提供的标准连接方式，以下是连接字符串的模板：

# 提供从 TiDB Cloud 控制台获取的连接信息
tidb_connection_string_template = "mysql+pymysql://<USER>:<PASSWORD>@<HOST>:4000/<DB>?ssl_ca=/etc/ssl/cert.pem&ssl_verify_cert=true&ssl_verify_identity=true"

# 输入 TiDB 密码
tidb_password = getpass.getpass("Input your TiDB password:")

# 替换模板中的密码
tidb_connection_string = tidb_connection_string_template.replace("<PASSWORD>", tidb_password)

4. 生成历史数据

我们首先创建一组初始的聊天记录，并存储到 TiDB 中。

from datetime import datetime
from langchain_community.chat_message_histories import TiDBChatMessageHistory

# 初始化 TiDBChatMessageHistory
history = TiDBChatMessageHistory(
    connection_string=tidb_connection_string,
    session_id="code_gen",  # 会话标识
    earliest_time=datetime.utcnow(),  # 可选：设置加载时间点
)

# 添加历史消息
history.add_user_message("How's our feature going?")
history.add_ai_message("It's going well. We are working on testing now. It will be released in Feb.")

# 查看存储的消息
print(history.messages)

运行此代码后，TiDB 中将存储以下消息记录：

[HumanMessage(content="How's our feature going?"),
 AIMessage(content="It's going well. We are working on testing now. It will be released in Feb.")]

5. 构建动态聊天交互

接下来，我们基于 LangChain 的 ChatPromptTemplate 和历史记录构建一个动态聊天链。

from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_openai import ChatOpenAI

# 定义对话模板
prompt = ChatPromptTemplate.from_messages(
    [
        ("system", "You're an assistant who's good at coding. You're helping a startup build"),
        MessagesPlaceholder(variable_name="history"),  # 动态历史消息填充
        ("human", "{question}"),  # 用户问题
    ]
)

# 创建聊天链
chain = prompt | ChatOpenAI()

我们将历史数据与聊天链结合，以实现实时的对话上下文管理。

from langchain_core.runnables.history import RunnableWithMessageHistory

# 定义历史数据驱动的链
chain_with_history = RunnableWithMessageHistory(
    chain,
    lambda session_id: TiDBChatMessageHistory(
        session_id=session_id,
        connection_string=tidb_connection_string
    ),
    input_messages_key="question",  # 用户输入
    history_messages_key="history",  # 历史记录
)

6. 开始对话

通过以下代码，发起一次对话，并结合历史记录生成回答。

response = chain_with_history.invoke(
    {"question": "Today is Jan 1st. How many days until our feature is released?"},
    config={"configurable": {"session_id": "code_gen"}},  # 使用会话 ID
)

print(response)

输出示例

AIMessage(content='There are 31 days in January, so there are 30 days until our feature is released in February.')

7. 检查历史数据

重新加载并验证历史记录数据，确保对话记录完整存储。

history.reload_cache()
print(history.messages)

输出示例

[HumanMessage(content="How's our feature going?"),
 AIMessage(content="It's going well. We are working on testing now. It will be released in Feb."),
 HumanMessage(content='Today is Jan 1st. How many days until our feature is released?'),
 AIMessage(content='There are 31 days in January, so there are 30 days until our feature is released in February.')]

---

应用场景分析

1. 客户服务机器人： 利用 TiDB 存储用户对话历史，提升个性化服务能力。
2. 协作工具集成： 在团队生产力工具中记录讨论上下文，供后续参考。
3. AI 辅助开发： 通过保存历史交互，构建基于上下文的代码生成或代码解析助手。

---

实践建议

1. 合理规划会话 ID： 在多用户场景中，确保每个用户的会话 ID 唯一，以便正确加载历史记录。
2. 优化存储性能： 配置 TiDB 的索引和连接池设置，确保高效的数据存取。
3. 定期清理旧数据： 对于长期积累的聊天记录，建议设计清理策略或归档方案。

---

如果遇到问题欢迎在评论区交流。