Merge pull request #786 from narengogi/enhancement/gemini-grounding

Enhancement: support grounding mode in gemini

Author: VisargD

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';
 
+export 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: {
@@ -778,6 +802,9 @@ export const GoogleChatCompleteStreamChunkTransform: (
           ...(!strictOpenAiCompliance && {
             safetyRatings: generation.safetyRatings,
           }),
+          ...(!strictOpenAiCompliance && generation.groundingMetadata
+            ? { groundingMetadata: generation.groundingMetadata }
+            : {}),
         };
       }) ?? [],
     usage: usageMetadata,

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/providers/google/chatComplete.ts

@@ -7,6 +7,7 @@ import {
   ToolCall,
   ToolChoice,
 } from '../../types/requestBody';
+import { buildGoogleSearchRetrievalTool } from '../google-vertex-ai/chatComplete';
 import { derefer, getMimeType } from '../google-vertex-ai/utils';
 import {
   ChatCompletionResponse,
@@ -325,12 +326,20 @@ export const GoogleChatCompleteConfig: 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: {
@@ -388,6 +397,24 @@ 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: {
@@ -423,8 +450,15 @@ export const GoogleErrorResponseTransform: (
 
 export const GoogleChatCompleteResponseTransform: (
   response: GoogleGenerateContentResponse | GoogleErrorResponse,
-  responseStatus: number
-) => ChatCompletionResponse | ErrorResponse = (response, responseStatus) => {
+  responseStatus: number,
+  responseHeaders: Headers,
+  strictOpenAiCompliance: boolean
+) => ChatCompletionResponse | ErrorResponse = (
+  response,
+  responseStatus,
+  _responseHeaders,
+  strictOpenAiCompliance
+) => {
   if (responseStatus !== 200) {
     const errorResponse = GoogleErrorResponseTransform(
       response as GoogleErrorResponse
@@ -468,6 +502,9 @@ export const GoogleChatCompleteResponseTransform: (
             message: message,
             index: generation.index ?? idx,
             finish_reason: generation.finishReason,
+            ...(!strictOpenAiCompliance && generation.groundingMetadata
+              ? { groundingMetadata: generation.groundingMetadata }
+              : {}),
           };
         }) ?? [],
       usage: {
@@ -483,8 +520,15 @@ export const GoogleChatCompleteResponseTransform: (
 
 export const GoogleChatCompleteStreamChunkTransform: (
   response: string,
-  fallbackId: string
-) => string = (responseChunk, fallbackId) => {
+  fallbackId: string,
+  streamState: any,
+  strictOpenAiCompliance: boolean
+) => string = (
+  responseChunk,
+  fallbackId,
+  _streamState,
+  strictOpenAiCompliance
+) => {
   let chunk = responseChunk.trim();
   if (chunk.startsWith('[')) {
     chunk = chunk.slice(1);
@@ -541,6 +585,9 @@ export const GoogleChatCompleteStreamChunkTransform: (
             delta: message,
             index: generation.index ?? index,
             finish_reason: generation.finishReason,
+            ...(!strictOpenAiCompliance && generation.groundingMetadata
+              ? { groundingMetadata: generation.groundingMetadata }
+              : {}),
           };
         }) ?? [],
       usage: {

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;
 }
 
 /**