Enable google search / retrieval for google endpoints

api/app/clients/GoogleClient.js

@@ -4,7 +4,7 @@ const { ChatVertexAI } = require('@langchain/google-vertexai');
 const { GoogleVertexAI } = require('@langchain/google-vertexai');
 const { ChatGoogleVertexAI } = require('@langchain/google-vertexai');
 const { ChatGoogleGenerativeAI } = require('@langchain/google-genai');
-const { GoogleGenerativeAI: GenAI } = require('@google/generative-ai');
+const { GoogleGenerativeAI: GenAI, DynamicRetrievalMode } = require('@google/generative-ai');
 const { AIMessage, HumanMessage, SystemMessage } = require('@langchain/core/messages');
 const { encoding_for_model: encodingForModel, get_encoding: getEncoding } = require('tiktoken');
 const {
@@ -36,6 +36,12 @@ const tokenizersCache = {};
 const settings = endpointSettings[EModelEndpoint.google];
 const EXCLUDED_GENAI_MODELS = /gemini-(?:1\.0|1-0|pro)/;
 
+
+const isNewGeminiModel = (model) => {
+  return model.includes('gemini-exp-1206') || 
+         model.includes('gemini-2.');
+};
+
 class GoogleClient extends BaseClient {
   constructor(credentials, options = {}) {
     super('apiKey', options);
@@ -124,7 +130,11 @@ class GoogleClient extends BaseClient {
       .filter((ex) => ex)
       .filter((obj) => obj.input.content !== '' && obj.output.content !== '');
 
-    this.modelOptions = this.options.modelOptions || {};
+    // Set modelOptions, ensuring enableSearch is included
+    this.modelOptions = {
+      ...(this.options.modelOptions || {}),
+      enableSearch: this.options.enableSearch,
+    };
 
     this.options.attachments?.then((attachments) => this.checkVisionRequest(attachments));
 
@@ -420,9 +430,9 @@ class GoogleClient extends BaseClient {
 
     logger.debug('[GoogleClient]', {
       orderedMessages,
-      parentMessageId,
+      parentMessageId, 
     });
-
+      
     const formattedMessages = orderedMessages.map((message) => ({
       author: message.isCreatedByUser ? this.userLabel : this.modelLabel,
       content: message?.content ?? message.text,
@@ -624,7 +634,29 @@ class GoogleClient extends BaseClient {
       return new ChatVertexAI(clientOptions);
     } else if (!EXCLUDED_GENAI_MODELS.test(model)) {
       logger.debug('Creating GenAI client');
-      return new GenAI(this.apiKey).getGenerativeModel({ ...clientOptions, model }, requestOptions);
+      const tools = [];
+      if (this.modelOptions.enableSearch) {
+        logger.debug('[GoogleClient] Adding search tool');
+        if (isNewGeminiModel(model)) {
+          tools.push({
+            googleSearch: {}
+          });
+        } else {
+          tools.push({
+            googleSearchRetrieval: {
+              dynamicRetrievalConfig: {
+                mode: DynamicRetrievalMode.MODE_DYNAMIC,
+                dynamicThreshold: 0.7,
+              },
+            },
+          });
+        }
+      }
+      return new GenAI(this.apiKey).getGenerativeModel({
+        ...clientOptions,
+        model,
+        tools,
+      }, requestOptions );
     }
 
     logger.debug('Creating Chat Google Generative AI client');

api/server/services/Endpoints/google/build.js

@@ -11,8 +11,10 @@ const buildOptions = (endpoint, parsedBody) => {
     greeting,
     spec,
     artifacts,
+    enableSearch,
     ...modelOptions
   } = parsedBody;
+
   const endpointOption = removeNullishValues({
     examples,
     endpoint,
@@ -22,6 +24,7 @@ const buildOptions = (endpoint, parsedBody) => {
     iconURL,
     greeting,
     spec,
+    enableSearch,
     modelOptions,
   });
 

client/src/components/Endpoints/Settings/Google.tsx

@@ -5,6 +5,7 @@ import {
   Input,
   Label,
   Slider,
+  Switch,
   HoverCard,
   InputNumber,
   SelectDropDown,
@@ -28,6 +29,7 @@ export default function Settings({ conversation, setOption, models, readonly }:
     topK,
     maxContextTokens,
     maxOutputTokens,
+    enableSearch,
   } = conversation ?? {};
 
   const [setMaxContextTokens, maxContextTokensValue] = useDebouncedInput<number | null | undefined>(
@@ -49,6 +51,7 @@ export default function Settings({ conversation, setOption, models, readonly }:
   const setTopP = setOption('topP');
   const setTopK = setOption('topK');
   const setMaxOutputTokens = setOption('maxOutputTokens');
+  const setEnableSearch = setOption('enableSearch');
 
   return (
     <div className="grid grid-cols-5 gap-6">
@@ -100,6 +103,28 @@ export default function Settings({ conversation, setOption, models, readonly }:
         </div>
       </div>
       <div className="col-span-5 flex flex-col items-center justify-start gap-6 px-3 sm:col-span-2">
+        <HoverCard openDelay={300}>
+          <HoverCardTrigger className="grid w-full items-center gap-2">
+            <div className="flex justify-between">
+              <Label htmlFor="enable-search" className="text-left text-sm font-medium">
+                Enable Search{' '}
+                <small className="opacity-40">
+                  ({localize('com_endpoint_default')}: {google.enableSearch.default ? 'Yes' : 'No'})
+                </small>
+              </Label>
+              <Switch
+                id="enable-search"
+                disabled={readonly}
+                checked={enableSearch ?? google.enableSearch.default}
+                onCheckedChange={(checked) => setEnableSearch(checked)}
+              />
+            </div>
+          </HoverCardTrigger>
+          <OptionHoverAlt
+            description="Enable Google Search functionality for enhanced responses"
+            side={ESide.Left}
+          />
+        </HoverCard>
         <HoverCard openDelay={300}>
           <HoverCardTrigger className="grid w-full items-center gap-2">
             <div className="mt-1 flex w-full justify-between">

packages/data-provider/src/schemas.ts

@@ -240,6 +240,9 @@ export const googleSettings = {
     step: 1,
     default: 40,
   },
+  enableSearch: {
+    default: false,
+  },
 };
 
 const ANTHROPIC_MAX_OUTPUT = 8192;
@@ -551,6 +554,7 @@ export const tConversationSchema = z.object({
   /* google */
   context: z.string().nullable().optional(),
   examples: z.array(tExampleSchema).optional(),
+  enableSearch: z.boolean().optional(),
   /* DB */
   tags: z.array(z.string()).optional(),
   createdAt: z.string(),
@@ -686,6 +690,7 @@ export const googleSchema = tConversationSchema
     greeting: true,
     spec: true,
     maxContextTokens: true,
+    enableSearch: true,
   })
   .transform((obj) => {
     return {
@@ -698,6 +703,7 @@ export const googleSchema = tConversationSchema
       maxOutputTokens: obj.maxOutputTokens ?? google.maxOutputTokens.default,
       topP: obj.topP ?? google.topP.default,
       topK: obj.topK ?? google.topK.default,
+      enableSearch: obj.enableSearch ?? google.enableSearch.default,
       iconURL: obj.iconURL ?? undefined,
       greeting: obj.greeting ?? undefined,
       spec: obj.spec ?? undefined,
@@ -713,6 +719,7 @@ export const googleSchema = tConversationSchema
     maxOutputTokens: google.maxOutputTokens.default,
     topP: google.topP.default,
     topK: google.topK.default,
+    enableSearch: google.enableSearch.default,
     iconURL: undefined,
     greeting: undefined,
     spec: undefined,
@@ -986,6 +993,7 @@ export const compactGoogleSchema = tConversationSchema
     greeting: true,
     spec: true,
     maxContextTokens: true,
+    enableSearch: true,
   })
   .transform((obj) => {
     const newObj: Partial<TConversation> = { ...obj };
@@ -1001,6 +1009,9 @@ export const compactGoogleSchema = tConversationSchema
     if (newObj.topK === google.topK.default) {
       delete newObj.topK;
     }
+    if (newObj.enableSearch === google.enableSearch.default) {
+      delete newObj.enableSearch;
+    }
 
     return removeNullishValues(newObj);
   })
@@ -1118,3 +1129,4 @@ export const compactAgentsSchema = tConversationSchema
   })
   .transform(removeNullishValues)
   .catch(() => ({}));
+