Vertex AI plugin

The Vertex AI plugin provides interfaces to several Google generative AI models through the Vertex AI API.

Requirements

If you want to locally run flows that use this plugin, you need the Google Cloud CLI tool installed.

Configuration

To use this plugin, import the vertexai package and call vertexai.Init():

import "github.com/firebase/genkit/go/plugins/vertexai"
if err := vertexai.Init(ctx, nil); err != nil {
    return err
}

The plugin requires you to specify your Google Cloud project ID, the region to which you want to make Vertex API requests, and your Google Cloud project credentials.

  • By default, vertexai.Init() gets your Google Cloud project ID from the GCLOUD_PROJECT environment variable.

    You can also pass this value directly:

    if err := vertexai.Init(ctx, &vertexai.Config{ProjectID: yourProjectID}); err != nil {
      return err
    }
    
  • By default, vertexai.Init() gets the Vertex AI API location from the GCLOUD_LOCATION environment variable.

    You can also pass this value directly:

    if err := vertexai.Init(ctx, &vertexai.Config{Location: "asia-south1"}); err != nil {
      return err
    }
    
  • To provide API credentials, you need to set up Google Cloud Application Default Credentials.

    1. To specify your credentials:

      • If you're running your flow from a Google Cloud environment (Cloud Functions, Cloud Run, and so on), this is set automatically.

      • On your local dev environment, do this by running:

      gcloud auth application-default login
      
    2. In addition, make sure the account is granted the Vertex AI User IAM role (roles/aiplatform.user). See the Vertex AI access control docs.

Usage

Generative models

To get a reference to a supported model, specify its identifier:

langModel := vertexai.Model("gemini-1.5-pro")

The following models are supported: gemini-1.0-pro, gemini-1.5-pro, and gemini-1.5-flash.

Model references have a Generate() method that calls the Vertex AI API:

genRes, err := langModel.Generate(ctx, ai.NewGenerateRequest(
    nil, ai.NewUserTextMessage("Tell me a joke.")), nil)
if err != nil {
    return err
}

See Generating content for more information.

Embedding models

To get a reference to a supported embedding model, specify its identifier:

embeddingModel := vertexai.Embedder("text-embedding-004")

The following models are supported: textembedding-gecko@003, textembedding-gecko@002, textembedding-gecko@001, text-embedding-004, textembedding-gecko-multilingual@001, text-multilingual-embedding-002, and multimodalembedding.

Embedder references have an Embed() method that calls the Vertex AI API:

embedRes, err := embeddingModel.Embed(ctx, &ai.EmbedRequest{
    Documents: []*ai.Document{ai.DocumentFromText(userInput, nil)},
})
if err != nil {
    return err
}

You can also pass an Embedder to an indexer's Index() method and a retriever's Retrieve() method:

if err := myIndexer.Index(ctx, &ai.IndexerRequest{Documents: docsToIndex}); err != nil {
    return err
}
retrieveRes, err := myRetriever.Retrieve(ctx, &ai.RetrieverRequest{
    Document: ai.DocumentFromText(userInput, nil),
})
if err != nil {
    return err
}

See Retrieval-augmented generation (RAG) for more information.