Auto Embedding 示例

本示例展示了如何结合 Auto Embedding 功能与 pytidb 客户端使用。

1. 使用 pytidb 客户端连接到 TiDB。
2. 定义一个表，并配置带有自动生成向量功能的 VectorField。
3. 插入纯文本数据：系统会在后台自动填充向量嵌入。
4. 使用自然语言查询执行向量搜索：系统会透明地生成向量嵌入。

前提条件

开始之前，请确保你具备以下条件：

• Python (>=3.10)：安装 Python 3.10 或以上版本。
• 一个 TiDB Cloud Starter 实例：你可以在 TiDB Cloud 上创建一个免费的 TiDB Cloud Starter 实例。

运行方式

第 1 步：克隆 pytidb 仓库

git clone https://github.com/pingcap/pytidb.git
cd pytidb/examples/auto_embedding/

第 2 步：安装所需依赖包

python -m venv .venv
source .venv/bin/activate
pip install -r reqs.txt

第 3 步：设置环境变量

1. 在 TiDB Cloud 控制台 中，进入 My TiDB 页面，然后点击目标 TiDB Cloud Starter 实例的名称，进入其实例概览页面。
2. 点击右上角的 Connect。此时会显示连接对话框，其中列出了连接参数。
3. 根据连接参数按如下方式设置环境变量：

cat > .env <<EOF
TIDB_HOST={gateway-region}.prod.aws.tidbcloud.com
TIDB_PORT=4000
TIDB_USERNAME={prefix}.root
TIDB_PASSWORD={password}
TIDB_DATABASE=test

# Using TiDB Cloud Free embedding model by default, which does not require setting up any API key
EMBEDDING_PROVIDER=tidbcloud_free
EOF

第 4 步：运行示例

python main.py

预期输出：

=== Define embedding function ===
Embedding function (model id: tidbcloud_free/amazon/titan-embed-text-v2) defined

=== Define table schema ===
Table created

=== Truncate table ===
Table truncated

=== Insert sample data ===
Inserted 3 chunks

=== Perform vector search ===
id: 1, text: TiDB is a distributed database that supports OLTP, OLAP, HTAP and AI workloads., distance: 0.30373281240458805
id: 2, text: PyTiDB is a Python library for developers to connect to TiDB., distance: 0.422506501973434
id: 3, text: LlamaIndex is a Python library for building AI-powered applications., distance: 0.5267239638442787

相关资源

• 源代码：在 GitHub 上查看