Consolidate various changes as of 2024/12/05

GOOGLE_EXTRA_HEADERS for GenAI 🛡️ feat: Google Reverse Proxy support, `CIVIC_INTEGRITY` harm category (danny-avila#5037) * 🛡️ feat: Google Reverse Proxy support, `CIVIC_INTEGRITY` harm category * 🔧 chore: Update @langchain/google-vertexai to version 0.1.4 in package.json and package-lock.json * fix: revert breaking Vertex AI changes --------- Co-authored-by: KiGamji <[email protected]> Enable google search / retrieval for google endpoints 🔐 feat: Implement Allowed Action Domains (danny-avila#4964) * chore: RequestExecutor typing * feat: allowed action domains * fix: rename TAgentsEndpoint to TAssistantEndpoint in typedefs * chore: update librechat-data-provider version to 0.7.62 🤖 feat: Support Google Agents, fix Various Provider Configurations (danny-avila#5126) * feat: Refactor ModelEndHandler to collect usage metadata only if it exists * feat: google tool end handling, custom anthropic class for better token ux * refactor: differentiate between client <> request options * feat: initial support for google agents * feat: only cache messages with non-empty text * feat: Cache non-empty messages in chatV2 controller * fix: anthropic llm client options llmConfig * refactor: streamline client options handling in LLM configuration * fix: VertexAI Agent Auth & Tool Handling * fix: additional fields for llmConfig, however customHeaders are not supported by langchain, requires PR * feat: set default location for vertexai LLM configuration * fix: outdated OpenAI Client options for getLLMConfig * chore: agent provider options typing * chore: add note about currently unsupported customHeaders in langchain GenAI client * fix: skip transaction creation when rawAmount is NaN 🔧 refactor: Improve Agent Context & Minor Fixes (danny-avila#5349) * refactor: Improve Context for Agents * 🔧 fix: Safeguard against undefined properties in OpenAIClient response handling * refactor: log error before re-throwing for original stack trace * refactor: remove toolResource state from useFileHandling, allow svg files * refactor: prevent verbose logs from axios errors when using actions * refactor: add silent method recordTokenUsage in AgentClient * refactor: streamline token count assignment in BaseClient * refactor: enhance safety settings handling for Gemini 2.0 model * fix: capabilities structure in MCPConnection * refactor: simplify civic integrity threshold handling in GoogleClient and llm * refactor: update token count retrieval method in BaseClient tests * ci: fix test for svg 📜 refactor: Log Error Messages when OAuth Fails (danny-avila#5337) 🔐 feat: Implement Allowed Action Domains (danny-avila#4964) * chore: RequestExecutor typing * feat: allowed action domains * fix: rename TAgentsEndpoint to TAssistantEndpoint in typedefs * chore: update librechat-data-provider version to 0.7.62 🛡️ feat: Google Reverse Proxy support, `CIVIC_INTEGRITY` harm category (danny-avila#5037) * 🛡️ feat: Google Reverse Proxy support, `CIVIC_INTEGRITY` harm category * 🔧 chore: Update @langchain/google-vertexai to version 0.1.4 in package.json and package-lock.json * fix: revert breaking Vertex AI changes --------- Co-authored-by: KiGamji <[email protected]> Enable google search / retrieval for google endpoints
openwengo · Feb 27, 2025 · 746f155 · 746f155
1 parent 34f967e
commit 746f155
Show file tree

Hide file tree

Showing 31 changed files with 544 additions and 40 deletions.
diff --git a/.env.example b/.env.example
@@ -400,6 +400,8 @@ GITHUB_CALLBACK_URL=/oauth/github/callback
 GOOGLE_CLIENT_ID=
 GOOGLE_CLIENT_SECRET=
 GOOGLE_CALLBACK_URL=/oauth/google/callback
+# Optional: restrict login to members of a specific Google Workspace group
+# [email protected]
 
 # Apple
 APPLE_CLIENT_ID=

diff --git a/api/app/clients/AnthropicClient.js b/api/app/clients/AnthropicClient.js
@@ -188,10 +188,15 @@ class AnthropicClient extends BaseClient {
     if (this.options.reverseProxyUrl) {
       options.baseURL = this.options.reverseProxyUrl;
     }
+    if (this.options.defaultHeaders) {
+      options.defaultHeaders = {
+        ...this.options.defaultHeaders,
+      };
+    }
 
     const headers = getClaudeHeaders(requestOptions?.model, this.supportsCacheControl);
     if (headers) {
-      options.defaultHeaders = headers;
+      options.defaultHeaders = { ...(options.defaultHeaders || {}), headers};
     }
 
     return new Anthropic(options);

diff --git a/api/app/clients/BaseClient.js b/api/app/clients/BaseClient.js
@@ -457,6 +457,21 @@ class BaseClient {
     }
 
     let orderedWithInstructions = this.addInstructions(orderedMessages, instructions);
+    if (this.clientName === EModelEndpoint.agents) {
+      const { dbMessages, editedIndices } = truncateToolCallOutputs(
+        orderedWithInstructions,
+        this.maxContextTokens,
+        this.getTokenCountForMessage.bind(this),
+      );
+
+      if (editedIndices.length > 0) {
+        logger.debug('[BaseClient] Truncated tool call outputs:', editedIndices);
+        for (const index of editedIndices) {
+          payload[index].content = dbMessages[index].content;
+        }
+        orderedWithInstructions = dbMessages;
+      }
+    }
 
     let { context, remainingContextTokens, messagesToRefine, summaryIndex } =
       await this.getMessagesWithinTokenLimit({

diff --git a/api/app/clients/GoogleClient.js b/api/app/clients/GoogleClient.js
@@ -2,8 +2,8 @@ const { google } = require('googleapis');
 const { concat } = require('@langchain/core/utils/stream');
 const { ChatVertexAI } = require('@langchain/google-vertexai');
 const { ChatGoogleGenerativeAI } = require('@langchain/google-genai');
-const { GoogleGenerativeAI: GenAI } = require('@google/generative-ai');
-const { HumanMessage, SystemMessage } = require('@langchain/core/messages');
+const { GoogleGenerativeAI: GenAI, DynamicRetrievalMode } = require('@google/generative-ai');
+const { AIMessage, HumanMessage, SystemMessage } = require('@langchain/core/messages');
 const {
   googleGenConfigSchema,
   validateVisionModel,
@@ -30,14 +30,20 @@ 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 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);
@@ -134,7 +140,11 @@ class GoogleClient extends BaseClient {
       this.options = options;
     }
 
-    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));
 
@@ -602,7 +612,29 @@ class GoogleClient extends BaseClient {
       return client;
     } else if (!EXCLUDED_GENAI_MODELS.test(model)) {
       logger.debug('Creating GenAI client');
-      return new GenAI(this.apiKey).getGenerativeModel({ 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,...this.options.customHeaders && { 'customHeaders': this.options.customHeaders }});
     }
 
     logger.debug('Creating Chat Google Generative AI client');
@@ -661,6 +693,7 @@ class GoogleClient extends BaseClient {
           };
         }
 
+
         const delay = modelName.includes('flash') ? 8 : 15;
         /** @type {GenAIUsageMetadata} */
         let usageMetadata;
@@ -676,16 +709,30 @@ class GoogleClient extends BaseClient {
         const result = await client.generateContentStream(requestOptions, {
           signal: abortController.signal,
         });
+        let lastGroundingMetadata = null;
+
         for await (const chunk of result.stream) {
           usageMetadata = !usageMetadata
             ? chunk?.usageMetadata
             : Object.assign(usageMetadata, chunk?.usageMetadata);
-          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);
+          }
         }
 
         if (usageMetadata) {
@@ -695,7 +742,25 @@ class GoogleClient extends BaseClient {
           };
         }
 
-        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 { instances } = _payload;
@@ -715,6 +780,7 @@ class GoogleClient extends BaseClient {
         safetySettings,
       });
 
+
       let delay = this.options.streamRate || 8;
 
       if (!this.options.streamRate) {
@@ -918,11 +984,26 @@ class GoogleClient extends BaseClient {
   }
 
   async sendCompletion(payload, opts = {}) {
-    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 '';
   }
 
+
   getEncoding() {
     return 'cl100k_base';
   }
@@ -955,3 +1036,4 @@ class GoogleClient extends BaseClient {
 }
 
 module.exports = GoogleClient;
+
diff --git a/api/app/clients/TextStream.js b/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', () => {

diff --git a/api/models/Message.js b/api/models/Message.js
@@ -210,6 +210,7 @@ async function updateMessage(req, message, metadata) {
       text: updatedMessage.text,
       isCreatedByUser: updatedMessage.isCreatedByUser,
       tokenCount: updatedMessage.tokenCount,
+      groundingMetadata: updatedMessage.groundingMetadata,
     };
   } catch (err) {
     logger.error('Error updating message:', err);

diff --git a/api/models/schema/defaults.js b/api/models/schema/defaults.js
@@ -130,6 +130,11 @@ const conversationPreset = {
   reasoning_effort: {
     type: String,
   },
+  // for google only
+  enableSearch: {
+    type: Boolean,
+    required: false,
+  },
 };
 
 module.exports = {

diff --git a/api/models/schema/messageSchema.js b/api/models/schema/messageSchema.js
@@ -112,6 +112,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: [

diff --git a/api/server/controllers/AskController.js b/api/server/controllers/AskController.js
@@ -1,3 +1,4 @@
+const throttle = require('lodash/throttle');
 const { getResponseSender, Constants } = require('librechat-data-provider');
 const { createAbortController, handleAbortError } = require('~/server/middleware');
 const { sendMessage, createOnProgress } = require('~/server/utils');
@@ -26,6 +27,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,
@@ -55,7 +57,17 @@ const AskController = async (req, res, next, initializeClient, addTitle) => {
 
   try {
     const { client } = await initializeClient({ req, res, endpointOption });
-    const { onProgress: progressCallback, getPartialText } = createOnProgress();
+    const { onProgress: progressCallback, getPartialText } = createOnProgress({
+      onProgress: throttle(
+        ({ text: partialText, metadata }) => {
+          if (metadata) {
+            currentMetadata = metadata;
+          }
+        },
+        3000,
+        { trailing: false },
+      ),
+    });
 
     getText = client.getStreamText != null ? client.getStreamText.bind(client) : getPartialText;
 
@@ -68,6 +80,7 @@ const AskController = async (req, res, next, initializeClient, addTitle) => {
       text: getText(),
       userMessage,
       promptTokens,
+      ...(currentMetadata ? { metadata: currentMetadata } : {}),
     });
 
     const { abortController, onStart } = createAbortController(req, res, getAbortData, getReqData);
@@ -105,6 +118,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';
@@ -116,19 +138,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' },
         );
       }

diff --git a/api/server/middleware/abortMiddleware.js b/api/server/middleware/abortMiddleware.js
@@ -156,7 +156,7 @@ const handleAbortError = async (res, req, error, data) => {
       message: truncateText(error.message, 350),
     });
   } else {
-    logger.error('[handleAbortError] AI response error; aborting request:', error);
+    logger.error(`[handleAbortError] AI response error; aborting request: ${error.message} , ${error?.stack} ::`, error, error?.stack);
   }
   const { sender, conversationId, messageId, parentMessageId, partialText } = data;