support grounding mode in gemini

refactoring

update request structure

Author: narengogi

src/providers/google-vertex-ai/chatComplete.ts

@@ -6,6 +6,7 @@ import {
   ContentType,
   Message,
   Params,
+  Tool,
   ToolCall,
 } from '../../types/requestBody';
 import {
@@ -36,9 +37,21 @@ import type {
   GoogleGenerateContentResponse,
   VertexLlamaChatCompleteStreamChunk,
   VertexLLamaChatCompleteResponse,
+  GoogleSearchRetrievalTool,
 } from './types';
 import { getMimeType } from './utils';
 
+const buildGoogleSearchRetrievalTool = (tool: Tool) => {
+  const googleSearchRetrievalTool: GoogleSearchRetrievalTool = {
+    googleSearchRetrieval: {},
+  };
+  if (tool.function.parameters?.dynamicRetrievalConfig) {
+    googleSearchRetrievalTool.googleSearchRetrieval.dynamicRetrievalConfig =
+      tool.function.parameters.dynamicRetrievalConfig;
+  }
+  return googleSearchRetrievalTool;
+};
+
 export const VertexGoogleChatCompleteConfig: ProviderConfig = {
   // https://cloud.google.com/vertex-ai/generative-ai/docs/learn/model-versioning#gemini-model-versions
   model: {
@@ -253,12 +266,20 @@ export const VertexGoogleChatCompleteConfig: ProviderConfig = {
     default: '',
     transform: (params: Params) => {
       const functionDeclarations: any = [];
+      const tools: any = [];
       params.tools?.forEach((tool) => {
         if (tool.type === 'function') {
-          functionDeclarations.push(tool.function);
+          if (tool.function.name === 'googleSearchRetrieval') {
+            tools.push(buildGoogleSearchRetrievalTool(tool));
+          } else {
+            functionDeclarations.push(tool.function);
+          }
         }
       });
-      return { functionDeclarations };
+      if (functionDeclarations.length) {
+        tools.push({ functionDeclarations });
+      }
+      return tools;
     },
   },
   tool_choice: {
@@ -648,6 +669,9 @@ export const GoogleChatCompleteResponseTransform: (
             ...(!strictOpenAiCompliance && {
               safetyRatings: generation.safetyRatings,
             }),
+            ...(!strictOpenAiCompliance && generation.groundingMetadata
+              ? { groundingMetadata: generation.groundingMetadata }
+              : {}),
           };
         }) ?? [],
       usage: {

src/providers/google-vertex-ai/types.ts

@@ -28,6 +28,24 @@ export interface GoogleGenerateContentResponse {
       category: string;
       probability: string;
     }[];
+    groundingMetadata?: {
+      webSearchQueries?: string[];
+      searchEntryPoint?: {
+        renderedContent: string;
+      };
+      groundingSupports?: Array<{
+        segment: {
+          startIndex: number;
+          endIndex: number;
+          text: string;
+        };
+        groundingChunkIndices: number[];
+        confidenceScores: number[];
+      }>;
+      retrievalMetadata?: {
+        webDynamicRetrievalScore: number;
+      };
+    };
   }[];
   promptFeedback: {
     safetyRatings: {
@@ -90,3 +108,12 @@ export interface GoogleEmbedResponse {
     billableCharacterCount: number;
   };
 }
+
+export interface GoogleSearchRetrievalTool {
+  googleSearchRetrieval: {
+    dynamicRetrievalConfig?: {
+      mode: string;
+      dynamicThreshold?: string;
+    };
+  };
+}

src/types/requestBody.ts

@@ -301,7 +301,7 @@ export interface Tool extends AnthropicPromptCache {
   /** The name of the function. */
   type: string;
   /** A description of the function. */
-  function?: Function;
+  function: Function;
 }
 
 /**