如何使用TiDBLoader加载数据:构建AI应用的第一步

78 阅读2分钟
# 引言

TiDB Cloud 提供了一种全面的数据库即服务解决方案,其中的 TiDB Serverless 正逐步将内置的向量搜索集成到 MySQL 领域。这一增强使得开发者可以在无需引入新数据库或额外技术栈的情况下轻松构建AI应用。本篇文章将介绍如何使用 TiDBLoader 从 TiDB 加载数据,帮助您快速上手。

# 主要内容

## 1. 环境准备

在使用 TiDBLoader 之前,需要安装相应的依赖:

```bash
%pip install --upgrade --quiet langchain

2. 配置数据库连接

通过 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
)

3. 使用 TiDBLoader 加载数据

TiDBLoader 提供了一些关键参数,以便用户自定义数据加载行为:

  • query(str): 执行的 SQL 查询。
  • page_content_columns(Optional[List[str]]): 需要包含在每个 Document 对象中的列。
  • metadata_columns(Optional[List[str]]): 包含在元数据中的列。

代码示例

以下是如何使用 TiDBLoader 加载数据的完整代码示例:

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}\nmetadata: {doc.metadata}")

# 清理测试表
test_table.drop(bind=engine)

常见问题和解决方案

  • 连接错误:确保您的连接字符串和密码正确。
  • 数据加载失败:检查 SQL 查询及字段名是否正确。

总结和进一步学习资源

通过 TiDBLoader,您可以无缝从 TiDB 加载数据用于 AI 应用的开发。同时,可以查看以下资源获取更多信息:

参考资料

如果这篇文章对你有帮助,欢迎点赞并关注我的博客。您的支持是我持续创作的动力!

---END---