|
const BaseClient = require('../BaseClient'); |
|
const { getModelMaxTokens } = require('../../../utils'); |
|
|
|
class FakeClient extends BaseClient { |
|
constructor(apiKey, options = {}) { |
|
super(apiKey, options); |
|
this.sender = 'AI Assistant'; |
|
this.setOptions(options); |
|
} |
|
setOptions(options) { |
|
if (this.options && !this.options.replaceOptions) { |
|
this.options.modelOptions = { |
|
...this.options.modelOptions, |
|
...options.modelOptions, |
|
}; |
|
delete options.modelOptions; |
|
this.options = { |
|
...this.options, |
|
...options, |
|
}; |
|
} else { |
|
this.options = options; |
|
} |
|
|
|
if (this.options.openaiApiKey) { |
|
this.apiKey = this.options.openaiApiKey; |
|
} |
|
|
|
const modelOptions = this.options.modelOptions || {}; |
|
if (!this.modelOptions) { |
|
this.modelOptions = { |
|
...modelOptions, |
|
model: modelOptions.model || 'gpt-3.5-turbo', |
|
temperature: |
|
typeof modelOptions.temperature === 'undefined' ? 0.8 : modelOptions.temperature, |
|
top_p: typeof modelOptions.top_p === 'undefined' ? 1 : modelOptions.top_p, |
|
presence_penalty: |
|
typeof modelOptions.presence_penalty === 'undefined' ? 1 : modelOptions.presence_penalty, |
|
stop: modelOptions.stop, |
|
}; |
|
} |
|
|
|
this.maxContextTokens = getModelMaxTokens(this.modelOptions.model) ?? 4097; |
|
} |
|
getCompletion() {} |
|
buildMessages() {} |
|
getTokenCount(str) { |
|
return str.length; |
|
} |
|
getTokenCountForMessage(message) { |
|
return message?.content?.length || message.length; |
|
} |
|
} |
|
|
|
const initializeFakeClient = (apiKey, options, fakeMessages) => { |
|
let TestClient = new FakeClient(apiKey); |
|
TestClient.options = options; |
|
TestClient.abortController = { abort: jest.fn() }; |
|
TestClient.saveMessageToDatabase = jest.fn(); |
|
TestClient.loadHistory = jest |
|
.fn() |
|
.mockImplementation((conversationId, parentMessageId = null) => { |
|
if (!conversationId) { |
|
TestClient.currentMessages = []; |
|
return Promise.resolve([]); |
|
} |
|
|
|
const orderedMessages = TestClient.constructor.getMessagesForConversation({ |
|
messages: fakeMessages, |
|
parentMessageId, |
|
}); |
|
|
|
TestClient.currentMessages = orderedMessages; |
|
return Promise.resolve(orderedMessages); |
|
}); |
|
|
|
TestClient.getSaveOptions = jest.fn().mockImplementation(() => { |
|
return {}; |
|
}); |
|
|
|
TestClient.getBuildMessagesOptions = jest.fn().mockImplementation(() => { |
|
return {}; |
|
}); |
|
|
|
TestClient.sendCompletion = jest.fn(async () => { |
|
return 'Mock response text'; |
|
}); |
|
|
|
TestClient.buildMessages = jest.fn(async (messages, parentMessageId) => { |
|
const orderedMessages = TestClient.constructor.getMessagesForConversation({ |
|
messages, |
|
parentMessageId, |
|
}); |
|
const formattedMessages = orderedMessages.map((message) => { |
|
let { role: _role, sender, text } = message; |
|
const role = _role ?? sender; |
|
const content = text ?? ''; |
|
return { |
|
role: role?.toLowerCase() === 'user' ? 'user' : 'assistant', |
|
content, |
|
}; |
|
}); |
|
return { |
|
prompt: formattedMessages, |
|
tokenCountMap: null, |
|
}; |
|
}); |
|
|
|
return TestClient; |
|
}; |
|
|
|
module.exports = { FakeClient, initializeFakeClient }; |
|
|