Tool Calling: add shouldRegister function to tool defintion

bc0f5bf4ceea7960ea574e291b576ea5115f42e6

Cohee <18619528+Cohee1207@users.noreply.github.com>

1 files changed, +30 -10Ignore whitespace
public/scripts/tool-calling.js+30 -10
@@ -32,6 +32,7 @@ import { slashCommandReturnHelper } from './slash-commands/SlashCommandReturnHel
3232 * @property {object} parameters - The parameters for the tool.
3333 * @property {function} action - The action to perform when the tool is invoked.
3434 * @property {function} formatMessage - A function to format the tool call message.
35+ * @property {function} shouldRegister - A function to determine if the tool should be registered.
3536 */
3637
3738/**
@@ -138,6 +139,12 @@ class ToolDefinition {
138139 #formatMessage;
139140
140141 /**
142+ * A function that will be called to determine if the tool should be registered.
143+ * @type {function}
144+ */
145+ #shouldRegister;
146+
147+ /**
141148 * Creates a new ToolDefinition.
142149 * @param {string} name A unique name for the tool.
143150 * @param {string} displayName A user-friendly display name for the tool.
@@ -145,14 +152,16 @@ class ToolDefinition {
145152 * @param {object} parameters A JSON schema for the parameters that the tool accepts.
146153 * @param {function} action A function that will be called when the tool is executed.
147154 * @param {function} formatMessage A function that will be called to format the tool call toast.
155+ * @param {function} shouldRegister A function that will be called to determine if the tool should be registered.
148156 */
149157 constructor(name, displayName, description, parameters, action, formatMessage, shouldRegister) {
150158 this.#name = name;
151159 this.#displayName = displayName;
152160 this.#description = description;
153161 this.#parameters = parameters;
154162 this.#action = action;
155163 this.#formatMessage = formatMessage;
164+ this.#shouldRegister = shouldRegister;
156165 }
157166
158167 /**
@@ -193,6 +202,12 @@ class ToolDefinition {
193202 : `Invoking tool: ${this.#displayName || this.#name}`;
194203 }
195204
205+ async shouldRegister() {
206+ return typeof this.#shouldRegister === 'function'
207+ ? await this.#shouldRegister()
208+ : true;
209+ }
210+
196211 get displayName() {
197212 return this.#displayName;
198213 }
@@ -228,17 +243,17 @@ export class ToolManager {
228243 * Registers a new tool with the tool registry.
229244 * @param {ToolRegistration} tool The tool to register.
230245 */
231246 static registerFunctionTool({ name, displayName, description, parameters, action, formatMessage, shouldRegister }) {
232247 // Convert WIP arguments
233248 if (typeof arguments[0] !== 'object') {
234249 [name, description, parameters, action] = arguments;
235250 }
236251
237252 if (this.#tools.has(name)) {
238253 console.warn(`[ToolManager] A tool with the name "${name}" has already been registered. The definition will be overwritten.`);
239254 }
240255
241256 const definition = new ToolDefinition(name, displayName, description, parameters, action, formatMessage, shouldRegister);
242257 this.#tools.set(name, definition);
243258 console.log('[ToolManager] Registered function tool:', definition);
244259 }
@@ -273,7 +288,7 @@ export class ToolManager {
273288 const result = await tool.invoke(invokeParameters);
274289 return typeof result === 'string' ? result : JSON.stringify(result);
275290 } catch (error) {
276291 console.error(`[ToolManager] An error occurred while invoking the tool "${name}":`, error);
277292
278293 if (error instanceof Error) {
279294 error.cause = name;
@@ -300,7 +315,7 @@ export class ToolManager {
300315 const formatParameters = typeof parameters === 'string' ? JSON.parse(parameters) : parameters;
301316 return tool.formatMessage(formatParameters);
302317 } catch (error) {
303318 console.error(`[ToolManager] An error occurred while formatting the tool call message for "${name}":`, error);
304319 return `Invoking tool: ${name}`;
305320 }
306321 }
@@ -327,11 +342,16 @@ export class ToolManager {
327342 const tools = [];
328343
329344 for (const tool of ToolManager.tools) {
345+ const register = await tool.shouldRegister();
346+ if (!register) {
347+ console.log('[ToolManager] Skipping tool registration:', tool);
348+ continue;
349+ }
330350 tools.push(tool.toFunctionOpenAI());
331351 }
332352
333353 if (tools.length) {
334354 console.log('[ToolManager] Registered function tools:', tools);
335355
336356 data['tools'] = tools;
337357 data['tool_choice'] = 'auto';
@@ -422,7 +442,7 @@ export class ToolManager {
422442 delete targetToolCall[this.#INPUT_DELTA_KEY];
423443 ToolManager.#applyToolCallDelta(targetToolCall, jsonDelta);
424444 } catch (error) {
425445 console.warn('[ToolManager] Failed to apply input JSON delta:', error);
426446 }
427447 }
428448 }
@@ -564,7 +584,7 @@ export class ToolManager {
564584 continue;
565585 }
566586
567587 console.log('[ToolManager] Function tool call:', toolCall);
568588 const id = toolCall.id;
569589 const parameters = toolCall.function.arguments;
570590 const name = toolCall.function.name;
@@ -574,7 +594,7 @@ export class ToolManager {
574594 const toast = message && toastr.info(message, 'Tool Calling', { timeOut: 0 });
575595 const toolResult = await ToolManager.invokeFunctionTool(name, parameters);
576596 toastr.clear(toast);
577597 console.log('[ToolManager] Function tool result:', result);
578598
579599 // Save a successful invocation
580600 if (toolResult instanceof Error) {