Switch branch
feat/multi-feature-pack maindefault
View all branches
Blame Raw
Cohee · 51ad27fb · · 605 lines (24.0 KB)
2 contributors 
1import path from 'node:path';


2import fs from 'node:fs';


3


4import vectra from 'vectra';


5import express from 'express';


6import sanitize from 'sanitize-filename';


7


8import { getConfigValue } from '../util.js';


9


10import { getNomicAIBatchVector, getNomicAIVector } from '../vectors/nomicai-vectors.js';


11import { getOpenAIVector, getOpenAIBatchVector } from '../vectors/openai-vectors.js';


12import { getTransformersVector, getTransformersBatchVector } from '../vectors/embedding.js';


13import { getExtrasVector, getExtrasBatchVector } from '../vectors/extras-vectors.js';


14import { getMakerSuiteVector, getMakerSuiteBatchVector } from '../vectors/google-vectors.js';


15import { getVertexVector, getVertexBatchVector } from '../vectors/google-vectors.js';


16import { getCohereVector, getCohereBatchVector } from '../vectors/cohere-vectors.js';


17import { getLlamaCppVector, getLlamaCppBatchVector } from '../vectors/llamacpp-vectors.js';


18import { getVllmVector, getVllmBatchVector } from '../vectors/vllm-vectors.js';


19import { getOllamaVector, getOllamaBatchVector } from '../vectors/ollama-vectors.js';


20


21// Don't forget to add new sources to the SOURCES array


22const SOURCES = [


23    'transformers',


24    'mistral',


25    'openai',


26    'extras',


27    'palm',


28    'togetherai',


29    'nomicai',


30    'cohere',


31    'ollama',


32    'llamacpp',


33    'vllm',


34    'webllm',


35    'koboldcpp',


36    'vertexai',


37    'electronhub',


38    'openrouter',


39    'chutes',


40    'nanogpt',


41    'siliconflow',


42    'workers_ai',


43];


44


45/**


46 * Gets the vector for the given text from the given source.


47 * @param {string} source - The source of the vector


48 * @param {Object} sourceSettings - Settings for the source, if it needs any


49 * @param {string} text - The text to get the vector for


50 * @param {boolean} isQuery - If the text is a query for embedding search


51 * @param {import('../users.js').UserDirectoryList} directories - The directories object for the user


52 * @returns {Promise<number[]>} - The vector for the text


53 */


54async function getVector(source, sourceSettings, text, isQuery, directories) {


55    switch (source) {


56        case 'nomicai':


57            return getNomicAIVector(text, source, directories);


58        case 'togetherai':


59        case 'mistral':


60        case 'openai':


61            return getOpenAIVector(text, source, directories, sourceSettings.model);


62        case 'electronhub':


63            return getOpenAIVector(text, source, directories, sourceSettings.model);


64        case 'openrouter':


65            return getOpenAIVector(text, source, directories, sourceSettings.model);


66        case 'transformers':


67            return getTransformersVector(text);


68        case 'extras':


69            return getExtrasVector(text, sourceSettings.extrasUrl, sourceSettings.extrasKey);


70        case 'palm':


71            return getMakerSuiteVector(text, sourceSettings.model, sourceSettings.request);


72        case 'vertexai':


73            return getVertexVector(text, sourceSettings.model, sourceSettings.request);


74        case 'cohere':


75            return getCohereVector(text, isQuery, directories, sourceSettings.model);


76        case 'llamacpp':


77            return getLlamaCppVector(text, sourceSettings.apiUrl, directories);


78        case 'vllm':


79            return getVllmVector(text, sourceSettings.apiUrl, sourceSettings.model, directories);


80        case 'ollama':


81            return getOllamaVector(text, sourceSettings.apiUrl, sourceSettings.model, sourceSettings.keep, directories);


82        case 'webllm':


83            return sourceSettings.embeddings[text];


84        case 'koboldcpp':


85            return sourceSettings.embeddings[text];


86        case 'chutes':


87            return getOpenAIVector(text, source, directories, sourceSettings.model);


88        case 'nanogpt':


89            return getOpenAIVector(text, source, directories, sourceSettings.model);


90        case 'siliconflow':


91            return getOpenAIVector(text, source, directories, sourceSettings.model, sourceSettings.urlOverride);


92        case 'workers_ai':


93            return getOpenAIVector(text, source, directories, sourceSettings.model, sourceSettings.urlOverride);


94    }


95


96    throw new Error(`Unknown vector source ${source}`);


97}


98


99/**


100 * Gets the vector for the given text batch from the given source.


101 * @param {string} source - The source of the vector


102 * @param {Object} sourceSettings - Settings for the source, if it needs any


103 * @param {string[]} texts - The array of texts to get the vector for


104 * @param {boolean} isQuery - If the text is a query for embedding search


105 * @param {import('../users.js').UserDirectoryList} directories - The directories object for the user


106 * @returns {Promise<number[][]>} - The array of vectors for the texts


107 */


108async function getBatchVector(source, sourceSettings, texts, isQuery, directories) {


109    const batchSize = 10;


110    const batches = Array(Math.ceil(texts.length / batchSize)).fill(undefined).map((_, i) => texts.slice(i * batchSize, i * batchSize + batchSize));


111


112    let results = [];


113    for (let batch of batches) {


114        switch (source) {


115            case 'nomicai':


116                results.push(...await getNomicAIBatchVector(batch, source, directories));


117                break;


118            case 'togetherai':


119            case 'mistral':


120            case 'openai':


121                results.push(...await getOpenAIBatchVector(batch, source, directories, sourceSettings.model));


122                break;


123            case 'electronhub':


124                results.push(...await getOpenAIBatchVector(batch, source, directories, sourceSettings.model));


125                break;


126            case 'openrouter':


127                results.push(...await getOpenAIBatchVector(batch, source, directories, sourceSettings.model));


128                break;


129            case 'transformers':


130                results.push(...await getTransformersBatchVector(batch));


131                break;


132            case 'extras':


133                results.push(...await getExtrasBatchVector(batch, sourceSettings.extrasUrl, sourceSettings.extrasKey));


134                break;


135            case 'palm':


136                results.push(...await getMakerSuiteBatchVector(batch, sourceSettings.model, sourceSettings.request));


137                break;


138            case 'vertexai':


139                results.push(...await getVertexBatchVector(batch, sourceSettings.model, sourceSettings.request));


140                break;


141            case 'cohere':


142                results.push(...await getCohereBatchVector(batch, isQuery, directories, sourceSettings.model));


143                break;


144            case 'llamacpp':


145                results.push(...await getLlamaCppBatchVector(batch, sourceSettings.apiUrl, directories));


146                break;


147            case 'vllm':


148                results.push(...await getVllmBatchVector(batch, sourceSettings.apiUrl, sourceSettings.model, directories));


149                break;


150            case 'ollama':


151                results.push(...await getOllamaBatchVector(batch, sourceSettings.apiUrl, sourceSettings.model, sourceSettings.keep, directories));


152                break;


153            case 'webllm':


154                results.push(...texts.map(x => sourceSettings.embeddings[x]));


155                break;


156            case 'koboldcpp':


157                results.push(...texts.map(x => sourceSettings.embeddings[x]));


158                break;


159            case 'chutes':


160                results.push(...await getOpenAIBatchVector(batch, source, directories, sourceSettings.model));


161                break;


162            case 'nanogpt':


163                results.push(...await getOpenAIBatchVector(batch, source, directories, sourceSettings.model));


164                break;


165            case 'siliconflow':


166                results.push(...await getOpenAIBatchVector(batch, source, directories, sourceSettings.model, sourceSettings.urlOverride));


167                break;


168            case 'workers_ai':


169                results.push(...await getOpenAIBatchVector(batch, source, directories, sourceSettings.model, sourceSettings.urlOverride));


170                break;


171            default:


172                throw new Error(`Unknown vector source ${source}`);


173        }


174    }


175


176    return results;


177}


178


179/**


180 * Extracts settings for the vectorization sources from the HTTP request headers.


181 * @param {string} source - Which source to extract settings for.


182 * @param {object} request - The HTTP request object.


183 * @returns {object} - An object that can be used as `sourceSettings` in functions that take that parameter.


184 */


185function getSourceSettings(source, request) {


186    switch (source) {


187        case 'togetherai':


188            return {


189                model: String(request.body.model),


190            };


191        case 'openai':


192            return {


193                model: String(request.body.model),


194            };


195        case 'electronhub':


196            return {


197                model: String(request.body.model || 'text-embedding-3-small'),


198            };


199        case 'openrouter':


200            return {


201                model: String(request.body.model) || 'openai/text-embedding-3-large',


202            };


203        case 'cohere':


204            return {


205                model: String(request.body.model),


206            };


207        case 'llamacpp':


208            return {


209                apiUrl: String(request.body.apiUrl),


210            };


211        case 'vllm':


212            return {


213                apiUrl: String(request.body.apiUrl),


214                model: String(request.body.model),


215            };


216        case 'ollama':


217            return {


218                apiUrl: String(request.body.apiUrl),


219                model: String(request.body.model),


220                keep: Boolean(request.body.keep),


221            };


222        case 'extras':


223            return {


224                extrasUrl: String(request.body.extrasUrl),


225                extrasKey: String(request.body.extrasKey),


226            };


227        case 'transformers':


228            return {


229                model: getConfigValue('extensions.models.embedding', ''),


230            };


231        case 'palm':


232        case 'vertexai':


233            return {


234                model: String(request.body.model || 'text-embedding-005'),


235                request: request, // Pass the request object to get API key and URL


236            };


237        case 'mistral':


238            return {


239                model: 'mistral-embed',


240            };


241        case 'nomicai':


242            return {


243                model: 'nomic-embed-text-v1.5',


244            };


245        case 'webllm':


246            return {


247                model: String(request.body.model),


248                embeddings: request.body.embeddings ?? {},


249            };


250        case 'koboldcpp':


251            return {


252                model: String(request.body.model),


253                embeddings: request.body.embeddings ?? {},


254            };


255        case 'chutes':


256            return {


257                model: String(request.body.model || 'chutes-qwen-qwen3-embedding-8b'),


258            };


259        case 'nanogpt':


260            return {


261                model: String(request.body.model || 'text-embedding-3-small'),


262            };


263        case 'siliconflow':


264            return {


265                model: String(request.body.model || 'Qwen/Qwen3-Embedding-0.6B'),


266                urlOverride: request.body.siliconflow_endpoint === 'cn'


267                    ? 'https://api.siliconflow.cn/v1' : null,


268            };


269        case 'workers_ai': {


270            const accountId = String(request.body.workers_ai_account_id || '').trim();


271            return {


272                model: String(request.body.model || '@cf/baai/bge-m3'),


273                urlOverride: accountId


274                    ? `https://api.cloudflare.com/client/v4/accounts/${encodeURIComponent(accountId)}/ai/v1`


275                    : null,


276            };


277        }


278        default:


279            return {};


280    }


281}


282


283/**


284 * Gets the model scope for the source.


285 * @param {object} sourceSettings - The settings for the source


286 * @returns {string} The model scope for the source


287 */


288function getModelScope(sourceSettings) {


289    return (sourceSettings?.model || '');


290}


291


292/**


293 * Gets the index for the vector collection


294 * @param {import('../users.js').UserDirectoryList} directories - User directories


295 * @param {string} collectionId - The collection ID


296 * @param {string} source - The source of the vector


297 * @param {object} sourceSettings - The model for the source


298 * @returns {Promise<vectra.LocalIndex>} - The index for the collection


299 */


300async function getIndex(directories, collectionId, source, sourceSettings) {


301    const model = getModelScope(sourceSettings);


302    const pathToFile = path.join(directories.vectors, sanitize(source), sanitize(collectionId), sanitize(model));


303    const store = new vectra.LocalIndex(pathToFile);


304


305    if (!await store.isIndexCreated()) {


306        await store.createIndex();


307    }


308


309    return store;


310}


311


312/**


313 * Inserts items into the vector collection


314 * @param {import('../users.js').UserDirectoryList} directories - User directories


315 * @param {string} collectionId - The collection ID


316 * @param {string} source - The source of the vector


317 * @param {Object} sourceSettings - Settings for the source, if it needs any


318 * @param {{ hash: number; text: string; index: number; }[]} items - The items to insert


319 */


320async function insertVectorItems(directories, collectionId, source, sourceSettings, items) {


321    const store = await getIndex(directories, collectionId, source, sourceSettings);


322


323    await store.beginUpdate();


324


325    const vectors = await getBatchVector(source, sourceSettings, items.map(x => x.text), false, directories);


326


327    for (let i = 0; i < items.length; i++) {


328        const item = items[i];


329        const vector = vectors[i];


330        await store.upsertItem({ vector: vector, metadata: { hash: item.hash, text: item.text, index: item.index } });


331    }


332


333    await store.endUpdate();


334}


335


336/**


337 * Gets the hashes of the items in the vector collection


338 * @param {import('../users.js').UserDirectoryList} directories - User directories


339 * @param {string} collectionId - The collection ID


340 * @param {string} source - The source of the vector


341 * @param {Object} sourceSettings - Settings for the source, if it needs any


342 * @returns {Promise<number[]>} - The hashes of the items in the collection


343 */


344async function getSavedHashes(directories, collectionId, source, sourceSettings) {


345    const store = await getIndex(directories, collectionId, source, sourceSettings);


346


347    const items = await store.listItems();


348    const hashes = items.map(x => Number(x.metadata.hash));


349


350    return hashes;


351}


352


353/**


354 * Deletes items from the vector collection by hash


355 * @param {import('../users.js').UserDirectoryList} directories - User directories


356 * @param {string} collectionId - The collection ID


357 * @param {string} source - The source of the vector


358 * @param {Object} sourceSettings - Settings for the source, if it needs any


359 * @param {number[]} hashes - The hashes of the items to delete


360 */


361async function deleteVectorItems(directories, collectionId, source, sourceSettings, hashes) {


362    const store = await getIndex(directories, collectionId, source, sourceSettings);


363    const items = await store.listItemsByMetadata({ hash: { '$in': hashes } });


364


365    await store.beginUpdate();


366


367    for (const item of items) {


368        await store.deleteItem(item.id);


369    }


370


371    await store.endUpdate();


372}


373


374/**


375 * Gets the hashes of the items in the vector collection that match the search text


376 * @param {import('../users.js').UserDirectoryList} directories - User directories


377 * @param {string} collectionId - The collection ID


378 * @param {string} source - The source of the vector


379 * @param {Object} sourceSettings - Settings for the source, if it needs any


380 * @param {string} searchText - The text to search for


381 * @param {number} topK - The number of results to return


382 * @param {number} threshold - The threshold for the search


383 * @returns {Promise<{hashes: number[], metadata: object[]}>} - The metadata of the items that match the search text


384 */


385async function queryCollection(directories, collectionId, source, sourceSettings, searchText, topK, threshold) {


386    const store = await getIndex(directories, collectionId, source, sourceSettings);


387    const vector = await getVector(source, sourceSettings, searchText, true, directories);


388


389    const result = await store.queryItems(vector, topK);


390    const metadata = result.filter(x => x.score >= threshold).map(x => x.item.metadata);


391    const hashes = result.map(x => Number(x.item.metadata.hash));


392    return { metadata, hashes };


393}


394


395/**


396 * Queries multiple collections for the given search queries. Returns the overall top K results.


397 * @param {import('../users.js').UserDirectoryList} directories - User directories


398 * @param {string[]} collectionIds - The collection IDs to query


399 * @param {string} source - The source of the vector


400 * @param {Object} sourceSettings - Settings for the source, if it needs any


401 * @param {string} searchText - The text to search for


402 * @param {number} topK - The number of results to return


403 * @param {number} threshold - The threshold for the search


404 *


405 * @returns {Promise<Record<string, { hashes: number[], metadata: object[] }>>} - The top K results from each collection


406 */


407async function multiQueryCollection(directories, collectionIds, source, sourceSettings, searchText, topK, threshold) {


408    const vector = await getVector(source, sourceSettings, searchText, true, directories);


409    const results = [];


410


411    for (const collectionId of collectionIds) {


412        const store = await getIndex(directories, collectionId, source, sourceSettings);


413        const result = await store.queryItems(vector, topK);


414        results.push(...result.map(result => ({ collectionId, result })));


415    }


416


417    // Sort results by descending similarity, apply threshold, and take top K


418    const sortedResults = results


419        .sort((a, b) => b.result.score - a.result.score)


420        .filter(x => x.result.score >= threshold)


421        .slice(0, topK);


422


423    /**


424     * Group the results by collection ID


425     * @type {Record<string, { hashes: number[], metadata: object[] }>}


426     */


427    const groupedResults = {};


428    for (const result of sortedResults) {


429        if (!groupedResults[result.collectionId]) {


430            groupedResults[result.collectionId] = { hashes: [], metadata: [] };


431        }


432


433        groupedResults[result.collectionId].hashes.push(Number(result.result.item.metadata.hash));


434        groupedResults[result.collectionId].metadata.push(result.result.item.metadata);


435    }


436


437    return groupedResults;


438}


439


440/**


441 * Performs a request to regenerate the index if it is corrupted.


442 * @param {import('express').Request} req Express request object


443 * @param {import('express').Response} res Express response object


444 * @param {Error} error Error object


445 * @returns {Promise<any>} Promise


446 */


447async function regenerateCorruptedIndexErrorHandler(req, res, error) {


448    if (error instanceof SyntaxError && !req.query.regenerated) {


449        const collectionId = String(req.body.collectionId);


450        const source = String(req.body.source) || 'transformers';


451        const sourceSettings = getSourceSettings(source, req);


452


453        if (collectionId && source) {


454            const index = await getIndex(req.user.directories, collectionId, source, sourceSettings);


455            const exists = await index.isIndexCreated();


456


457            if (exists) {


458                const path = index.folderPath;


459                console.warn(`Corrupted index detected at ${path}, regenerating...`);


460                await index.deleteIndex();


461                return res.redirect(307, req.originalUrl + '?regenerated=true');


462            }


463        }


464    }


465


466    console.error(error);


467    return res.sendStatus(500);


468}


469


470export const router = express.Router();


471


472router.post('/query', async (req, res) => {


473    try {


474        if (!req.body.collectionId || !req.body.searchText) {


475            return res.sendStatus(400);


476        }


477


478        const collectionId = String(req.body.collectionId);


479        const searchText = String(req.body.searchText);


480        const topK = Number(req.body.topK) || 10;


481        const threshold = Number(req.body.threshold) || 0.0;


482        const source = String(req.body.source) || 'transformers';


483        const sourceSettings = getSourceSettings(source, req);


484


485        const results = await queryCollection(req.user.directories, collectionId, source, sourceSettings, searchText, topK, threshold);


486        return res.json(results);


487    } catch (error) {


488        return regenerateCorruptedIndexErrorHandler(req, res, error);


489    }


490});


491


492router.post('/query-multi', async (req, res) => {


493    try {


494        if (!Array.isArray(req.body.collectionIds) || !req.body.searchText) {


495            return res.sendStatus(400);


496        }


497


498        const collectionIds = req.body.collectionIds.map(x => String(x));


499        const searchText = String(req.body.searchText);


500        const topK = Number(req.body.topK) || 10;


501        const threshold = Number(req.body.threshold) || 0.0;


502        const source = String(req.body.source) || 'transformers';


503        const sourceSettings = getSourceSettings(source, req);


504


505        const results = await multiQueryCollection(req.user.directories, collectionIds, source, sourceSettings, searchText, topK, threshold);


506        return res.json(results);


507    } catch (error) {


508        return regenerateCorruptedIndexErrorHandler(req, res, error);


509    }


510});


511


512router.post('/insert', async (req, res) => {


513    try {


514        if (!Array.isArray(req.body.items) || !req.body.collectionId) {


515            return res.sendStatus(400);


516        }


517


518        const collectionId = String(req.body.collectionId);


519        const items = req.body.items.map(x => ({ hash: x.hash, text: x.text, index: x.index }));


520        const source = String(req.body.source) || 'transformers';


521        const sourceSettings = getSourceSettings(source, req);


522


523        await insertVectorItems(req.user.directories, collectionId, source, sourceSettings, items);


524        return res.sendStatus(200);


525    } catch (error) {


526        return regenerateCorruptedIndexErrorHandler(req, res, error);


527    }


528});


529


530router.post('/list', async (req, res) => {


531    try {


532        if (!req.body.collectionId) {


533            return res.sendStatus(400);


534        }


535


536        const collectionId = String(req.body.collectionId);


537        const source = String(req.body.source) || 'transformers';


538        const sourceSettings = getSourceSettings(source, req);


539


540        const hashes = await getSavedHashes(req.user.directories, collectionId, source, sourceSettings);


541        return res.json(hashes);


542    } catch (error) {


543        return regenerateCorruptedIndexErrorHandler(req, res, error);


544    }


545});


546


547router.post('/delete', async (req, res) => {


548    try {


549        if (!Array.isArray(req.body.hashes) || !req.body.collectionId) {


550            return res.sendStatus(400);


551        }


552


553        const collectionId = String(req.body.collectionId);


554        const hashes = req.body.hashes.map(x => Number(x));


555        const source = String(req.body.source) || 'transformers';


556        const sourceSettings = getSourceSettings(source, req);


557


558        await deleteVectorItems(req.user.directories, collectionId, source, sourceSettings, hashes);


559        return res.sendStatus(200);


560    } catch (error) {


561        return regenerateCorruptedIndexErrorHandler(req, res, error);


562    }


563});


564


565router.post('/purge-all', async (req, res) => {


566    try {


567        for (const source of SOURCES) {


568            const sourcePath = path.join(req.user.directories.vectors, sanitize(source));


569            if (!fs.existsSync(sourcePath)) {


570                continue;


571            }


572            await fs.promises.rm(sourcePath, { recursive: true });


573            console.info(`Deleted vector source store at ${sourcePath}`);


574        }


575


576        return res.sendStatus(200);


577    } catch (error) {


578        console.error(error);


579        return res.sendStatus(500);


580    }


581});


582


583router.post('/purge', async (req, res) => {


584    try {


585        if (!req.body.collectionId) {


586            return res.sendStatus(400);


587        }


588


589        const collectionId = String(req.body.collectionId);


590


591        for (const source of SOURCES) {


592            const sourcePath = path.join(req.user.directories.vectors, sanitize(source), sanitize(collectionId));


593            if (!fs.existsSync(sourcePath)) {


594                continue;


595            }


596            await fs.promises.rm(sourcePath, { recursive: true });


597            console.info(`Deleted vector index at ${sourcePath}`);


598        }


599


600        return res.sendStatus(200);


601    } catch (error) {


602        console.error(error);


603        return res.sendStatus(500);


604    }


605});