使用 Vertex AI 执行向量相似度搜索

欢迎使用 Firebase Data Connect 的向量相似度搜索功能，这是 Firebase 对语义搜索的实现，可与 Google Vertex AI 集成。

此功能的核心是向量嵌入，即表示文本或媒体的语义含义的浮点数数组。通过使用输入向量嵌入运行最近邻搜索，您可以找到所有语义相似的内容。Data Connect 使用 PostgreSQL 的 pgvector 扩展程序来实现此功能。

这种强大的语义搜索功能可以支持推荐引擎和搜索引擎等使用场景。它也是生成式 AI 流程中检索增强生成的关键组件。如需了解详情，请参阅 Vertex AI 文档。

您可以依赖 Data Connect 的内置支持，使用 Vertex AI 的 Embeddings API 自动生成向量嵌入，也可以使用该 API 手动生成向量嵌入。

前提条件

为您的项目设置 Data Connect。

**注意**：在完成设置流程时，请注意 Data Connect的 Vertex AI 集成仅支持某些 Cloud SQL for PostgreSQL 位置。请参阅位置列表。
启用 Vertex AI API。

**注意**：您无需创建单独的 Google Cloud 云项目或安装 Google Cloud CLI。

设置

您可以选择本地开发流程（如果您是 Web、Kotlin Android 或 iOS 开发者）或 IDX 流程（适用于 Web 开发者）。您可以使用本地数据库或生产 Data Connect 项目及其 Cloud SQL for PostgreSQL 实例进行开发。

这些说明假定您已按照快速入门指南创建您的 Data Connect 项目。

与本地 PostgreSQL 集成

设置本地 PostgreSQL 实例。
授予自己 Vertex AI User IAM 角色。
在您的环境中设置 Google Cloud 应用默认凭证。
在本地 PostgreSQL 实例中安装 pgvector 扩展程序。
按照 pgvector 代码库说明，使用 CREATE EXTENSION vector 启用该扩展程序。

与 IDX 集成

使用 Data Connect 模板设置 IDX 工作区。
授予自己 Vertex AI User IAM 角色。
按照 pgvector 代码库说明，使用 CREATE EXTENSION vector 启用该扩展程序。

设计架构

如需执行向量搜索，请在架构中添加 Vector 类型的新字段。例如，如果您想使用电影说明执行语义搜索，请添加一个字段来保存与电影说明关联的向量嵌入。在此架构中，添加了 descriptionEmbedding 来存储 description 字段的向量嵌入。

type Movie @table {
 id: ID! @col(name: "movie_id") @default(id: ID! @col(name: "movie_id") @default(expr: "uuidV4()")
 title: String!
 description: String
 descriptionEmbedding: Vector! @col(size:768)
 // ...
}

生成和检索嵌入

Data Connect 通过 _embed 服务器值提供对向量嵌入的集成支持。这会指示 Data Connect 通过在内部调用 Vertex AI 的 Embedding API 来生成向量嵌入。_embed 服务器值可用于变更和查询。

变更

通过 Data Connect 生成和存储嵌入

在向量搜索应用中，您可能希望在向数据库添加记录后立即请求生成嵌入。以下 createMovie 变更会将电影记录添加到 Movie 表，并传递具有指定嵌入 model 的电影说明。

mutation createMovie($title: String!, $description: String!) {
  movie_insert(data: {
    title: $title,
    description: $description,
    descriptionEmbedding_embed: {model: "textembedding-gecko@003", text: $description}
  })
}

在某些情况下，您可能需要更新电影说明和嵌入。

mutation updateDescription($id: String!, $description: String!) {
  movie_update(id: $id, data: {
    description: $description,
    descriptionEmbedding_embed: {model: "textembedding-gecko@003", text: $description}
  })
}

如需从客户端调用后一种变更，请执行以下操作：

import { updateMovieDescription } from 'lib/dataconnect-sdk/';

await updateMovieDescription({ id: movieId, description: description});

// Use the response

查询

使用如下所示的查询提取向量嵌入。请注意，查询返回的 descriptionEmbedding 是浮点数数组，通常不便于用户理解。因此，Data Connect 生成的 SDK 不支持直接返回它。

您可以使用返回的向量嵌入执行相似度搜索，如下一部分所述。

query getMovieDescription($id: String!) @auth(level: PUBLIC) {
 movie(id: $id)
   id
   description
   descriptionEmbedding
}

执行相似度搜索

现在，我们可以执行相似度搜索了。

对于每个 Vector 字段，Data Connect 都会生成一个实现相似度搜索的 GraphQL 函数。此生成的函数的名称为 ${pluralType}_${vectorFieldName}_similarity。它支持一些参数，如以下示例和参考列表所示。

您可以定义调用相似度搜索的 GraphQL 函数。如上所述，_embed 服务器值会指示 Data Connect 使用 Vertex AI 的 Embedding API 生成向量嵌入，在本例中，用于为用于与电影说明嵌入进行比较的搜索字符串创建嵌入。

在此示例中，相似度搜索将返回最多 5 部电影，这些电影的说明在语义上与输入查询最接近。结果集按距离升序排序，即从最近到最远。

query searchMovieDescriptionUsingL2Similarity ($query: String!) @auth(level: PUBLIC) {
    movies_descriptionEmbedding_similarity(
      compare_embed: {model: "textembedding-gecko@003", text: $query},
      where: {content: {ne: "No info available for this movie."}}, limit: 5)
      {
        id
        title
        description
      }
  }

调整相似度查询

method 和 within 等搜索参数的默认值适用于大多数使用场景。但是，如果您发现查询返回的结果过于不同，或者缺少您想要包含的结果，请尝试调整这些参数。

如需查找 within 的适当值，我们可以将 _metadata.distance 添加到所选字段，以查看每个结果与查询向量的距离。根据返回的 distance 值，我们可以设置 within 参数；只有距离小于 within 值的结果才会包含在内：

query searchMovieDescriptionUsingL2Similarity ($query: String!) @auth(level: PUBLIC) {
    movies_descriptionEmbedding_similarity(
      compare_embed: {model: "textembedding-gecko@003", text: $query},
      within: 2,
      where: {content: {ne: "No info available for this movie."}}, limit: 5)
      {
        id
        title
        description
        _metadata {
          distance
        }
      }
  }

您还可以通过设置 method 参数来尝试不同的距离函数。

query searchMovieDescriptionUsingL2Similarity ($query: String!) @auth(level: PUBLIC) {
    movies_descriptionEmbedding_similarity(
      compare_embed: {model: "textembedding-gecko@003", text: $query},
      within: .5,
      method: COSINE,
      where: {content: {ne: "No info available for this movie."}}, limit: 5)
      {
        id
        title
        description
        _metadata {
          distance
        }
      }
  }

请注意，不同的方法返回的距离值差异很大：如果您已设置 within，则需要在更改 method 后再次调整该值。

调用相似度查询

如需从客户端代码调用相似度搜索，请执行以下操作：

import { searchMovieDescriptionUsingL2similarity} from 'lib/dataconnect-sdk';

const response = await searchMovieDescriptionUsingL2similarity({ query });

// Use the response

使用自定义嵌入

Data Connect 还允许您直接将嵌入作为 Vectors 使用，而不是使用 _embed 服务器值来生成嵌入。

存储自定义嵌入

使用 Vertex Embeddings API，指定匹配模型并请求正确维度的嵌入结果。

然后，将返回的浮点数数组转换为 Vector，以传递给更新操作进行存储。

mutation updateDescription($id: String!, $description: String!, $descriptionEmbedding: Vector!) {
  movie_update(id: $id, data: {
    // title, genre...
    description: $description,
    descriptionEmbedding: $descriptionEmbedding
  })
}

使用自定义嵌入执行相似度搜索

执行相同的操作来检索搜索字词的嵌入，并将其转换为 Vectors。

然后，调用 _similarity 查询来执行每次搜索。

query searchMovieDescriptionUsingL2Similarity($compare: Vector!, $within: Float, $excludesContent: String, $limit: Int) @auth(level: PUBLIC) {
    movies_descriptionEmbedding_similarity(
      compare: $compare,
      method: L2,
      within: $within,
      where: {content: {ne: $excludesContent}}, limit: $limit)
      {
        id
        title
        description
      }
  }

部署到生产环境

部署架构和连接器

典型 Data Connect 迭代的最后一步是将您的资产部署到生产环境。

使用 firebase deploy 命令将包含 Vector 类型的架构部署到 CloudSQL 时，Firebase CLI 会采取必要的步骤，在 CloudSQL 实例上启用基于 Vertex AI 的嵌入生成。

firebase deploy --only dataconnect

如果您希望在 CloudSQL 实例中手动启用嵌入支持，或者遇到 CLI 错误，请按照以下说明操作。

向量搜索语法

架构扩展程序

Data Connect 的 Vector 数据类型会映射到 pgvector 扩展程序定义的 PostgreSQL 的 vector 类型。pgvector 的 vector 类型在 PostgreSQL 中存储为单精度浮点数数组。

在 Data Connect 中，Vector 类型表示为 JSON 数字数组。输入会被强制转换为 float32 值数组。如果强制转换失败，则会引发错误。

使用 @col 指令的 size 参数设置向量的维度。

type Question @table {
    text: String!
    category: String!
    textEmbedding: Vector! @col(size: 768)
}

size 仅支持 Vector 类型。Vector 操作（例如相似度搜索）要求所有 Vector 都具有相同的维度数。

directive @col(
  # … existing args
  """
  Defines a fixed column size for certain scalar types.

  - For Vector, size is required.
  - For all other types, size is currently unsupported and hence supplying it will result in a schema error.
  """
  size: Int
) on FIELD_DEFINITION

查询和变更的 `_embed` 服务器值

_embed

此服务器值会指示 Data Connect 服务使用 Vertex AI 的 Embedding API 生成和存储嵌入。此服务器值可用于查询和变更。

相似度搜索的参数

method: COSINE|INNER_PRODUCT|L2

用于搜索附近邻居的距离函数。目前支持的算法是 pgvector 搜索算法的子集。

within: float

对执行最近邻搜索的距离的限制。

where: FDC filter condition

请参阅架构、查询和变更指南。

limit: int

要返回的结果数。