Andreas/LibreChat
1
0
Fork 0
mirror of https://github.com/danny-avila/LibreChat.git synced 2025-12-17 17:00:15 +01:00
Code Issues Projects Releases Packages Wiki Activity Actions 4

🧩 feat: Support Alternate API Keys for Plugins (#1760)

Browse source 
* refactor(DALL-E): retrieve env variables at runtime and not from memory

* feat(plugins): add alternate env variable handling to allow setting one api key for multiple plugins

* docs: update docs
This commit is contained in:
Danny Avila 2024-02-09 10:38:50 -05:00 • committed by GitHub
parent 927ce5395b
commit 39caeb2027
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
10 changed files with 328 additions and 113 deletions
Download patch file Download diff file Expand all files Collapse all files

120
api/app/clients/tools/util/handleTools.test.js
View file

@ -4,26 +4,33 @@ const mockUser = {
findByIdAndDelete: jest.fn(),
};
var mockPluginService = {
const mockPluginService = {
updateUserPluginAuth: jest.fn(),
deleteUserPluginAuth: jest.fn(),
getUserPluginAuthValue: jest.fn(),
};
jest.mock('../../../../models/User', () => {
jest.mock('~/models/User', () => {
return function () {
return mockUser;
};
});
jest.mock('../../../../server/services/PluginService', () => mockPluginService);
jest.mock('~/server/services/PluginService', () => mockPluginService);
const User = require('../../../../models/User');
const { validateTools, loadTools } = require('./');
const PluginService = require('../../../../server/services/PluginService');
const { BaseChatModel } = require('langchain/chat_models/openai');
const { Calculator } = require('langchain/tools/calculator');
const { availableTools, OpenAICreateImage, GoogleSearchAPI, StructuredSD } = require('../');
const { BaseChatModel } = require('langchain/chat_models/openai');
const User = require('~/models/User');
const PluginService = require('~/server/services/PluginService');
const { validateTools, loadTools, loadToolWithAuth } = require('./handleTools');
const {
availableTools,
OpenAICreateImage,
GoogleSearchAPI,
StructuredSD,
WolframAlphaAPI,
} = require('../');
describe('Tool Handlers', () => {
let fakeUser;
@ -44,7 +51,10 @@ describe('Tool Handlers', () => {
});
mockPluginService.updateUserPluginAuth.mockImplementation(
(userId, authField, _pluginKey, credential) => {
userAuthValues[`${userId}-${authField}`] = credential;
const fields = authField.split('||');
fields.forEach((field) => {
userAuthValues[`${userId}-${field}`] = credential;
});
},
);
@ -134,6 +144,18 @@ describe('Tool Handlers', () => {
loadTool2 = toolFunctions[sampleTools[1]];
loadTool3 = toolFunctions[sampleTools[2]];
});
let originalEnv;
beforeEach(() => {
originalEnv = process.env;
process.env = { ...originalEnv };
});
afterEach(() => {
process.env = originalEnv;
});
it('returns the expected load functions for requested tools', async () => {
expect(loadTool1).toBeDefined();
expect(loadTool2).toBeDefined();
@ -150,6 +172,86 @@ describe('Tool Handlers', () => {
expect(authTool).toBeInstanceOf(ToolClass);
expect(tool).toBeInstanceOf(ToolClass2);
});
it('should initialize an authenticated tool with primary auth field', async () => {
process.env.DALLE2_API_KEY = 'mocked_api_key';
const initToolFunction = loadToolWithAuth(
'userId',
['DALLE2_API_KEY||DALLE_API_KEY'],
ToolClass,
);
const authTool = await initToolFunction();
expect(authTool).toBeInstanceOf(ToolClass);
expect(mockPluginService.getUserPluginAuthValue).not.toHaveBeenCalled();
});
it('should initialize an authenticated tool with alternate auth field when primary is missing', async () => {
delete process.env.DALLE2_API_KEY; // Ensure the primary key is not set
process.env.DALLE_API_KEY = 'mocked_alternate_api_key';
const initToolFunction = loadToolWithAuth(
'userId',
['DALLE2_API_KEY||DALLE_API_KEY'],
ToolClass,
);
const authTool = await initToolFunction();
expect(authTool).toBeInstanceOf(ToolClass);
expect(mockPluginService.getUserPluginAuthValue).toHaveBeenCalledTimes(1);
expect(mockPluginService.getUserPluginAuthValue).toHaveBeenCalledWith(
'userId',
'DALLE2_API_KEY',
);
});
it('should fallback to getUserPluginAuthValue when env vars are missing', async () => {
mockPluginService.updateUserPluginAuth('userId', 'DALLE_API_KEY', 'dalle', 'mocked_api_key');
const initToolFunction = loadToolWithAuth(
'userId',
['DALLE2_API_KEY||DALLE_API_KEY'],
ToolClass,
);
const authTool = await initToolFunction();
expect(authTool).toBeInstanceOf(ToolClass);
expect(mockPluginService.getUserPluginAuthValue).toHaveBeenCalledTimes(2);
});
it('should initialize an authenticated tool with singular auth field', async () => {
process.env.WOLFRAM_APP_ID = 'mocked_app_id';
const initToolFunction = loadToolWithAuth('userId', ['WOLFRAM_APP_ID'], WolframAlphaAPI);
const authTool = await initToolFunction();
expect(authTool).toBeInstanceOf(WolframAlphaAPI);
expect(mockPluginService.getUserPluginAuthValue).not.toHaveBeenCalled();
});
it('should initialize an authenticated tool when env var is set', async () => {
process.env.WOLFRAM_APP_ID = 'mocked_app_id';
const initToolFunction = loadToolWithAuth('userId', ['WOLFRAM_APP_ID'], WolframAlphaAPI);
const authTool = await initToolFunction();
expect(authTool).toBeInstanceOf(WolframAlphaAPI);
expect(mockPluginService.getUserPluginAuthValue).not.toHaveBeenCalledWith(
'userId',
'WOLFRAM_APP_ID',
);
});
it('should fallback to getUserPluginAuthValue when singular env var is missing', async () => {
delete process.env.WOLFRAM_APP_ID; // Ensure the environment variable is not set
mockPluginService.getUserPluginAuthValue.mockResolvedValue('mocked_user_auth_value');
const initToolFunction = loadToolWithAuth('userId', ['WOLFRAM_APP_ID'], WolframAlphaAPI);
const authTool = await initToolFunction();
expect(authTool).toBeInstanceOf(WolframAlphaAPI);
expect(mockPluginService.getUserPluginAuthValue).toHaveBeenCalledTimes(1);
expect(mockPluginService.getUserPluginAuthValue).toHaveBeenCalledWith(
'userId',
'WOLFRAM_APP_ID',
);
});
it('should throw an error for an unauthenticated tool', async () => {
try {
await loadTool2();