Skip to content

Commit

Permalink
🔧 refactor: Optimize Agent Tool Loading and Fix Bedrock Tool Handling (…
Browse files Browse the repository at this point in the history
…#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
  • Loading branch information
danny-avila authored Nov 5, 2024
1 parent 3428c3c commit 0c2a583
Show file tree
Hide file tree
Showing 9 changed files with 11 additions and 18 deletions.
2 changes: 1 addition & 1 deletion api/app/clients/prompts/formatAgentMessages.spec.js
Original file line number Diff line number Diff line change
Expand Up @@ -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', () => {
Expand Down
7 changes: 5 additions & 2 deletions api/app/clients/prompts/formatMessages.js
Original file line number Diff line number Diff line change
Expand Up @@ -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);

Expand Down
1 change: 0 additions & 1 deletion api/server/controllers/agents/client.js
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
3 changes: 0 additions & 3 deletions api/server/controllers/agents/run.js
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -29,7 +28,6 @@ async function createRun({
runId,
tools,
agent,
toolMap,
modelOptions,
customHandlers,
streaming = true,
Expand All @@ -47,7 +45,6 @@ async function createRun({

const graphConfig = {
tools,
toolMap,
llmConfig,
instructions: agent.instructions,
additional_instructions: agent.additional_instructions,
Expand Down
5 changes: 2 additions & 3 deletions api/server/services/ActionService.js
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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, '.');
}
Expand Down Expand Up @@ -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,
});
Expand Down
3 changes: 1 addition & 2 deletions api/server/services/Endpoints/agents/initialize.js
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -100,7 +100,6 @@ const initializeClient = async ({ req, res, endpointOption }) => {
agent,
tools,
sender,
toolMap,
contentParts,
modelOptions,
eventHandlers,
Expand Down
1 change: 0 additions & 1 deletion api/server/services/Endpoints/bedrock/initialize.js
Original file line number Diff line number Diff line change
Expand Up @@ -60,7 +60,6 @@ const initializeClient = async ({ req, res, endpointOption }) => {
agent,
sender,
// tools,
// toolMap,
modelOptions,
contentParts,
eventHandlers,
Expand Down
6 changes: 2 additions & 4 deletions api/server/services/ToolService.js
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down Expand Up @@ -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,
};
}

Expand Down
1 change: 0 additions & 1 deletion api/typedefs.js
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down

0 comments on commit 0c2a583

Please sign in to comment.