利用TiDB存储聊天消息历史并进行AI交互

在现代AI应用中，存储和处理大量的聊天消息是非常常见的需求。TiDB Cloud 提供了一种便捷的方式，通过其无服务器（Serverless）集成的向量搜索功能，可以在MySQL环境中开发AI应用。本文将介绍如何利用TiDB来存储聊天消息历史，并使用它们进行动态的AI交互。

技术背景介绍

TiDB 是一款融合了OLTP和OLAP的分布式数据库，支持水平弹性扩展、强一致性、MySQL 兼容以及高可用性。其云服务TiDB Cloud提供了数据库即服务（DBaaS）的解决方案，其中无服务器选项大大简化了数据库管理。

核心原理解析

本次实现的核心在于利用TiDB的集成向量搜索和开放API接口，将聊天记录持久化，并结合AI模型进行智能对话。这一过程需要：

• 配置数据库连接
• 存储和提取聊天数据
• 建立AI交互模型

代码实现演示

接下来，我将分步骤展示如何在Python中实现这些功能。

依赖安装

首先，安装必要的Python库：

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

配置OpenAI和TiDB

确保你的API Key和数据库连接信息设置正确：

import getpass
import os

# 设置OpenAI API密钥
os.environ["OPENAI_API_KEY"] = getpass.getpass("Input your OpenAI API key:")

# 配置TiDB连接
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_password = getpass.getpass("Input your TiDB password:")
tidb_connection_string = tidb_connection_string_template.replace("<PASSWORD>", tidb_password)

存储历史数据

创建并存储一组历史数据：

from datetime import datetime
from langchain_community.chat_message_histories import TiDBChatMessageHistory

# 初始化历史数据存储
history = TiDBChatMessageHistory(
    connection_string=tidb_connection_string,
    session_id="code_gen",
    earliest_time=datetime.utcnow(),  # 可选：限制加载时间点之后的消息
)

# 添加用户与AI的消息
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.")

使用历史数据进行对话

利用LangChain库，通过历史数据进行AI互动：

from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_openai import ChatOpenAI
from langchain_core.runnables.history import RunnableWithMessageHistory

# 创建对话提示模板
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()
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",
)

# 启动对话
response = chain_with_history.invoke(
    {"question": "Today is Jan 1st. How many days until our feature is released?"},
    config={"configurable": {"session_id": "code_gen"}},
)
print(response.content)

查看历史记录

# 重新加载缓存并查看消息历史
history.reload_cache()
print(history.messages)

应用场景分析

这种实现方案特别适合需要长期记录并分析客户对话的大型企业应用，可以帮助企业提供更精准的客户支持服务。

实践建议

• 确保数据库连接的安全性，切勿在代码中明文保存敏感信息。
• 通过TiDB的无服务器架构，可以降低成本，实现资源的自动调配。
• 定期备份聊天记录，以免数据泄露或丢失。

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