Support Vertex AI Gemini grounding with Google Search and Vertex data stores

[chrisraygill]

Is your feature request related to a problem? Please describe.

Vertex AI offers the ability to ground responses for Gemini models using Google Search or your own data in within Vertex AI Search datas stores.

It appears to supported in the Python and Node.js SDK, but not yet in the Go SDK.

In the Vertex AI Gemini API, it works by passing in the grounding mechanism as a tool to the model, like so (Node.js):

async function generateContentWithVertexAISearchGrounding() {
  const generativeModelPreview = vertexAI.preview.getGenerativeModel({
    model: textModel,
    // The following parameters are optional
    // They can also be passed to individual content generation requests
    safetySettings: [{category: HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT, threshold: HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE}],
    generationConfig: {maxOutputTokens: 256},
  });

  const vertexAIRetrievalTool = {
    retrieval: {
      vertexAiSearch: {
        datastore: 'projects/.../locations/.../collections/.../dataStores/...',
      },
      disableAttribution: false,
    },
  };
  const result = await generativeModelPreview.generateContent({
    contents: [{role: 'user', parts: [{text: 'Why is the sky blue?'}]}],
    tools: [vertexAIRetrievalTool],
  })
  const response = result.response;
  const groundingMetadata = response.candidates[0].groundingMetadata;
  console.log("Grounding metadata is: ", JSON.stringify(groundingMetadata));
}
generateContentWithVertexAISearchGrounding();

I'm not sure how this would work with Genkit's current design for specifying tools in a generation request.

[chrisraygill]

@pavelgj from a design perspective, any thoughts on how we might do this?

Maybe just making them custom configuration parameters for the Vertex Gemini models rather than tools? Not sure if that would have odd downstream effects when it comes to tool handling stuff like forced function calling.

[chrisraygill]

@pr-Mais any proposals from you on how this might look?

[pr-Mais]

@chrisraygill I was thinking, we could turn it into a generic tool, as a start it will only support vertex ai models through this API, later we add a custom implementation for it to be used with other models.

Update: after looking deeper into genkit's tools, this could be tricky. As you mentioned, a better solution is to add it as a custom param in the generate request specific to vertexai plugin. Something like:

const GeminiConfigSchema = GenerationCommonConfigSchema.extend({
  safetySettings: z.array(SafetySettingsSchema).optional(),
  vertexRetriever: z
    .object({
      vertexAiSearch: z
        .object({
          datastore: z.string(),
        })
        .optional(),
      googleSearchRetrieval: z.boolean().optional(),
      disableAttribution: z.boolean().optional(),
    })
    .optional(),
});