Perform vector similarity search with Vertex AI

Welcome to Firebase Data Connect's vector similarity search — Firebase's implementation of semantic search that integrates with Google Vertex AI.

At the core of this feature are vector embeddings, which are arrays of floating point numbers representing the semantic meaning of text or media. By running a nearest neighbor search using an input vector embedding, you can find all semantically similar content. Data Connect uses PostgreSQL's pgvector extension for this capability.

This powerful semantic search can drive use cases like recommendation engines and search engines. It's also a key component in retrieval-augmented generation in generative AI flows. The Vertex AI documentation is a great place to learn more.

You can rely on Data Connect's built-in support for generating vector embeddings automatically using Vertex AI's Embeddings API, or use that API to generate them manually.

Prerequisites

Set up Data Connect for your project.

Note: As you work through the setup flow, be aware that Data Connect's Vertex AI integration is supported only for certain Cloud SQL for PostgreSQL locations. See the location list.
Enable Vertex AI APIs.

Note: You don't need to create a separate Google Cloud project or install the Google Cloud CLI.

Setup

You can choose between a local development flow (if you're a web, Kotlin Android, or iOS developer) or an IDX flow (for web developers). You can use a local database or your production Data Connect project and its Cloud SQL for PostgreSQL instance for development.

These instructions assume you have created your Data Connect project following the quickstart guide.

Integrate with local PostgreSQL

Set up a local PostgreSQL instance.
Grant yourself the Vertex AI user IAM role.
Set up Google Cloud Application Default Credentials in your environment.
Install the pgvector extension in your local PostgreSQL instance.
Enable the extension using CREATE EXTENSION vector per the pgvector repository instructions.

Integrate with IDX

Set up your IDX workspace using the Data Connect template.
Grant yourself the Vertex AI user IAM role.
Enable the extension using CREATE EXTENSION vector per the pgvector repository instructions.

Design your schema

To perform vector search, add a new field of Vector type in your schema. For example, if you want to do a semantic search using movie descriptions, add a field to hold the vector embeddings associated with the movie description. In this schema, descriptionEmbedding is added to store vector embeddings for the description field.

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)
 // ...
}

Generate and retrieve embeddings

Data Connect brings integrated support for vector embeddings with the _embed server value. This directs Data Connect to generate vector embeddings by internally calling Vertex AI's Embedding APIs. The _embed server value can be used in both mutations and queries.

Mutations

Generate and store an embedding through Data Connect

In your vector search app, you'll likely want to request that embeddings be generated as soon as you add records to your database. Here's a createMovie mutation adds a movie record to the Movie table and also passes a movie description with a specified embedding model.

mutation createMovie($movieData: Movie_Data! @pick(fields: ["title", "genre", "description"])) {
  movie_insert(data: {
    ...movieData,
    descriptionEmbedding_embed: {model: "textembedding-gecko@003", text: $movieData.description}
  })
}

In some cases, you may want to update the movie description and embedding.

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

To call the latter mutation from a client:

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

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

// Use the response

Queries

Fetch vector embeddings using a query like the following. Note that the descriptionEmbedding returned by the query is an array of floats, which is typically not human-readable. Thus, Data Connect generated SDKs don't support returning it directly.

You can use returned vector embeddings to do similarity search, as described in the next section.

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

Perform similarity search

Now we can perform similarity search.

For each Vector field, Data Connect generates a GraphQL function that implements the similarity search. The name of this generated function is ${pluralType}_${vectorFieldName}_similarity. It supports a few parameters as shown in the following examples and in the reference list.

You can define a GraphQL function that invokes the similarity search. As mentioned above, the _embed server value directs Data Connect to generate the vector embeddings using Vertex AI's Embedding APIs, in this case to create embeddings for the search string used for comparison with movie description embeddings.

In this example, the similarity search will return up to 5 movies whose description is semantically closest to the input query. The result set is sorted in the ascending order of the distance - closest to furthest.

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

Call the similarity query

To call a similarity search from client code:

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

const response = await searchMovieDescriptionUsingL2similarity({ query });

// Use the response

Use custom embeddings

Data Connect also lets you work with embeddings directly as Vectors rather than using the _embed server value to generate them.

Store a custom embedding

Using the Vertex Embeddings API, specify a matching model and request embedding results of the correct dimension.

Then, cast the returned array of floats to a Vector to pass to the update operation for storage.

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

Perform similarity search using custom embeddings

Carry out the same operation to retrieve embeddings for search terms and cast them to Vectors.

Then, call the _similarity query to perform each search.

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

Deploy to production

Deploy your schema and connector

The last step in a typical Data Connect iteration is to deploy your assets to production.

When deploying your schema containing Vector types to CloudSQL using the firebase deploy command, the Firebase CLI takes the necessary steps to enable Vertex AI-based embedding generation on your CloudSQL instance.

firebase deploy --only dataconnect

if you wish to enable embedding support in your CloudSQL instance manually, or encounter a CLI error, follow these instructions.

Vector search syntax

Schema extensions

Data Connect's Vector data type maps to PostgreSQL's vector type as defined by the pgvector extension. pgvector's vector type is stored as an array of single precision floating point numbers in PostgreSQL.

In Data Connect, the Vector type is represented as an array of JSON numbers. Inputs are coerced into an array of float32 values. If coercion fails, an error is raised.

Use the size parameter of the @col directive to set the dimensions of the vector.

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

size is only supported for Vector types. Vector operations, such as similarity-search, necessitate that all the Vectors have the same number of dimensions.

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` server value for queries and mutations

_embed

This server value directs the Data Connect service to generate and store embeddings using Vertex AI's Embedding APIs. This server value can be used on both queries and mutations.

Parameters For similarity search

method: COSINE|INNER_PRODUCT|L2

The distance function used to search for nearby neighbors. Currently-supported algorithms are a subset of pgvector search algorithms.

within: float

A constraint on the distance within which the nearest neighbor search is performed.

where: FDC filter condition

See the schemas, queries and mutations guide.

limit: int

The number of results to return.