📣

TiDB Cloud Essential 开放公测中。此页面由 AI 自动翻译，英文原文请见此处。

Gemini 向量嵌入

本文档介绍如何在 TiDB Cloud 中结合 Auto Embedding 使用 Gemini 向量嵌入模型，通过文本查询实现语义搜索。

注意

Auto Embedding 仅适用于托管在 AWS 上的 TiDB Cloud Starter 集群。

可用模型

如果你自带 Gemini API 密钥（BYOK），则所有 Gemini 模型均可通过 gemini/ 前缀使用。例如：

gemini-embedding-001

名称：gemini/gemini-embedding-001
维度：128–3072（默认：3072）
距离度量：Cosine，L2
最大输入文本 token 数：2,048
价格：由 Google 收费
由 TiDB Cloud 托管：❌
支持 Bring Your Own Key（BYOK，由用户自行提供 API 密钥）：✅

完整可用模型列表请参见 Gemini 文档。

使用示例

以下示例展示如何创建向量表、插入文档，并使用 Google Gemini 向量嵌入模型进行相似度搜索。

步骤 1：连接数据库

from pytidb import TiDBClient

tidb_client = TiDBClient.connect(
    host="{gateway-region}.prod.aws.tidbcloud.com",
    port=4000,
    username="{prefix}.root",
    password="{password}",
    database="{database}",
    ensure_db=True,
)

mysql -h {gateway-region}.prod.aws.tidbcloud.com \
    -P 4000 \
    -u {prefix}.root \
    -p{password} \
    -D {database}

步骤 2：配置 API 密钥

在 Google AI Studio 创建你的 API 密钥，并自带密钥（BYOK）以使用向量嵌入服务。

使用 TiDB Client 配置 Google Gemini 向量嵌入提供方的 API 密钥：

tidb_client.configure_embedding_provider(
    provider="google_gemini",
    api_key="{your-google-api-key}",
)

通过 SQL 设置 Google Gemini 向量嵌入提供方的 API 密钥：

SET @@GLOBAL.TIDB_EXP_EMBED_GEMINI_API_KEY = "{your-google-api-key}";

步骤 3：创建向量表

创建一个包含向量字段的表，使用 gemini-embedding-001 模型生成 3072 维（默认）向量：

from pytidb.schema import TableModel, Field
from pytidb.embeddings import EmbeddingFunction
from pytidb.datatype import TEXT

class Document(TableModel):
    __tablename__ = "sample_documents"
    id: int = Field(primary_key=True)
    content: str = Field(sa_type=TEXT)
    embedding: list[float] = EmbeddingFunction(
        model_name="gemini-embedding-001"
    ).VectorField(source_field="content")

table = tidb_client.create_table(schema=Document, if_exists="overwrite")

CREATE TABLE sample_documents (
    `id`        INT PRIMARY KEY,
    `content`   TEXT,
    `embedding` VECTOR(3072) GENERATED ALWAYS AS (EMBED_TEXT(
        "gemini-embedding-001",
        `content`
    )) STORED
);

步骤 4：向表中插入数据

使用 table.insert() 或 table.bulk_insert() API 添加数据：

documents = [
    Document(id=1, content="Java: Object-oriented language for cross-platform development."),
    Document(id=2, content="Java coffee: Bold Indonesian beans with low acidity."),
    Document(id=3, content="Java island: Densely populated, home to Jakarta."),
    Document(id=4, content="Java's syntax is used in Android apps."),
    Document(id=5, content="Dark roast Java beans enhance espresso blends."),
]
table.bulk_insert(documents)

使用 INSERT INTO 语句插入数据：

INSERT INTO sample_documents (id, content)
VALUES
    (1, "Java: Object-oriented language for cross-platform development."),
    (2, "Java coffee: Bold Indonesian beans with low acidity."),
    (3, "Java island: Densely populated, home to Jakarta."),
    (4, "Java's syntax is used in Android apps."),
    (5, "Dark roast Java beans enhance espresso blends.");

步骤 5：搜索相似文档

使用 table.search() API 进行向量搜索：

results = table.search("How to start learning Java programming?") \
    .limit(2) \
    .to_list()
print(results)

使用 VEC_EMBED_COSINE_DISTANCE 函数基于余弦距离进行向量搜索：

SELECT
    `id`,
    `content`,
    VEC_EMBED_COSINE_DISTANCE(embedding, "How to start learning Java programming?") AS _distance
FROM sample_documents
ORDER BY _distance ASC
LIMIT 2;

自定义向量嵌入维度

gemini-embedding-001 模型通过 Matryoshka Representation Learning (MRL) 支持灵活的维度。你可以在向量嵌入函数中指定所需的维度：

# 1536 维
embedding: list[float] = EmbeddingFunction(
    model_name="gemini-embedding-001",
    dimensions=1536
).VectorField(source_field="content")

# 768 维
embedding: list[float] = EmbeddingFunction(
    model_name="gemini-embedding-001",
    dimensions=768
).VectorField(source_field="content")

-- 1536 维
`embedding` VECTOR(1536) GENERATED ALWAYS AS (EMBED_TEXT(
    "gemini-embedding-001",
    `content`,
    '{"embedding_config": {"output_dimensionality": 1536}}'
)) STORED

-- 768 维
`embedding` VECTOR(768) GENERATED ALWAYS AS (EMBED_TEXT(
    "gemini-embedding-001",
    `content`,
    '{"embedding_config": {"output_dimensionality": 768}}'
)) STORED

请根据你的性能需求和存储约束选择合适的维度。更高的维度可以提升准确性，但会占用更多存储和计算资源。

选项

所有 Gemini 选项均可通过 EMBED_TEXT() 函数的 additional_json_options 参数进行设置。

示例：指定任务类型以提升质量

CREATE TABLE sample (
  `id`        INT,
  `content`   TEXT,
  `embedding` VECTOR(1024) GENERATED ALWAYS AS (EMBED_TEXT(
                "gemini/gemini-embedding-001",
                `content`,
                '{"task_type": "SEMANTIC_SIMILARITY"}'
              )) STORED
);

示例：使用不同维度

CREATE TABLE sample (
  `id`        INT,
  `content`   TEXT,
  `embedding` VECTOR(768) GENERATED ALWAYS AS (EMBED_TEXT(
                "gemini/gemini-embedding-001",
                `content`,
                '{"output_dimensionality": 768}'
              )) STORED
);

所有可用选项请参见 Gemini 文档。