🤖 feat: Support Google Agents, fix Various Provider Configurations (d…

…anny-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
jmaddington · Dec 28, 2024 · 24cad6b · 24cad6b
1 parent a423eb8
commit 24cad6b
Show file tree

Hide file tree

Showing 18 changed files with 428 additions and 362 deletions.
diff --git a/api/app/clients/BaseClient.js b/api/app/clients/BaseClient.js
@@ -649,15 +649,17 @@ class BaseClient {
 
     this.responsePromise = this.saveMessageToDatabase(responseMessage, saveOptions, user);
     this.savedMessageIds.add(responseMessage.messageId);
-    const messageCache = getLogStores(CacheKeys.MESSAGES);
-    messageCache.set(
-      responseMessageId,
-      {
-        text: responseMessage.text,
-        complete: true,
-      },
-      Time.FIVE_MINUTES,
-    );
+    if (responseMessage.text) {
+      const messageCache = getLogStores(CacheKeys.MESSAGES);
+      messageCache.set(
+        responseMessageId,
+        {
+          text: responseMessage.text,
+          complete: true,
+        },
+        Time.FIVE_MINUTES,
+      );
+    }
     delete responseMessage.tokenCount;
     return responseMessage;
   }

diff --git a/api/app/clients/PluginsClient.js b/api/app/clients/PluginsClient.js
@@ -256,15 +256,17 @@ class PluginsClient extends OpenAIClient {
     }
 
     this.responsePromise = this.saveMessageToDatabase(responseMessage, saveOptions, user);
-    const messageCache = getLogStores(CacheKeys.MESSAGES);
-    messageCache.set(
-      responseMessage.messageId,
-      {
-        text: responseMessage.text,
-        complete: true,
-      },
-      Time.FIVE_MINUTES,
-    );
+    if (responseMessage.text) {
+      const messageCache = getLogStores(CacheKeys.MESSAGES);
+      messageCache.set(
+        responseMessage.messageId,
+        {
+          text: responseMessage.text,
+          complete: true,
+        },
+        Time.FIVE_MINUTES,
+      );
+    }
     delete responseMessage.tokenCount;
     return { ...responseMessage, ...result };
   }

diff --git a/api/models/Transaction.js b/api/models/Transaction.js
@@ -27,6 +27,9 @@ transactionSchema.methods.calculateTokenValue = function () {
  */
 transactionSchema.statics.create = async function (txData) {
   const Transaction = this;
+  if (txData.rawAmount != null && isNaN(txData.rawAmount)) {
+    return;
+  }
 
   const transaction = new Transaction(txData);
   transaction.endpointTokenConfig = txData.endpointTokenConfig;

diff --git a/api/models/Transaction.spec.js b/api/models/Transaction.spec.js
@@ -1,5 +1,6 @@
 const mongoose = require('mongoose');
 const { MongoMemoryServer } = require('mongodb-memory-server');
+const { Transaction } = require('./Transaction');
 const Balance = require('./Balance');
 const { spendTokens, spendStructuredTokens } = require('./spendTokens');
 const { getMultiplier, getCacheMultiplier } = require('./tx');
@@ -346,3 +347,28 @@ describe('Structured Token Spending Tests', () => {
     expect(result.completion.completion).toBeCloseTo(-50 * 15 * 1.15, 0); // Assuming multiplier is 15 and cancelRate is 1.15
   });
 });
+
+describe('NaN Handling Tests', () => {
+  test('should skip transaction creation when rawAmount is NaN', async () => {
+    const userId = new mongoose.Types.ObjectId();
+    const initialBalance = 10000000;
+    await Balance.create({ user: userId, tokenCredits: initialBalance });
+
+    const model = 'gpt-3.5-turbo';
+    const txData = {
+      user: userId,
+      conversationId: 'test-conversation-id',
+      model,
+      context: 'test',
+      endpointTokenConfig: null,
+      rawAmount: NaN,
+      tokenType: 'prompt',
+    };
+
+    const result = await Transaction.create(txData);
+    expect(result).toBeUndefined();
+
+    const balance = await Balance.findOne({ user: userId });
+    expect(balance.tokenCredits).toBe(initialBalance);
+  });
+});
diff --git a/api/package.json b/api/package.json
@@ -44,7 +44,7 @@
     "@langchain/google-genai": "^0.1.4",
     "@langchain/google-vertexai": "^0.1.4",
     "@langchain/textsplitters": "^0.1.0",
-    "@librechat/agents": "^1.8.8",
+    "@librechat/agents": "^1.9.7",
     "axios": "^1.7.7",
     "bcryptjs": "^2.4.3",
     "cheerio": "^1.0.0-rc.12",

diff --git a/api/server/controllers/agents/callbacks.js b/api/server/controllers/agents/callbacks.js
@@ -1,8 +1,10 @@
 const { Tools, StepTypes, imageGenTools, FileContext } = require('librechat-data-provider');
 const {
   EnvVar,
+  Providers,
   GraphEvents,
   ToolEndHandler,
+  handleToolCalls,
   ChatModelStreamHandler,
 } = require('@librechat/agents');
 const { processCodeOutput } = require('~/server/services/Files/Code/process');
@@ -57,13 +59,22 @@ class ModelEndHandler {
       return;
     }
 
-    const usage = data?.output?.usage_metadata;
-    if (metadata?.model) {
-      usage.model = metadata.model;
-    }
+    try {
+      if (metadata.provider === Providers.GOOGLE) {
+        handleToolCalls(data?.output?.tool_calls, metadata, graph);
+      }
+
+      const usage = data?.output?.usage_metadata;
+      if (!usage) {
+        return;
+      }
+      if (metadata?.model) {
+        usage.model = metadata.model;
+      }
 
-    if (usage) {
       this.collectedUsage.push(usage);
+    } catch (error) {
+      logger.error('Error handling model end event:', error);
     }
   }
 }

diff --git a/api/server/controllers/assistants/chatV2.js b/api/server/controllers/assistants/chatV2.js
@@ -398,15 +398,17 @@ const chatV2 = async (req, res) => {
       response = streamRunManager;
       response.text = streamRunManager.intermediateText;
 
-      const messageCache = getLogStores(CacheKeys.MESSAGES);
-      messageCache.set(
-        responseMessageId,
-        {
-          complete: true,
-          text: response.text,
-        },
-        Time.FIVE_MINUTES,
-      );
+      if (response.text) {
+        const messageCache = getLogStores(CacheKeys.MESSAGES);
+        messageCache.set(
+          responseMessageId,
+          {
+            complete: true,
+            text: response.text,
+          },
+          Time.FIVE_MINUTES,
+        );
+      }
     };
 
     await processRun();

diff --git a/api/server/services/Endpoints/agents/initialize.js b/api/server/services/Endpoints/agents/initialize.js
@@ -12,6 +12,7 @@ const initAnthropic = require('~/server/services/Endpoints/anthropic/initialize'
 const getBedrockOptions = require('~/server/services/Endpoints/bedrock/options');
 const initOpenAI = require('~/server/services/Endpoints/openAI/initialize');
 const initCustom = require('~/server/services/Endpoints/custom/initialize');
+const initGoogle = require('~/server/services/Endpoints/google/initialize');
 const { getCustomEndpointConfig } = require('~/server/services/Config');
 const { loadAgentTools } = require('~/server/services/ToolService');
 const AgentClient = require('~/server/controllers/agents/client');
@@ -24,6 +25,7 @@ const providerConfigMap = {
   [EModelEndpoint.azureOpenAI]: initOpenAI,
   [EModelEndpoint.anthropic]: initAnthropic,
   [EModelEndpoint.bedrock]: getBedrockOptions,
+  [EModelEndpoint.google]: initGoogle,
   [Providers.OLLAMA]: initCustom,
 };
 
@@ -116,6 +118,10 @@ const initializeAgentOptions = async ({
     endpointOption: _endpointOption,
   });
 
+  if (options.provider != null) {
+    agent.provider = options.provider;
+  }
+
   agent.model_parameters = Object.assign(model_parameters, options.llmConfig);
   if (options.configOptions) {
     agent.model_parameters.configuration = options.configOptions;

diff --git a/api/server/services/Endpoints/anthropic/initialize.js b/api/server/services/Endpoints/anthropic/initialize.js
@@ -20,7 +20,7 @@ const initializeClient = async ({ req, res, endpointOption, overrideModel, optio
     checkUserKeyExpiry(expiresAt, EModelEndpoint.anthropic);
   }
 
-  const clientOptions = {};
+  let clientOptions = {};
 
   /** @type {undefined | TBaseEndpoint} */
   const anthropicConfig = req.app.locals[EModelEndpoint.anthropic];
@@ -36,7 +36,7 @@ const initializeClient = async ({ req, res, endpointOption, overrideModel, optio
   }
 
   if (optionsOnly) {
-    const requestOptions = Object.assign(
+    clientOptions = Object.assign(
       {
         reverseProxyUrl: ANTHROPIC_REVERSE_PROXY ?? null,
         proxy: PROXY ?? null,
@@ -45,9 +45,9 @@ const initializeClient = async ({ req, res, endpointOption, overrideModel, optio
       clientOptions,
     );
     if (overrideModel) {
-      requestOptions.modelOptions.model = overrideModel;
+      clientOptions.modelOptions.model = overrideModel;
     }
-    return getLLMConfig(anthropicApiKey, requestOptions);
+    return getLLMConfig(anthropicApiKey, clientOptions);
   }
 
   const client = new AnthropicClient(anthropicApiKey, {

diff --git a/api/server/services/Endpoints/anthropic/llm.js b/api/server/services/Endpoints/anthropic/llm.js
@@ -28,28 +28,32 @@ function getLLMConfig(apiKey, options = {}) {
 
   const mergedOptions = Object.assign(defaultOptions, options.modelOptions);
 
+  /** @type {AnthropicClientOptions} */
   const requestOptions = {
     apiKey,
     model: mergedOptions.model,
     stream: mergedOptions.stream,
     temperature: mergedOptions.temperature,
-    top_p: mergedOptions.topP,
-    top_k: mergedOptions.topK,
-    stop_sequences: mergedOptions.stop,
-    max_tokens:
+    topP: mergedOptions.topP,
+    topK: mergedOptions.topK,
+    stopSequences: mergedOptions.stop,
+    maxTokens:
       mergedOptions.maxOutputTokens || anthropicSettings.maxOutputTokens.reset(mergedOptions.model),
+    clientOptions: {},
   };
 
-  const configOptions = {};
   if (options.proxy) {
-    configOptions.httpAgent = new HttpsProxyAgent(options.proxy);
+    requestOptions.clientOptions.httpAgent = new HttpsProxyAgent(options.proxy);
   }
 
   if (options.reverseProxyUrl) {
-    configOptions.baseURL = options.reverseProxyUrl;
+    requestOptions.clientOptions.baseURL = options.reverseProxyUrl;
   }
 
-  return { llmConfig: removeNullishValues(requestOptions), configOptions };
+  return {
+    /** @type {AnthropicClientOptions} */
+    llmConfig: removeNullishValues(requestOptions),
+  };
 }
 
 module.exports = { getLLMConfig };
diff --git a/api/server/services/Endpoints/bedrock/options.js b/api/server/services/Endpoints/bedrock/options.js
@@ -60,42 +60,41 @@ const getOptions = async ({ req, endpointOption }) => {
     streamRate = allConfig.streamRate;
   }
 
-  /** @type {import('@librechat/agents').BedrockConverseClientOptions} */
-  const requestOptions = Object.assign(
-    {
-      model: endpointOption.model,
-      region: BEDROCK_AWS_DEFAULT_REGION,
-      streaming: true,
-      streamUsage: true,
-      callbacks: [
-        {
-          handleLLMNewToken: async () => {
-            if (!streamRate) {
-              return;
-            }
-            await sleep(streamRate);
-          },
+  /** @type {BedrockClientOptions} */
+  const requestOptions = {
+    model: endpointOption.model,
+    region: BEDROCK_AWS_DEFAULT_REGION,
+    streaming: true,
+    streamUsage: true,
+    callbacks: [
+      {
+        handleLLMNewToken: async () => {
+          if (!streamRate) {
+            return;
+          }
+          await sleep(streamRate);
         },
-      ],
-    },
-    endpointOption.model_parameters,
-  );
+      },
+    ],
+  };
 
   if (credentials) {
     requestOptions.credentials = credentials;
   }
 
+  if (BEDROCK_REVERSE_PROXY) {
+    requestOptions.endpointHost = BEDROCK_REVERSE_PROXY;
+  }
+
   const configOptions = {};
   if (PROXY) {
+    /** NOTE: NOT SUPPORTED BY BEDROCK */
     configOptions.httpAgent = new HttpsProxyAgent(PROXY);
   }
 
-  if (BEDROCK_REVERSE_PROXY) {
-    configOptions.endpointHost = BEDROCK_REVERSE_PROXY;
-  }
-
   return {
-    llmConfig: removeNullishValues(requestOptions),
+    /** @type {BedrockClientOptions} */
+    llmConfig: removeNullishValues(Object.assign(requestOptions, endpointOption.model_parameters)),
     configOptions,
   };
 };

diff --git a/api/server/services/Endpoints/custom/initialize.js b/api/server/services/Endpoints/custom/initialize.js
@@ -123,7 +123,7 @@ const initializeClient = async ({ req, res, endpointOption, optionsOnly, overrid
     customOptions.streamRate = allConfig.streamRate;
   }
 
-  const clientOptions = {
+  let clientOptions = {
     reverseProxyUrl: baseURL ?? null,
     proxy: PROXY ?? null,
     req,
@@ -135,13 +135,13 @@ const initializeClient = async ({ req, res, endpointOption, optionsOnly, overrid
   if (optionsOnly) {
     const modelOptions = endpointOption.model_parameters;
     if (endpoint !== Providers.OLLAMA) {
-      const requestOptions = Object.assign(
+      clientOptions = Object.assign(
         {
           modelOptions,
         },
         clientOptions,
       );
-      const options = getLLMConfig(apiKey, requestOptions);
+      const options = getLLMConfig(apiKey, clientOptions);
       if (!customOptions.streamRate) {
         return options;
       }