集成 TiDB 向量检索与 Jina AI Embeddings API

本教程将指导你如何使用 Jina AI 为文本数据生成向量，并将这些向量存储到 TiDB 向量存储中，基于向量实现相似文本的检索。

前置条件

完成本教程，你需要：

• 已安装 Python 3.8 或更高版本。
• 已安装 Git。
• 一个 TiDB 集群。

如果你还没有 TiDB 集群，可以按如下方式创建：

• （推荐）参考 创建 TiDB Cloud Starter 集群 创建属于你自己的 TiDB Cloud 集群。
• 参考 部署本地测试 TiDB 集群 或 部署生产环境 TiDB 集群 创建 v8.4.0 或更高版本的本地集群。

运行示例应用

你可以按照以下步骤，快速了解如何将 TiDB 向量检索与 JinaAI Embedding 集成。

步骤 1. 克隆代码仓库

将 tidb-vector-python 仓库克隆到本地：

git clone https://github.com/pingcap/tidb-vector-python.git

步骤 2. 创建虚拟环境

为你的项目创建一个虚拟环境：

cd tidb-vector-python/examples/jina-ai-embeddings-demo
python3 -m venv .venv
source .venv/bin/activate

步骤 3. 安装依赖

为示例项目安装所需依赖：

pip install -r requirements.txt

步骤 4. 配置环境变量

从 Jina AI Embeddings API 页面获取 Jina AI API key，然后根据你选择的 TiDB 部署方式配置环境变量。

export JINA_API_KEY="****"
export TIDB_DATABASE_URL="mysql+pymysql://<USERNAME>:<PASSWORD>@<HOST>:<PORT>/<DATABASE>"
# 例如：export TIDB_DATABASE_URL="mysql+pymysql://root@127.0.0.1:4000/test"

步骤 5. 运行示例

python jina-ai-embeddings-demo.py

示例输出：

- Inserting Data to TiDB...
  - Inserting: Jina AI offers best-in-class embeddings, reranker and prompt optimizer, enabling advanced multimodal AI.
  - Inserting: TiDB is an open-source MySQL-compatible database that supports Hybrid Transactional and Analytical Processing (HTAP) workloads.
- List All Documents and Their Distances to the Query:
  - distance: 0.3585317326132522
    content: Jina AI offers best-in-class embeddings, reranker and prompt optimizer, enabling advanced multimodal AI.
  - distance: 0.10858102967720984
    content: TiDB is an open-source MySQL-compatible database that supports Hybrid Transactional and Analytical Processing (HTAP) workloads.
- The Most Relevant Document and Its Distance to the Query:
  - distance: 0.10858102967720984
    content: TiDB is an open-source MySQL-compatible database that supports Hybrid Transactional and Analytical Processing (HTAP) workloads.

示例代码片段

从 Jina AI 获取向量

定义 generate_embeddings 辅助函数，调用 Jina AI embeddings API：

import os
import requests
import dotenv

dotenv.load_dotenv()

JINAAI_API_KEY = os.getenv('JINAAI_API_KEY')

def generate_embeddings(text: str):
    JINAAI_API_URL = 'https://api.jina.ai/v1/embeddings'
    JINAAI_HEADERS = {
        'Content-Type': 'application/json',
        'Authorization': f'Bearer {JINAAI_API_KEY}'
    }
    JINAAI_REQUEST_DATA = {
        'input': [text],
        'model': 'jina-embeddings-v2-base-en'  # with dimension 768.
    }
    response = requests.post(JINAAI_API_URL, headers=JINAAI_HEADERS, json=JINAAI_REQUEST_DATA)
    return response.json()['data'][0]['embedding']

连接 TiDB 集群

通过 SQLAlchemy 连接 TiDB 集群：

import os
import dotenv

from tidb_vector.sqlalchemy import VectorType
from sqlalchemy.orm import Session, declarative_base

dotenv.load_dotenv()

TIDB_DATABASE_URL = os.getenv('TIDB_DATABASE_URL')
assert TIDB_DATABASE_URL is not None
engine = create_engine(url=TIDB_DATABASE_URL, pool_recycle=300)

定义向量表结构

创建名为 jinaai_tidb_demo_documents 的表，包含用于存储文本的 content 字段和用于存储向量的 content_vec 字段：

from sqlalchemy import Column, Integer, String, create_engine
from sqlalchemy.orm import declarative_base

Base = declarative_base()

class Document(Base):
    __tablename__ = "jinaai_tidb_demo_documents"

    id = Column(Integer, primary_key=True)
    content = Column(String(255), nullable=False)
    content_vec = Column(
        # DIMENSIONS is determined by the embedding model,
        # for Jina AI's jina-embeddings-v2-base-en model it's 768.
        VectorType(dim=768),
        comment="hnsw(distance=cosine)"

使用 Jina AI 生成向量并存储到 TiDB

使用 Jina AI Embeddings API 为每条文本生成向量，并将其存储到 TiDB：

TEXTS = [
   'Jina AI offers best-in-class embeddings, reranker and prompt optimizer, enabling advanced multimodal AI.',
   'TiDB is an open-source MySQL-compatible database that supports Hybrid Transactional and Analytical Processing (HTAP) workloads.',
]
data = []

for text in TEXTS:
    # Generate embeddings for the texts via Jina AI API.
    embedding = generate_embeddings(text)
    data.append({
        'text': text,
        'embedding': embedding
    })

with Session(engine) as session:
   print('- Inserting Data to TiDB...')
   for item in data:
      print(f'  - Inserting: {item["text"]}')
      session.add(Document(
         content=item['text'],
         content_vec=item['embedding']
      ))
   session.commit()

在 TiDB 中基于 Jina AI 向量进行语义检索

通过 Jina AI embeddings API 为查询文本生成向量，然后基于 查询文本的向量 与 向量表中每条数据的向量 的余弦距离，检索最相关的文档：

query = 'What is TiDB?'
# Generate the embedding for the query via Jina AI API.
query_embedding = generate_embeddings(query)

with Session(engine) as session:
    print('- The Most Relevant Document and Its Distance to the Query:')
    doc, distance = session.query(
        Document,
        Document.content_vec.cosine_distance(query_embedding).label('distance')
    ).order_by(
        'distance'
    ).limit(1).first()
    print(f'  - distance: {distance}\n'
          f'    content: {doc.content}')

相关阅读

• 向量数据类型
• 向量检索索引