Enable google search for gemini-1.5-pro and gemini-2.0-flash-exp

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 {
@@ -27,15 +27,18 @@ const {
   truncateText,
 } = require('./prompts');
 const BaseClient = require('./BaseClient');
-
 const loc = process.env.GOOGLE_LOC || 'us-central1';
 const publisher = 'google';
 const endpointPrefix = `${loc}-aiplatform.googleapis.com`;
 const tokenizersCache = {};
-
 const settings = endpointSettings[EModelEndpoint.google];
 const EXCLUDED_GENAI_MODELS = /gemini-(?:1\.0|1-0|pro)/;
 
+
+const isNewGeminiModel = (model) => {
+  return model.includes('gemini-2.');
+};
+
 class GoogleClient extends BaseClient {
   constructor(credentials, options = {}) {
     super('apiKey', options);
@@ -124,7 +127,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 +427,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 +631,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');
@@ -711,17 +740,48 @@ class GoogleClient extends BaseClient {
 
       const delay = modelName.includes('flash') ? 8 : 15;
       const result = await client.generateContentStream(requestOptions);
+      let lastGroundingMetadata = null;
+
       for await (const chunk of result.stream) {
-        const chunkText = chunk.text();
-        await this.generateTextStream(chunkText, onProgress, {
-          delay,
-        });
-        reply += chunkText;
-        await sleep(streamRate);
+        // Get the text content from the first candidate's content parts
+        const text = chunk.candidates?.[0]?.content?.parts?.[0]?.text ?? '';
+
+        // Store the grounding metadata from the last chunk that has it
+        if (chunk.candidates?.[0]?.groundingMetadata) {
+          lastGroundingMetadata = chunk.candidates[0].groundingMetadata;
+        }
+
+        // Only send text content if there is any
+        if (text) {
+          await this.generateTextStream(text, onProgress, {
+            delay,
+            metadata: lastGroundingMetadata ? { groundingMetadata: lastGroundingMetadata } : undefined
+          });
+          reply += text;
+          await sleep(streamRate);
+        }
       }
-      return reply;
-    }
 
+      // Send final completion message with metadata
+      const finalMessage = {
+        text: reply,
+        isComplete: true,
+        metadata: lastGroundingMetadata ? { groundingMetadata: lastGroundingMetadata } : undefined
+      };
+
+      await onProgress(finalMessage);
+
+      // Set metadata for BaseClient to save
+      if (lastGroundingMetadata) {
+        this.metadata = { groundingMetadata: lastGroundingMetadata };
+      }
+
+      return {
+        text: reply,
+        groundingMetadata: lastGroundingMetadata
+      };
+
+    }
     const stream = await model.stream(messages, {
       signal: abortController.signal,
       safetySettings: _payload.safetySettings,
@@ -890,11 +950,26 @@ class GoogleClient extends BaseClient {
   async sendCompletion(payload, opts = {}) {
     payload.safetySettings = this.getSafetySettings();
 
-    let reply = '';
-    reply = await this.getCompletion(payload, opts);
-    return reply.trim();
+    const response = await this.getCompletion(payload, opts);
+
+    // Handle both string and object responses
+    if (typeof response === 'string') {
+      return response.trim();
+    }
+
+    // If response is an object with text and metadata
+    if (response && typeof response === 'object') {
+      const { text, groundingMetadata } = response;
+      if (groundingMetadata) {
+        this.metadata = { groundingMetadata };
+      }
+      return text.trim();
+    }
+
+    return '';
   }
 
+
   getSafetySettings() {
     return [
       {
@@ -943,3 +1018,4 @@ class GoogleClient extends BaseClient {
 }
 
 module.exports = GoogleClient;
+

api/app/clients/TextStream.js

@@ -9,6 +9,7 @@ class TextStream extends Readable {
     this.minChunkSize = options.minChunkSize ?? 2;
     this.maxChunkSize = options.maxChunkSize ?? 4;
     this.delay = options.delay ?? 20; // Time in milliseconds
+    this.metadata = options.metadata;
   }
 
   _read() {
@@ -35,7 +36,13 @@ class TextStream extends Readable {
   async processTextStream(onProgressCallback) {
     const streamPromise = new Promise((resolve, reject) => {
       this.on('data', (chunk) => {
-        onProgressCallback(chunk.toString());
+        const payload = {
+          text: chunk.toString(),
+        };
+        if (this.metadata) {
+          payload.metadata = this.metadata;
+        }
+        onProgressCallback(payload);
       });
 
       this.on('end', () => {

api/models/Message.js

@@ -204,6 +204,7 @@ async function updateMessage(req, message, metadata) {
       isCreatedByUser: updatedMessage.isCreatedByUser,
       tokenCount: updatedMessage.tokenCount,
       isEdited: true,
+      groundingMetadata: updatedMessage.groundingMetadata,
     };
   } catch (err) {
     logger.error('Error updating message:', err);

api/models/schema/defaults.js

@@ -130,6 +130,11 @@ const conversationPreset = {
   max_tokens: {
     type: Number,
   },
+  // for google only
+  enableSearch: {
+    type: Boolean,
+    required: false,
+  },
 };
 
 const agentOptions = {

api/models/schema/messageSchema.js

@@ -116,6 +116,7 @@ const messageSchema = mongoose.Schema(
       type: String,
     },
     attachments: { type: [{ type: mongoose.Schema.Types.Mixed }], default: undefined },
+    groundingMetadata: { type: mongoose.Schema.Types.Mixed, default: undefined },
     /*
     attachments: {
       type: [

api/server/controllers/AskController.js

@@ -28,6 +28,7 @@ const AskController = async (req, res, next, initializeClient, addTitle) => {
   let promptTokens;
   let userMessageId;
   let responseMessageId;
+  let currentMetadata = null;
   const sender = getResponseSender({
     ...endpointOption,
     model: endpointOption.modelOptions.model,
@@ -60,7 +61,7 @@ const AskController = async (req, res, next, initializeClient, addTitle) => {
     const messageCache = getLogStores(CacheKeys.MESSAGES);
     const { onProgress: progressCallback, getPartialText } = createOnProgress({
       onProgress: throttle(
-        ({ text: partialText }) => {
+        ({ text: partialText, metadata }) => {
           /*
               const unfinished = endpointOption.endpoint === EModelEndpoint.google ? false : true;
           messageCache.set(responseMessageId, {
@@ -76,6 +77,9 @@ const AskController = async (req, res, next, initializeClient, addTitle) => {
           }, Time.FIVE_MINUTES);
           */
 
+          if (metadata) {
+            currentMetadata = metadata;
+          }
           messageCache.set(responseMessageId, partialText, Time.FIVE_MINUTES);
         },
         3000,
@@ -94,6 +98,7 @@ const AskController = async (req, res, next, initializeClient, addTitle) => {
       text: getPartialText(),
       userMessage,
       promptTokens,
+      ...(currentMetadata ? { metadata: currentMetadata } : {}),
     });
 
     const { abortController, onStart } = createAbortController(req, res, getAbortData, getReqData);
@@ -131,6 +136,15 @@ const AskController = async (req, res, next, initializeClient, addTitle) => {
     let response = await client.sendMessage(text, messageOptions);
     response.endpoint = endpointOption.endpoint;
 
+    // Add metadata to the final response if available
+    if (currentMetadata) {
+      if (currentMetadata.metadata?.groundingMetadata) {
+        response.groundingMetadata = currentMetadata.metadata.groundingMetadata;
+      } else if (currentMetadata.groundingMetadata) {
+        response.groundingMetadata = currentMetadata.groundingMetadata;
+      }
+    }
+
     const { conversation = {} } = await client.responsePromise;
     conversation.title =
       conversation && !conversation.title ? null : conversation?.title || 'New Chat';
@@ -142,19 +156,21 @@ const AskController = async (req, res, next, initializeClient, addTitle) => {
     }
 
     if (!abortController.signal.aborted) {
+      const finalResponse = { ...response };
+
       sendMessage(res, {
         final: true,
         conversation,
         title: conversation.title,
         requestMessage: userMessage,
-        responseMessage: response,
+        responseMessage: finalResponse,
       });
       res.end();
 
       if (!client.savedMessageIds.has(response.messageId)) {
         await saveMessage(
           req,
-          { ...response, user },
+          { ...finalResponse, user },
           { context: 'api/server/controllers/AskController.js - response end' },
         );
       }

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,
   });
 

api/server/utils/handleText.js

@@ -21,13 +21,32 @@ const base = { message: true, initial: true };
 const createOnProgress = ({ generation = '', onProgress: _onProgress }) => {
   let i = 0;
   let tokens = addSpaceIfNeeded(generation);
+  let currentMetadata = null;
 
   const basePayload = Object.assign({}, base, { text: tokens || '' });
 
   const progressCallback = (chunk, { res, ...rest }) => {
-    basePayload.text = basePayload.text + chunk;
+    // If chunk is an object with text and metadata
+    if (typeof chunk === 'object' && chunk.text !== undefined) {
+      basePayload.text = basePayload.text + chunk.text;
+      if (chunk.metadata) {
+        currentMetadata = chunk.metadata;
+      }
+    } else {
+      // Handle plain text chunks
+      basePayload.text = basePayload.text + chunk;
+    }
 
+    // Always include current metadata in the payload if available
     const payload = Object.assign({}, basePayload, rest);
+    if (currentMetadata) {
+      if (currentMetadata.metadata?.groundingMetadata) {
+        payload.groundingMetadata = currentMetadata.metadata.groundingMetadata;
+      } else if (currentMetadata.groundingMetadata) {
+        payload.groundingMetadata = currentMetadata.groundingMetadata;
+      }
+    }
+
     sendMessage(res, payload);
     if (_onProgress) {
       _onProgress(payload);
@@ -41,6 +60,13 @@ const createOnProgress = ({ generation = '', onProgress: _onProgress }) => {
   const sendIntermediateMessage = (res, payload, extraTokens = '') => {
     basePayload.text = basePayload.text + extraTokens;
     const message = Object.assign({}, basePayload, payload);
+    if (currentMetadata) {
+      if (currentMetadata.metadata?.groundingMetadata) {
+        message.groundingMetadata = currentMetadata.metadata.groundingMetadata;
+      } else if (currentMetadata.groundingMetadata) {
+        message.groundingMetadata = currentMetadata.groundingMetadata;
+      }
+    }
     sendMessage(res, message);
     if (i === 0) {
       basePayload.initial = false;

client/src/components/Chat/Messages/Content/MessageContent.tsx

@@ -80,11 +80,22 @@ const DisplayMessage = ({ text, isCreatedByUser, message, showCursor }: TDisplay
     () => message.messageId === latestMessage?.messageId,
     [message.messageId, latestMessage?.messageId],
   );
-
   let content: React.ReactElement;
   if (!isCreatedByUser) {
     content = (
-      <Markdown content={text} showCursor={showCursorState} isLatestMessage={isLatestMessage} />
+      <>
+        <Markdown content={text} showCursor={showCursorState} isLatestMessage={isLatestMessage} />
+        {message.groundingMetadata?.searchEntryPoint?.renderedContent && (
+          <div className="mt-4 rounded-lg border border-token-border-light bg-token-surface-secondary p-4">
+            <div 
+              className="prose dark:prose-invert"
+              dangerouslySetInnerHTML={{ 
+                __html: message.groundingMetadata.searchEntryPoint.renderedContent 
+              }} 
+            />
+          </div>
+        )}
+      </>
     );
   } else if (enableUserMsgMarkdown) {
     content = <MarkdownLite content={text} />;