Skip to main content
TiDB Cloud 是一款全面的数据库即服务(DBaaS)解决方案,提供专用和无服务器两种选项。TiDB Serverless 现已将内置向量搜索集成到 MySQL 生态中。借助这一增强功能,您无需新建数据库或引入额外技术栈,即可使用 TiDB Serverless 无缝开发 AI 应用。请访问 tidb.cloud/ai 加入私测候补名单,抢先体验。
本笔记本介绍如何在 LangChain 中使用 TiDBLoader 从 TiDB 加载数据。

前提条件

在使用 TiDBLoader 之前,我们需要安装以下依赖:
pip install -qU langchain
然后,配置与 TiDB 的连接。本笔记本采用 TiDB Cloud 提供的标准连接方式来建立安全高效的数据库连接。
import getpass

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

从 TiDB 加载数据

以下是 TiDBLoader 的一些关键参数说明:
  • query(str):在 TiDB 数据库中执行的 SQL 查询。查询应选择您想要加载到 Document 对象中的数据。例如,可以使用 "SELECT * FROM my_table" 获取 my_table 中的所有数据。
  • page_content_columns(Optional[List[str]]):指定其值应包含在每个 Document 对象 page_content 中的列名列表。若设为 None(默认值),查询返回的所有列都将包含在 page_content 中。这允许您根据特定列定制每个文档的内容。
  • metadata_columns(Optional[List[str]]):指定其值应包含在每个 Document 对象 metadata 中的列名列表。默认为空列表,除非明确指定,否则不会包含任何元数据。这可用于为每个文档添加不构成主要内容但对处理或分析仍有价值的附加信息。
from sqlalchemy import Column, Integer, MetaData, String, Table, create_engine

# 连接数据库
engine = create_engine(tidb_connection_string)
metadata = MetaData()
table_name = "test_tidb_loader"

# 创建表
test_table = Table(
    table_name,
    metadata,
    Column("id", Integer, primary_key=True),
    Column("name", String(255)),
    Column("description", String(255)),
)
metadata.create_all(engine)


with engine.connect() as connection:
    transaction = connection.begin()
    try:
        connection.execute(
            test_table.insert(),
            [
                {"name": "Item 1", "description": "Description of Item 1"},
                {"name": "Item 2", "description": "Description of Item 2"},
                {"name": "Item 3", "description": "Description of Item 3"},
            ],
        )
        transaction.commit()
    except:
        transaction.rollback()
        raise
from langchain_community.document_loaders import TiDBLoader

# 设置 TiDBLoader 以检索数据
loader = TiDBLoader(
    connection_string=tidb_connection_string,
    query=f"SELECT * FROM {table_name};",
    page_content_columns=["name", "description"],
    metadata_columns=["id"],
)

# 加载数据
documents = loader.load()

# 显示已加载的文档
for doc in documents:
    print("-" * 30)
    print(f"content: {doc.page_content}\nmetada: {doc.metadata}")
------------------------------
content: name: Item 1
description: Description of Item 1
metada: {'id': 1}
------------------------------
content: name: Item 2
description: Description of Item 2
metada: {'id': 2}
------------------------------
content: name: Item 3
description: Description of Item 3
metada: {'id': 3}
test_table.drop(bind=engine)