refactor(handleTools.js): change loadTools function signature to include functions parameter

feat(handleTools.test.js): add test for loading StructuredSD tool with functions parameter
This commit is contained in:
Daniel Avila 2023-06-14 13:23:02 -04:00 committed by Danny Avila
parent d339c291fa
commit bffa9ad016
3 changed files with 15 additions and 3 deletions

View file

@ -451,6 +451,7 @@ Only respond with your conversational reply to the following User Message:
user, user,
model, model,
tools: this.options.tools, tools: this.options.tools,
functions: this.functionsAgent,
options: { options: {
openAIApiKey: this.openAIApiKey openAIApiKey: this.openAIApiKey
} }

View file

@ -72,8 +72,7 @@ const loadToolWithAuth = async (user, authFields, ToolConstructor, options = {})
}; };
}; };
const loadTools = async ({ user, model, tools = [], options = {} }) => { const loadTools = async ({ user, model, functions = null, tools = [], options = {} }) => {
const { functions } = options;
const toolConstructors = { const toolConstructors = {
calculator: Calculator, calculator: Calculator,
google: GoogleSearchAPI, google: GoogleSearchAPI,

View file

@ -24,7 +24,7 @@ const { validateTools, loadTools } = require('./');
const PluginService = require('../../../../server/services/PluginService'); const PluginService = require('../../../../server/services/PluginService');
const { BaseChatModel } = require('langchain/chat_models/openai'); const { BaseChatModel } = require('langchain/chat_models/openai');
const { Calculator } = require('langchain/tools/calculator'); const { Calculator } = require('langchain/tools/calculator');
const { availableTools, OpenAICreateImage, GoogleSearchAPI } = require('../'); const { availableTools, OpenAICreateImage, GoogleSearchAPI, StructuredSD } = require('../');
describe('Tool Handlers', () => { describe('Tool Handlers', () => {
let fakeUser; let fakeUser;
@ -174,5 +174,17 @@ describe('Tool Handlers', () => {
}); });
expect(toolFunctions).toEqual({}); expect(toolFunctions).toEqual({});
}); });
it('should return the StructuredTool version when using functions', async () => {
process.env.SD_WEBUI_URL = mockCredential;
toolFunctions = await loadTools({
user: fakeUser._id,
model: BaseChatModel,
tools: ['stable-diffusion'],
functions: true
});
const structuredTool = await toolFunctions['stable-diffusion']();
expect(structuredTool).toBeInstanceOf(StructuredSD);
delete process.env.SD_WEBUI_URL;
});
}); });
}); });