使用向量嵌入搜索

本页面介绍如何通过以下方法使用 Firestore 执行 K 最近邻 (KNN) 向量搜索：

  • 存储向量值
  • 创建和管理 KNN 向量索引
  • 使用一个支持的向量距离衡量方式执行 K 最近邻 (KNN) 查询

准备工作

在将嵌入存储到 Firestore 中之前，您必须生成向量嵌入。Firestore 不会生成嵌入。您可以使用 Vertex AI 等服务从 Firestore 数据创建向量值（例如文本嵌入）。然后，您可以将这些嵌入重新存储在 Firestore 文档中。

如需详细了解嵌入，请参阅什么是嵌入？

如需了解如何使用 Vertex AI 获取文本嵌入，请参阅获取文本嵌入。

存储向量嵌入

以下示例演示了如何在 Firestore 中存储向量嵌入。

通过向量嵌入执行写入操作

以下示例展示了如何将向量嵌入存储在 Firestore 文档中：

Python
from google.cloud import firestore
from google.cloud.firestore_v1.vector import Vector

firestore_client = firestore.Client()
collection = firestore_client.collection("coffee-beans")
doc = {
    "name": "Kahawa coffee beans",
    "description": "Information about the Kahawa coffee beans.",
    "embedding_field": Vector([0.18332680, 0.24160706, 0.3416704]),
}

collection.add(doc)
Node.js
import {
  Firestore,
  FieldValue,
} from "@google-cloud/firestore";

const db = new Firestore();
const coll = db.collection('coffee-beans');
await coll.add({
  name: "Kahawa coffee beans",
  description: "Information about the Kahawa coffee beans.",
  embedding_field: FieldValue.vector([1.0 , 2.0, 3.0])
});
Go
import (
	"context"
	"fmt"
	"io"

	"cloud.google.com/go/firestore"
)

type CoffeeBean struct {
	Name           string             `firestore:"name,omitempty"`
	Description    string             `firestore:"description,omitempty"`
	EmbeddingField firestore.Vector32 `firestore:"embedding_field,omitempty"`
	Color          string             `firestore:"color,omitempty"`
}

func storeVectors(w io.Writer, projectID string) error {
	ctx := context.Background()

	// Create client
	client, err := firestore.NewClient(ctx, projectID)
	if err != nil {
		return fmt.Errorf("firestore.NewClient: %w", err)
	}
	defer client.Close()

	// Vector can be represented by Vector32 or Vector64
	doc := CoffeeBean{
		Name:           "Kahawa coffee beans",
		Description:    "Information about the Kahawa coffee beans.",
		EmbeddingField: []float32{1.0, 2.0, 3.0},
		Color:          "red",
	}
	ref := client.Collection("coffee-beans").NewDoc()
	if _, err = ref.Set(ctx, doc); err != nil {
		fmt.Fprintf(w, "failed to upsert: %v", err)
		return err
	}

	return nil
}
Java
import com.google.cloud.firestore.CollectionReference;
import com.google.cloud.firestore.DocumentReference;
import com.google.cloud.firestore.FieldValue;
import com.google.cloud.firestore.VectorQuery;

CollectionReference coll = firestore.collection("coffee-beans");

Map<String, Object> docData = new HashMap<>();
docData.put("name", "Kahawa coffee beans");
docData.put("description", "Information about the Kahawa coffee beans.");
docData.put("embedding_field", FieldValue.vector(new double[] {1.0, 2.0, 3.0}));

ApiFuture<DocumentReference> future = coll.add(docData);
DocumentReference documentReference = future.get();

使用 Cloud Functions 函数计算向量嵌入

如要在文档每次创建或更新时便计算并存储相应的向量嵌入，您可以设置一个 Cloud Run 函数：

Python
@functions_framework.cloud_event
def store_embedding(cloud_event) -> None:
  """Triggers by a change to a Firestore document.
  """
  firestore_payload = firestore.DocumentEventData()
  payload = firestore_payload._pb.ParseFromString(cloud_event.data)

  collection_id, doc_id = from_payload(payload)
  # Call a function to calculate the embedding
  embedding = calculate_embedding(payload)
  # Update the document
  doc = firestore_client.collection(collection_id).document(doc_id)
  doc.set({"embedding_field": embedding}, merge=True)
Node.js
/**
 * A vector embedding will be computed from the
 * value of the `content` field. The vector value
 * will be stored in the `embedding` field. The
 * field names `content` and `embedding` are arbitrary
 * field names chosen for this example.
 */
async function storeEmbedding(event: FirestoreEvent<any>): Promise<void> {
  // Get the previous value of the document's `content` field.
  const previousDocumentSnapshot = event.data.before as QueryDocumentSnapshot;
  const previousContent = previousDocumentSnapshot.get("content");

  // Get the current value of the document's `content` field.
  const currentDocumentSnapshot = event.data.after as QueryDocumentSnapshot;
  const currentContent = currentDocumentSnapshot.get("content");

  // Don't update the embedding if the content field did not change
  if (previousContent === currentContent) {
    return;
  }

  // Call a function to calculate the embedding for the value
  // of the `content` field.
  const embeddingVector = calculateEmbedding(currentContent);

  // Update the `embedding` field on the document.
  await currentDocumentSnapshot.ref.update({
    embedding: embeddingVector,
  });
}
Go
  // Not yet supported in the Go client library
Java
  // Not yet supported in the Java client library

创建和管理向量索引

您必须先创建相应的索引，然后才能通过向量嵌入执行最近邻搜索。以下示例演示了如何使用 Google Cloud CLI 和控制台创建和管理向量索引。您还可以使用 Firebase CLI 和 Terraform 管理向量索引。

创建向量索引

Google Cloud 控制台

如需在 Google Cloud 控制台中手动创建新索引，请执行以下操作：

  1. 在 Google Cloud 控制台中，前往数据库页面。

    前往“数据库”

  2. 从数据库列表中选择所需的数据库。
  3. 在导航菜单中，点击索引，然后点击手动标签页。
  4. 点击创建索引。

    如需为向量搜索将向量字段编入索引，请选择创建向量索引。

  5. 输入集合 ID。 输入向量字段路径和向量嵌入维度数。添加要编入索引的任何其他字段的名称，并为每个字段添加索引模式。

    点击保存索引。

您的新索引将显示在手动索引列表中，并且 Firestore 将开始创建索引。索引创建完成后，您会在索引旁边看到一个绿色对勾标记。

gcloud

在创建向量索引之前，请先升级到最新版本的 Google Cloud CLI：

gcloud components update

如需创建向量索引，请使用 gcloud firestore indexes composite create：

gcloud firestore indexes composite create \
--collection-group=collection-group \
--query-scope=COLLECTION \
--field-config field-path=vector-field,vector-config='vector-configuration' \
--database=database-id

其中：

  • collection-group 是集合组的 ID。
  • vector-field 是包含向量嵌入的字段的名称。
  • database-id 是相应数据库的 ID。
  • vector-configuration 包含向量 dimension 和索引类型。dimension 是一个不超过 2,048 的整数。索引类型必须为 flat。按如下方式设置索引配置的格式：{"dimension":"DIMENSION", "flat": "{}"}。

以下示例创建了一个复合索引，其中包含字段 vector-field 的向量索引和字段 color 的升序索引。您可以在执行最近邻搜索之前使用此类索引预先过滤数据。

gcloud firestore indexes composite create \
--collection-group=collection-group \
--query-scope=COLLECTION \
--field-config=order=ASCENDING,field-path="color" \
--field-config field-path=vector-field,vector-config='{"dimension":"1024", "flat": "{}"}' \
--database=database-id

列出所有向量索引

Google Cloud 控制台

  1. 在 Google Cloud 控制台中，前往数据库页面。

    前往“数据库”

  2. 从数据库列表中选择所需的数据库。

  3. 在导航菜单中，点击索引，然后点击手动标签页。

    索引表会列出数据库的所有索引。向量索引包含带有 图标的向量字段。

gcloud

如需列出所有索引并检索索引 ID，请运行以下命令：

gcloud firestore indexes composite list --database=database-id

将 database-id 替换为相应数据库的 ID。

您可以使用索引 ID 查看有关索引的更多详细信息：

gcloud firestore indexes composite describe index-id --database=database-id

其中：

  • index-id 是要描述的索引的 ID。
  • database-id 是相应数据库的 ID。

删除矢量索引

Google Cloud 控制台

  1. 在 Google Cloud 控制台中，前往数据库页面。

    前往“数据库”

  2. 从数据库列表中选择所需的数据库。

  3. 在导航菜单中，点击索引，然后点击手动标签页。

  4. 在手动索引列表中，针对要删除的索引点击更多按钮 。点击删除。
  5. 点击提醒中的删除索引，确认您要删除此索引。
gcloud
gcloud firestore indexes composite delete index-id --database=database-id

其中：

  • index-id 是要删除的索引的 ID。可使用 indexes composite list 检索索引 ID。
  • database-id 是相应数据库的 ID。

执行最近邻查询

您可以执行相似度搜索来查找向量嵌入的最近邻。相似度搜索需要使用向量索引。如果没有现成的索引，Firestore 会使用 gcloud CLI 建议一个可创建的索引。

以下示例会查找查询向量的 10 个最近邻。

Python
from google.cloud.firestore_v1.base_vector_query import DistanceMeasure
from google.cloud.firestore_v1.vector import Vector

collection = db.collection("coffee-beans")

# Requires a single-field vector index
vector_query = collection.find_nearest(
    vector_field="embedding_field",
    query_vector=Vector([0.3416704,