Quick dip: Use Google Cloud Vertex AI SDK with the Vercel AI SDK

🏊‍♀️ This post is a quick dip. The opposite of a deep dive.

The Vercel AI SDK documentation includes code samples for hooking up the SDK to different providers like Anthropic, OpenAI, and others. This list includes Google, but only shows you how to use the Google AI JavaScript SDK. Here's some sample code of how to hook it up to the Vertex AI SDK:

// app/api/chat/route.ts
import { VertexAI } from "@google-cloud/vertexai";
import { GoogleGenerativeAIStream, Message, StreamingTextResponse } from "ai";

const project = process.env.GOOGLE_CLOUD_PROJECT || "";
const location = "us-central1";

const vertexAI = new VertexAI({ project: project, location: location });

// Google GenAI API can use the edge runtime. This will not work with Vertex AI
// export const runtime = 'edge';

// Instantiate Gemini models
const generativeModel = vertexAI.getGenerativeModel({
  model: "gemini-pro",
});

const buildGoogleGenAIPrompt = (messages: Message[]) => ({
  contents: messages
    .filter(
      (message) => message.role === "user" || message.role === "assistant"
    )
    .map((message) => ({
      role: message.role === "user" ? "user" : "model",
      parts: [{ text: message.content }],
    })),
});

export async function POST(req: Request) {
  // Extract the `prompt` from the body of the request
  const { messages } = await req.json();

  const geminiStream = await generativeModel.generateContentStream(
    buildGoogleGenAIPrompt(messages)
  );

  // Convert the response into a friendly text-stream
  // GoogleGenerativeAIStream class decodes/extracts the text tokens in the
  // response and then re-encodes them properly for simple consumption.
  const stream = GoogleGenerativeAIStream(geminiStream);

  // Respond with the stream
  return new StreamingTextResponse(stream);
}

To run the above code, follow the instructions for the Getting Started tutorial and use it for the app/api/chat/route.ts file. (Skip the API key step).

Notes

Some differences between this Vertex AI SDK example and the one for the Google AI JavaScript SDK:

• We don't indicate the runtime as 'edge' because it is not compatible with the Vertex AI SDK. Click here for more information on the Edge Runtime and Node.js modules.
• Instead of API keys, Vertex uses a project ID indicated via the GOOGLE_CLOUD_PROJECT environment variable and application default credentials indicated via the GOOGLE_APPLICATION_CREDENTIALS environment variable.