rixty/SillyTavern
1
0
Fork 0
mirror of https://github.com/SillyTavern/SillyTavern.git synced 2025-06-05 21:59:27 +02:00
Code Issues Packages Projects Releases Wiki Activity

added mistral vector support (off the back of oai's)

Browse Source
This commit is contained in:
based
2023-12-16 08:37:39 +10:00
parent c517483141
commit 5dd2e8cd88
4 changed files with 12 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
public/scripts/extensions/vectors/index.js
View File

@ -394,7 +394,8 @@ async function getSavedHashes(collectionId) {
*/ */
async function insertVectorItems(collectionId, items) { async function insertVectorItems(collectionId, items) {
if (settings.source === 'openai' && !secret_state[SECRET_KEYS.OPENAI] || if (settings.source === 'openai' && !secret_state[SECRET_KEYS.OPENAI] ||
settings.source === 'palm' && !secret_state[SECRET_KEYS.MAKERSUITE]) { settings.source === 'palm' && !secret_state[SECRET_KEYS.MAKERSUITE] ||
settings.source === 'mistral' && !secret_state[SECRET_KEYS.MISTRALAI]) {
throw new Error('Vectors: API key missing', { cause: 'api_key_missing' }); throw new Error('Vectors: API key missing', { cause: 'api_key_missing' });
} }

1
public/scripts/extensions/vectors/settings.html
View File

@ -13,6 +13,7 @@
<option value="transformers">Local (Transformers)</option> <option value="transformers">Local (Transformers)</option>
<option value="openai">OpenAI</option> <option value="openai">OpenAI</option>
<option value="palm">Google MakerSuite (PaLM)</option> <option value="palm">Google MakerSuite (PaLM)</option>
<option value="mistral">MistralAI</option>
</select> </select>
</div> </div>

3
src/endpoints/vectors.js
View File

@ -12,8 +12,9 @@ const { jsonParser } = require('../express-common');
*/ */
async function getVector(source, text) { async function getVector(source, text) {
switch (source) { switch (source) {
case 'mistral':
case 'openai': case 'openai':
return require('../openai-vectors').getOpenAIVector(text); return require('../openai-vectors').getOpenAIVector(text, source);
case 'transformers': case 'transformers':
return require('../embedding').getTransformersVector(text); return require('../embedding').getTransformersVector(text);
case 'palm': case 'palm':

12
src/openai-vectors.js
View File

@ -2,19 +2,21 @@ const fetch = require('node-fetch').default;
const { SECRET_KEYS, readSecret } = require('./endpoints/secrets'); const { SECRET_KEYS, readSecret } = require('./endpoints/secrets');
/** /**
* Gets the vector for the given text from OpenAI ada model * Gets the vector for the given text from an OpenAI compatible endpoint.
* @param {string} text - The text to get the vector for * @param {string} text - The text to get the vector for
* @returns {Promise<number[]>} - The vector for the text * @returns {Promise<number[]>} - The vector for the text
*/ */
async function getOpenAIVector(text) { async function getOpenAIVector(text, source) {
const key = readSecret(SECRET_KEYS.OPENAI); const isMistral = source === 'mistral';
const key = readSecret(isMistral ? SECRET_KEYS.MISTRALAI : SECRET_KEYS.OPENAI);
if (!key) { if (!key) {
console.log('No OpenAI key found'); console.log('No OpenAI key found');
throw new Error('No OpenAI key found'); throw new Error('No OpenAI key found');
} }
const response = await fetch('https://api.openai.com/v1/embeddings', { const url = isMistral ? 'api.mistral.ai' : 'api.openai.com';
const response = await fetch(`https://${url}/v1/embeddings`, {
method: 'POST', method: 'POST',
headers: { headers: {
'Content-Type': 'application/json', 'Content-Type': 'application/json',
@ -22,7 +24,7 @@ async function getOpenAIVector(text) {
}, },
body: JSON.stringify({ body: JSON.stringify({
input: text, input: text,
model: 'text-embedding-ada-002', model: isMistral ? 'mistral-embed' : 'text-embedding-ada-002',
}), }),
}); });