结合 Jina AI 嵌入模型 API 使用 TiDB 向量搜索

本文档将展示如何使用 Jina AI 为文本数据生成向量嵌入，然后将向量嵌入存储在 TiDB 中，并根据向量嵌入搜索相似文本。

前置需求

为了能够顺利完成本文中的操作，你需要提前：

• 在你的机器上安装 Python 3.8 或更高版本
• 在你的机器上安装 Git
• 准备一个 TiDB 集群

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

• 参考部署本地测试 TiDB 集群或部署正式 TiDB 集群，创建本地集群。
• 参考创建 TiDB Cloud Serverless 集群，创建 TiDB Cloud 集群。

运行示例应用程序

您可以通过以下步骤快速了解如何结合 Jina AI 嵌入模型 API 使用 TiDB 向量搜索。

第 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 密钥，然后根据 TiDB 集群的部署方式不同，选择对应的环境变量配置方式。

• 本地部署 TiDB
• TiDB Cloud Serverless

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 的嵌入 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)
    )

使用 Jina AI 生成向量嵌入并存入 TiDB

使用 Jina AI 嵌入 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:
    # 通过 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()

使用 Jina AI 生成的向量嵌入在 TiDB 中执行语义搜索

通过 Jina AI 的嵌入 API 生成查询文本的向量嵌入，然后根据查询文本的向量嵌入和向量表中各个向量嵌入之间的余弦距离搜索最相关的 document：

query = 'What is TiDB?'
# 通过 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}')

另请参阅

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