🔧 refactor: Optimize Agent Tool Loading and Fix Bedrock Tool Handling (#4641)

* fix: bedrock tool name regex

* fix: pass args as single input, attempt json first.

* refactor: remove toolMap from agent tool load as is not used

* fix: update formatAgentMessages test to use strictEqual for args comparison, testing new behavior

api/app/clients/prompts/formatAgentMessages.spec.js

@@ -120,7 +120,7 @@ describe('formatAgentMessages', () => {
     ];
     const result = formatAgentMessages(payload);
     expect(result).toHaveLength(2);
-    expect(result[0].tool_calls[0].args).toBe('non-json-string');
+    expect(result[0].tool_calls[0].args).toStrictEqual({ input: 'non-json-string' });
   });
 
   it('should handle complex tool calls with multiple steps', () => {

api/app/clients/prompts/formatMessages.js

@@ -189,10 +189,13 @@ const formatAgentMessages = (payload) => {
         // TODO: investigate; args as dictionary may need to be provider-or-tool-specific
         let args = _args;
         try {
-          args = JSON.parse(args);
+          args = JSON.parse(_args);
         } catch (e) {
-          // failed to parse, leave as is
+          if (typeof _args === 'string') {
+            args = { input: _args };
           }
+        }
+
         tool_call.args = args;
         lastAIMessage.tool_calls.push(tool_call);
 

api/server/controllers/agents/client.js

@@ -461,7 +461,6 @@ class AgentClient extends BaseClient {
         req: this.options.req,
         agent: this.options.agent,
         tools: this.options.tools,
-        toolMap: this.options.toolMap,
         runId: this.responseMessageId,
         modelOptions: this.modelOptions,
         customHandlers: this.options.eventHandlers,

api/server/controllers/agents/run.js

@@ -18,7 +18,6 @@ const { providerEndpointMap } = require('librechat-data-provider');
  * @param {string | undefined} [options.runId] - Optional run ID; otherwise, a new run ID will be generated.
  * @param {Agent} options.agent - The agent for this run.
  * @param {StructuredTool[] | undefined} [options.tools] - The tools to use in the run.
- * @param {Record<string, StructuredTool[]> | undefined} [options.toolMap] - The tool map for the run.
  * @param {Record<GraphEvents, EventHandler> | undefined} [options.customHandlers] - Custom event handlers.
  * @param {ClientOptions} [options.modelOptions] - Optional model to use; if not provided, it will use the default from modelMap.
  * @param {boolean} [options.streaming=true] - Whether to use streaming.
@@ -29,7 +28,6 @@ async function createRun({
   runId,
   tools,
   agent,
-  toolMap,
   modelOptions,
   customHandlers,
   streaming = true,
@@ -47,7 +45,6 @@ async function createRun({
 
   const graphConfig = {
     tools,
-    toolMap,
     llmConfig,
     instructions: agent.instructions,
     additional_instructions: agent.additional_instructions,

api/server/services/ActionService.js

@@ -14,6 +14,7 @@ const { getLogStores } = require('~/cache');
 const { logger } = require('~/config');
 
 const toolNameRegex = /^[a-zA-Z0-9_-]+$/;
+const replaceSeparatorRegex = new RegExp(actionDomainSeparator, 'g');
 
 /**
  * Validates tool name against regex pattern and updates if necessary.
@@ -83,8 +84,6 @@ async function domainParser(req, domain, inverse = false) {
     return key;
   }
 
-  const replaceSeparatorRegex = new RegExp(actionDomainSeparator, 'g');
-
   if (!cachedDomain) {
     return domain.replace(replaceSeparatorRegex, '.');
   }
@@ -156,7 +155,7 @@ async function createActionTool({ action, requestBuilder, zodSchema, name, descr
 
   if (name) {
     return tool(_call, {
-      name,
+      name: name.replace(replaceSeparatorRegex, '_'),
       description: description || '',
       schema: zodSchema,
     });

api/server/services/Endpoints/agents/initialize.js

@@ -54,7 +54,7 @@ const initializeClient = async ({ req, res, endpointOption }) => {
     throw new Error('Agent not found');
   }
 
-  const { tools, toolMap } = await loadAgentTools({
+  const { tools } = await loadAgentTools({
     req,
     tools: agent.tools,
     agent_id: agent.id,
@@ -100,7 +100,6 @@ const initializeClient = async ({ req, res, endpointOption }) => {
     agent,
     tools,
     sender,
-    toolMap,
     contentParts,
     modelOptions,
     eventHandlers,

api/server/services/Endpoints/bedrock/initialize.js

@@ -60,7 +60,6 @@ const initializeClient = async ({ req, res, endpointOption }) => {
     agent,
     sender,
     // tools,
-    // toolMap,
     modelOptions,
     contentParts,
     eventHandlers,

api/server/services/ToolService.js

@@ -373,14 +373,14 @@ async function processRequiredActions(client, requiredActions) {
 }
 
 /**
- * Processes the runtime tool calls and returns a combined toolMap.
+ * Processes the runtime tool calls and returns the tool classes.
  * @param {Object} params - Run params containing user and request information.
  * @param {ServerRequest} params.req - The request object.
  * @param {string} params.agent_id - The agent ID.
  * @param {Agent['tools']} params.tools - The agent's available tools.
  * @param {Agent['tool_resources']} params.tool_resources - The agent's available tool resources.
  * @param {string | undefined} [params.openAIApiKey] - The OpenAI API key.
- * @returns {Promise<{ tools?: StructuredTool[]; toolMap?: Record<string, StructuredTool>}>} The combined toolMap.
+ * @returns {Promise<{ tools?: StructuredTool[] }>} The agent tools.
  */
 async function loadAgentTools({ req, agent_id, tools, tool_resources, openAIApiKey }) {
   if (!tools || tools.length === 0) {
@@ -482,10 +482,8 @@ async function loadAgentTools({ req, agent_id, tools, tool_resources, openAIApiK
     throw new Error('No tools found for the specified tool calls.');
   }
 
-  const toolMap = { ...ToolMap, ...ActionToolMap };
   return {
     tools: agentTools,
-    toolMap,
   };
 }
 

api/typedefs.js

@@ -958,7 +958,6 @@
  * @property {Object} [headers] - Additional headers for requests
  * @property {string} [proxy] - Proxy configuration
  * @property {Object} [tools] - Available tools for the agent
- * @property {Object} [toolMap] - Mapping of tool configurations
  * @property {Object} [eventHandlers] - Custom event handlers
  * @property {Object} [addParams] - Additional parameters to add to requests
  * @property {string[]} [dropParams] - Parameters to remove from requests