PingKai Logo下载
搜索

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

本教程将指导你如何使用 Jina AI 生成文本 embedding，将其存储到 TiDB 中，并基于 embedding 实现相似文本搜索。

注意

前置条件

完成本教程，你需要：

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

运行示例应用

你可以按照以下步骤，快速学习如何将 TiDB 向量搜索与 Jina AI 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 部署方式配置环境变量。

对于 TiDB Cloud Starter 或 Essential 实例，按以下步骤获取连接字符串并配置环境变量：

  1. 进入 My TiDB 页面，点击目标 TiDB Cloud Starter 或 Essential 实例的名称进入其概览页。

  2. 点击右上角的 Connect，弹出连接对话框。

  3. 确保连接对话框中的配置与你的运行环境一致。

    • Connection Type 设置为 Public
    • Branch 设置为 main
    • Connect With 设置为 SQLAlchemy
    • Operating System 与你的环境一致

    提示：

    如果你的程序运行在 Windows Subsystem for Linux (WSL) 中，请切换到对应的 Linux 发行版。

  4. 切换到 PyMySQL 标签页，点击 Copy 图标复制连接字符串。

    提示：

    如果你还未设置密码，请点击 Create password 生成随机密码。

  5. 在终端中将 Jina AI API key 和 TiDB 连接字符串设置为环境变量，或创建 .env 文件，内容如下：

    JINAAI_API_KEY="****"
TIDB_DATABASE_URL="{tidb_connection_string}"

    以下是 macOS 下的连接字符串示例：

    TIDB_DATABASE_URL="mysql+pymysql://<prefix>.root:<password>@gateway01.<region>.prod.aws.tidbcloud.com:4000/test?ssl_ca=/etc/ssl/cert.pem&ssl_verify_cert=true&ssl_verify_identity=true"

步骤 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 获取 embedding

定义 generate_embeddings 辅助函数，调用 Jina AI embedding 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 字段和用于存储 embedding 的向量字段 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)"

注意

  • 向量字段的维度必须与 embedding 模型生成的 embedding 维度一致。
  • 本示例中，jina-embeddings-v2-base-en 模型生成的 embedding 维度为 768。

使用 Jina AI 生成 embedding 并存储到 TiDB

使用 Jina AI Embeddings API 为每条文本生成 embedding，并将 embedding 存储到 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 embedding 进行语义搜索

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

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}')

另请参阅