透過 Vertex AI 執行向量相似度搜尋

歡迎使用 Firebase Data Connect 的向量相似度搜尋功能,這是 Firebase 實作語意搜尋功能,並與 Google Vertex AI 整合。

這項功能的核心是向量嵌入,也就是代表文字或媒體語意意義的浮點數陣列。您可以使用輸入向量嵌入執行最近鄰搜尋,找出所有語意相似的內容。Data Connect 會使用 PostgreSQL 的 pgvector 擴充功能來提供這項功能。

這項強大的語意搜尋功能可用於推薦引擎和搜尋引擎等用途。這也是生成式 AI 流程中擷取輔助產生的關鍵元件。如要進一步瞭解,請參閱 Vertex AI 說明文件

您可以使用 Vertex AI 的 Embeddings API,利用 Data Connect 內建的支援功能自動產生向量嵌入資料,也可以使用該 API 手動產生。

事前準備

  • 為專案設定 Data Connect

  • 啟用 Vertex AI API

設定

您可以選擇本地開發流程 (如果您是網頁、Kotlin Android 或 iOS 開發人員),或是 IDX 流程 (適用於網頁開發人員)。您可以使用本機資料庫或實際工作環境的 Data Connect 專案,以及 PostgreSQL 適用的 Cloud SQL 執行個體進行開發。

這些操作說明假設您已按照快速入門指南建立 Data Connect 專案。

與本機 PostgreSQL 整合

  1. 設定本機 PostgreSQL 執行個體。
  2. 授予自己「Vertex AI 使用者」這個 IAM 角色
  3. 在環境中設定 Google Cloud 應用程式預設憑證
  4. 在本機 PostgreSQL 例項中安裝 pgvector 擴充功能
  5. 按照 pgvector 存放區操作說明,使用 CREATE EXTENSION vector 啟用擴充功能。

與 IDX 整合

  1. 使用 Data Connect 範本設定 IDX 工作區。
  2. 授予自己 Vertex AI 使用者 IAM 角色
  3. 按照 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($movieData: Movie_Data! @pick(fields: ["title", "genre", "description"])) {
  movie_insert(data: {
    ...movieData,
    descriptionEmbedding_embed: {model: "textembedding-gecko@003", text: $movieData.description}
  })
}

在某些情況下,你可能需要更新電影說明和嵌入內容。

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

如要從用戶端呼叫後者變異,請按照下列步驟操作:

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

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

// Use the response

查詢

使用類似下方的查詢擷取向量嵌入資料,請注意,查詢傳回的 descriptionEmbedding 是浮點陣列,通常無法供人類閱讀。因此,Data Connect 產生的 SDK 不支援直接傳回。

您可以使用傳回的向量嵌入項目進行相似度搜尋,詳情請參閱下一節。

query getMovieDescription($id: String!) @auth(is: 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},
      method: L2,
      within: 2,
      where: {content: {ne: "No info available for this movie."}}, limit: 5)
      {
        id
        title
        description
      }
  }

呼叫相似度查詢

如要從用戶端程式碼呼叫相似度搜尋功能,請按照下列步驟操作:

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

const response = await searchMovieDescriptionUsingL2similarity({ query });

// Use the response

使用自訂嵌入

Data Connect 也讓您可以直接以 Vector 的形式使用嵌入內容,而非使用 _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 ConnectVector 資料類型會對應至 PostgreSQL 的 vector 類型,如 pgvector 擴充功能所定義。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 的嵌入 API 產生及儲存嵌入項目。這個伺服器值可用於查詢和 mutation。

相似度搜尋參數

method: COSINE|INNER_PRODUCT|L2

用於搜尋鄰近項目的距離函式。目前支援的演算法是 pgvector 搜尋演算法的子集。

within: float

限制執行最鄰近搜尋的距離。

where: FDC filter condition

請參閱結構定義、查詢和變異型指南

limit: int

要傳回的結果數量。