2023-12-10 14:54:13 -05:00
|
|
|
const { OpenAIClient } = require('~/app');
|
2023-11-04 15:03:31 -04:00
|
|
|
const initializeClient = require('./initializeClient');
|
2023-12-10 14:54:13 -05:00
|
|
|
const { getUserKey } = require('~/server/services/UserService');
|
2023-11-04 15:03:31 -04:00
|
|
|
|
|
|
|
// Mock getUserKey since it's the only function we want to mock
|
2023-12-10 14:54:13 -05:00
|
|
|
jest.mock('~/server/services/UserService', () => ({
|
2023-11-04 15:03:31 -04:00
|
|
|
getUserKey: jest.fn(),
|
2023-12-10 14:54:13 -05:00
|
|
|
checkUserKeyExpiry: jest.requireActual('~/server/services/UserService').checkUserKeyExpiry,
|
2023-11-04 15:03:31 -04:00
|
|
|
}));
|
|
|
|
|
|
|
|
describe('initializeClient', () => {
|
|
|
|
// Set up environment variables
|
|
|
|
const originalEnvironment = process.env;
|
|
|
|
|
|
|
|
beforeEach(() => {
|
|
|
|
jest.resetModules(); // Clears the cache
|
|
|
|
process.env = { ...originalEnvironment }; // Make a copy
|
|
|
|
});
|
|
|
|
|
|
|
|
afterAll(() => {
|
|
|
|
process.env = originalEnvironment; // Restore original env vars
|
|
|
|
});
|
|
|
|
|
|
|
|
test('should initialize client with OpenAI API key and default options', async () => {
|
|
|
|
process.env.OPENAI_API_KEY = 'test-openai-api-key';
|
|
|
|
process.env.DEBUG_OPENAI = 'false';
|
|
|
|
process.env.OPENAI_SUMMARIZE = 'false';
|
|
|
|
|
|
|
|
const req = {
|
|
|
|
body: { key: null, endpoint: 'openAI' },
|
|
|
|
user: { id: '123' },
|
|
|
|
};
|
|
|
|
const res = {};
|
|
|
|
const endpointOption = {};
|
|
|
|
|
|
|
|
const client = await initializeClient({ req, res, endpointOption });
|
|
|
|
|
|
|
|
expect(client.openAIApiKey).toBe('test-openai-api-key');
|
|
|
|
expect(client.client).toBeInstanceOf(OpenAIClient);
|
|
|
|
});
|
|
|
|
|
|
|
|
test('should initialize client with Azure credentials when endpoint is azureOpenAI', async () => {
|
|
|
|
process.env.AZURE_API_KEY = 'test-azure-api-key';
|
2023-11-10 09:58:17 -05:00
|
|
|
(process.env.AZURE_OPENAI_API_INSTANCE_NAME = 'some-value'),
|
|
|
|
(process.env.AZURE_OPENAI_API_DEPLOYMENT_NAME = 'some-value'),
|
|
|
|
(process.env.AZURE_OPENAI_API_VERSION = 'some-value'),
|
|
|
|
(process.env.AZURE_OPENAI_API_COMPLETIONS_DEPLOYMENT_NAME = 'some-value'),
|
|
|
|
(process.env.AZURE_OPENAI_API_EMBEDDINGS_DEPLOYMENT_NAME = 'some-value'),
|
|
|
|
(process.env.OPENAI_API_KEY = 'test-openai-api-key');
|
2023-11-04 15:03:31 -04:00
|
|
|
process.env.DEBUG_OPENAI = 'false';
|
|
|
|
process.env.OPENAI_SUMMARIZE = 'false';
|
|
|
|
|
|
|
|
const req = {
|
|
|
|
body: { key: null, endpoint: 'azureOpenAI' },
|
|
|
|
user: { id: '123' },
|
|
|
|
};
|
|
|
|
const res = {};
|
|
|
|
const endpointOption = { modelOptions: { model: 'test-model' } };
|
|
|
|
|
|
|
|
const client = await initializeClient({ req, res, endpointOption });
|
|
|
|
|
|
|
|
expect(client.openAIApiKey).toBe('test-azure-api-key');
|
|
|
|
expect(client.client).toBeInstanceOf(OpenAIClient);
|
|
|
|
});
|
|
|
|
|
|
|
|
test('should use the debug option when DEBUG_OPENAI is enabled', async () => {
|
|
|
|
process.env.OPENAI_API_KEY = 'test-openai-api-key';
|
|
|
|
process.env.DEBUG_OPENAI = 'true';
|
|
|
|
|
|
|
|
const req = {
|
|
|
|
body: { key: null, endpoint: 'openAI' },
|
|
|
|
user: { id: '123' },
|
|
|
|
};
|
|
|
|
const res = {};
|
|
|
|
const endpointOption = {};
|
|
|
|
|
|
|
|
const client = await initializeClient({ req, res, endpointOption });
|
|
|
|
|
|
|
|
expect(client.client.options.debug).toBe(true);
|
|
|
|
});
|
|
|
|
|
|
|
|
test('should set contextStrategy to summarize when OPENAI_SUMMARIZE is enabled', async () => {
|
|
|
|
process.env.OPENAI_API_KEY = 'test-openai-api-key';
|
|
|
|
process.env.OPENAI_SUMMARIZE = 'true';
|
|
|
|
|
|
|
|
const req = {
|
|
|
|
body: { key: null, endpoint: 'openAI' },
|
|
|
|
user: { id: '123' },
|
|
|
|
};
|
|
|
|
const res = {};
|
|
|
|
const endpointOption = {};
|
|
|
|
|
|
|
|
const client = await initializeClient({ req, res, endpointOption });
|
|
|
|
|
|
|
|
expect(client.client.options.contextStrategy).toBe('summarize');
|
|
|
|
});
|
|
|
|
|
|
|
|
test('should set reverseProxyUrl and proxy when they are provided in the environment', async () => {
|
|
|
|
process.env.OPENAI_API_KEY = 'test-openai-api-key';
|
|
|
|
process.env.OPENAI_REVERSE_PROXY = 'http://reverse.proxy';
|
|
|
|
process.env.PROXY = 'http://proxy';
|
|
|
|
|
|
|
|
const req = {
|
|
|
|
body: { key: null, endpoint: 'openAI' },
|
|
|
|
user: { id: '123' },
|
|
|
|
};
|
|
|
|
const res = {};
|
|
|
|
const endpointOption = {};
|
|
|
|
|
|
|
|
const client = await initializeClient({ req, res, endpointOption });
|
|
|
|
|
|
|
|
expect(client.client.options.reverseProxyUrl).toBe('http://reverse.proxy');
|
|
|
|
expect(client.client.options.proxy).toBe('http://proxy');
|
|
|
|
});
|
|
|
|
|
|
|
|
test('should throw an error if the user-provided key has expired', async () => {
|
|
|
|
process.env.OPENAI_API_KEY = 'user_provided';
|
|
|
|
process.env.AZURE_API_KEY = 'user_provided';
|
|
|
|
process.env.DEBUG_OPENAI = 'false';
|
|
|
|
process.env.OPENAI_SUMMARIZE = 'false';
|
|
|
|
|
|
|
|
const expiresAt = new Date(Date.now() - 10000).toISOString(); // Expired
|
|
|
|
const req = {
|
|
|
|
body: { key: expiresAt, endpoint: 'openAI' },
|
|
|
|
user: { id: '123' },
|
|
|
|
};
|
|
|
|
const res = {};
|
|
|
|
const endpointOption = {};
|
|
|
|
|
|
|
|
await expect(initializeClient({ req, res, endpointOption })).rejects.toThrow(
|
|
|
|
'Your OpenAI API key has expired. Please provide your API key again.',
|
|
|
|
);
|
|
|
|
});
|
|
|
|
|
|
|
|
test('should throw an error if no API keys are provided in the environment', async () => {
|
|
|
|
// Clear the environment variables for API keys
|
|
|
|
delete process.env.OPENAI_API_KEY;
|
|
|
|
delete process.env.AZURE_API_KEY;
|
|
|
|
|
|
|
|
const req = {
|
|
|
|
body: { key: null, endpoint: 'openAI' },
|
|
|
|
user: { id: '123' },
|
|
|
|
};
|
|
|
|
const res = {};
|
|
|
|
const endpointOption = {};
|
|
|
|
|
|
|
|
await expect(initializeClient({ req, res, endpointOption })).rejects.toThrow(
|
|
|
|
'API key not provided.',
|
|
|
|
);
|
|
|
|
});
|
|
|
|
|
|
|
|
it('should handle user-provided keys and check expiry', async () => {
|
|
|
|
// Set up the req.body to simulate user-provided key scenario
|
|
|
|
const req = {
|
|
|
|
body: {
|
|
|
|
key: new Date(Date.now() + 10000).toISOString(),
|
|
|
|
endpoint: 'openAI',
|
|
|
|
},
|
|
|
|
user: {
|
|
|
|
id: '123',
|
|
|
|
},
|
|
|
|
};
|
|
|
|
|
|
|
|
const res = {};
|
|
|
|
const endpointOption = {};
|
|
|
|
|
|
|
|
// Ensure the environment variable is set to 'user_provided' to match the isUserProvided condition
|
|
|
|
process.env.OPENAI_API_KEY = 'user_provided';
|
|
|
|
|
|
|
|
// Mock getUserKey to return the expected key
|
|
|
|
getUserKey.mockResolvedValue('test-user-provided-openai-api-key');
|
|
|
|
|
|
|
|
// Call the initializeClient function
|
|
|
|
const result = await initializeClient({ req, res, endpointOption });
|
|
|
|
|
|
|
|
// Assertions
|
|
|
|
expect(result.openAIApiKey).toBe('test-user-provided-openai-api-key');
|
|
|
|
});
|
|
|
|
|
|
|
|
test('should throw an error if the user-provided key is invalid', async () => {
|
|
|
|
const invalidKey = new Date(Date.now() - 100000).toISOString();
|
|
|
|
const req = {
|
|
|
|
body: { key: invalidKey, endpoint: 'openAI' },
|
|
|
|
user: { id: '123' },
|
|
|
|
};
|
|
|
|
const res = {};
|
|
|
|
const endpointOption = {};
|
|
|
|
|
|
|
|
// Ensure the environment variable is set to 'user_provided' to match the isUserProvided condition
|
|
|
|
process.env.OPENAI_API_KEY = 'user_provided';
|
|
|
|
|
|
|
|
// Mock getUserKey to return an invalid key
|
|
|
|
getUserKey.mockResolvedValue(invalidKey);
|
|
|
|
|
|
|
|
await expect(initializeClient({ req, res, endpointOption })).rejects.toThrow(
|
|
|
|
/Your OpenAI API key has expired/,
|
|
|
|
);
|
|
|
|
});
|
|
|
|
});
|