Switch branch
feat/multi-feature-pack maindefault
View all branches
Blame Raw
Cohee · 51ad27fb · · 265 lines (9.8 KB)
2 contributors 
1import { TEXTGEN_TYPES, OPENROUTER_HEADERS, FEATHERLESS_HEADERS } from './constants.js';


2import { SECRET_KEYS, readSecret } from './endpoints/secrets.js';


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


4


5/**


6 * Gets the headers for the Mancer API.


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


8 * @param {string|null} secretId Secret ID for the request (optional, used to determine which secret to use)


9 * @returns {object} Headers for the request


10 */


11function getMancerHeaders(directories, secretId = null) {


12    const apiKey = readSecret(directories, SECRET_KEYS.MANCER, secretId);


13


14    return apiKey ? ({


15        'X-API-KEY': apiKey,


16        'Authorization': `Bearer ${apiKey}`,


17    }) : {};


18}


19


20/**


21 * Gets the headers for the TogetherAI API.


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


23 * @param {string|null} secretId Secret ID for the request (optional, used to determine which secret to use)


24 * @returns {object} Headers for the request


25 */


26function getTogetherAIHeaders(directories, secretId = null) {


27    const apiKey = readSecret(directories, SECRET_KEYS.TOGETHERAI, secretId);


28


29    return apiKey ? ({


30        'Authorization': `Bearer ${apiKey}`,


31    }) : {};


32}


33


34/**


35 * Gets the headers for the InfermaticAI API.


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


37 * @param {string|null} secretId Secret ID for the request (optional, used to determine which secret to use)


38 * @returns {object} Headers for the request


39 */


40function getInfermaticAIHeaders(directories, secretId = null) {


41    const apiKey = readSecret(directories, SECRET_KEYS.INFERMATICAI, secretId);


42


43    return apiKey ? ({


44        'Authorization': `Bearer ${apiKey}`,


45    }) : {};


46}


47


48/**


49 * Gets the headers for the DreamGen API.


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


51 * @param {string|null} secretId Secret ID for the request (optional, used to determine which secret to use)


52 * @returns {object} Headers for the request


53 */


54function getDreamGenHeaders(directories, secretId = null) {


55    const apiKey = readSecret(directories, SECRET_KEYS.DREAMGEN, secretId);


56


57    return apiKey ? ({


58        'Authorization': `Bearer ${apiKey}`,


59    }) : {};


60}


61


62/**


63 * Gets the headers for the OpenRouter API.


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


65 * @param {string|null} secretId Secret ID for the request (optional, used to determine which secret to use)


66 * @returns {object} Headers for the request


67 */


68function getOpenRouterHeaders(directories, secretId = null) {


69    const apiKey = readSecret(directories, SECRET_KEYS.OPENROUTER, secretId);


70    const baseHeaders = { ...OPENROUTER_HEADERS };


71


72    return apiKey ? Object.assign(baseHeaders, { 'Authorization': `Bearer ${apiKey}` }) : baseHeaders;


73}


74


75/**


76 * Gets the headers for the vLLM API.


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


78 * @param {string|null} secretId Secret ID for the request (optional, used to determine which secret to use)


79 * @returns {object} Headers for the request


80 */


81function getVllmHeaders(directories, secretId = null) {


82    const apiKey = readSecret(directories, SECRET_KEYS.VLLM, secretId);


83


84    return apiKey ? ({


85        'Authorization': `Bearer ${apiKey}`,


86    }) : {};


87}


88


89/**


90 * Gets the headers for the Aphrodite API.


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


92 * @param {string|null} secretId Secret ID for the request (optional, used to determine which secret to use)


93 * @returns {object} Headers for the request


94 */


95function getAphroditeHeaders(directories, secretId = null) {


96    const apiKey = readSecret(directories, SECRET_KEYS.APHRODITE, secretId);


97


98    return apiKey ? ({


99        'X-API-KEY': apiKey,


100        'Authorization': `Bearer ${apiKey}`,


101    }) : {};


102}


103


104/**


105 * Gets the headers for the Tabby API.


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


107 * @param {string|null} secretId Secret ID for the request (optional, used to determine which secret to use)


108 * @returns {object} Headers for the request


109 */


110function getTabbyHeaders(directories, secretId = null) {


111    const apiKey = readSecret(directories, SECRET_KEYS.TABBY, secretId);


112


113    return apiKey ? ({


114        'x-api-key': apiKey,


115        'Authorization': `Bearer ${apiKey}`,


116    }) : {};


117}


118


119/**


120 * Gets the headers for the LlamaCPP API.


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


122 * @param {string|null} secretId Secret ID for the request (optional, used to determine which secret to use)


123 * @returns {object} Headers for the request


124 */


125function getLlamaCppHeaders(directories, secretId = null) {


126    const apiKey = readSecret(directories, SECRET_KEYS.LLAMACPP, secretId);


127


128    return apiKey ? ({


129        'Authorization': `Bearer ${apiKey}`,


130    }) : {};


131}


132


133/**


134 * Gets the headers for the Ooba API.


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


136 * @param {string|null} secretId Secret ID for the request (optional, used to determine which secret to use)


137 * @returns {object} Headers for the request


138 */


139function getOobaHeaders(directories, secretId = null) {


140    const apiKey = readSecret(directories, SECRET_KEYS.OOBA, secretId);


141


142    return apiKey ? ({


143        'Authorization': `Bearer ${apiKey}`,


144    }) : {};


145}


146


147/**


148 * Gets the headers for the KoboldCpp API.


149 * @param {import('./users.js').UserDirectoryList} directories


150 * @param {string|null} secretId Secret ID for the request (optional, used to determine which secret to use)


151 * @returns {object} Headers for the request


152 */


153function getKoboldCppHeaders(directories, secretId = null) {


154    const apiKey = readSecret(directories, SECRET_KEYS.KOBOLDCPP, secretId);


155


156    return apiKey ? ({


157        'Authorization': `Bearer ${apiKey}`,


158    }) : {};


159}


160


161/**


162 * Gets the headers for the Featherless API.


163 * @param {import('./users.js').UserDirectoryList} directories


164 * @param {string|null} secretId Secret ID for the request (optional, used to determine which secret to use)


165 * @returns {object} Headers for the request


166 */


167function getFeatherlessHeaders(directories, secretId = null) {


168    const apiKey = readSecret(directories, SECRET_KEYS.FEATHERLESS, secretId);


169    const baseHeaders = { ...FEATHERLESS_HEADERS };


170


171    return apiKey ? Object.assign(baseHeaders, { 'Authorization': `Bearer ${apiKey}` }) : baseHeaders;


172}


173


174/**


175 * Gets the headers for the HuggingFace API.


176 * @param {import('./users.js').UserDirectoryList} directories


177 * @param {string|null} secretId Secret ID for the request (optional, used to determine which secret to use)


178 * @returns {object} Headers for the request


179 */


180function getHuggingFaceHeaders(directories, secretId = null) {


181    const apiKey = readSecret(directories, SECRET_KEYS.HUGGINGFACE, secretId);


182


183    return apiKey ? ({


184        'Authorization': `Bearer ${apiKey}`,


185    }) : {};


186}


187


188/**


189 * Gets the headers for the Generic text completion API.


190 * @param {import('./users.js').UserDirectoryList} directories


191 * @param {string|null} secretId Secret ID for the request (optional, used to determine which secret to use)


192 * @returns {object} Headers for the request


193 */


194function getGenericHeaders(directories, secretId = null) {


195    const apiKey = readSecret(directories, SECRET_KEYS.GENERIC, secretId);


196


197    return apiKey ? ({


198        'Authorization': `Bearer ${apiKey}`,


199    }) : {};


200}


201


202export function getOverrideHeaders(urlHost) {


203    const requestOverrides = getConfigValue('requestOverrides', []);


204    const overrideHeaders = requestOverrides?.find((e) => e.hosts?.includes(urlHost))?.headers;


205    if (overrideHeaders && urlHost) {


206        return overrideHeaders;


207    } else {


208        return {};


209    }


210}


211


212/**


213 * Sets additional headers for the request.


214 * @param {import('express').Request} request Original request body


215 * @param {object} args New request arguments


216 * @param {string|null} server API server for new request


217 */


218export function setAdditionalHeaders(request, args, server) {


219    setAdditionalHeadersByType(args.headers, request.body.api_type, server, request.user.directories, request.body.secret_id);


220}


221


222/**


223 *


224 * @param {object} requestHeaders Request headers


225 * @param {string} type API type


226 * @param {string|null} server API server for new request


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


228 * @param {string|null} secretId Secret ID for the request (optional, used for some API types to determine which secret to use)


229 */


230export function setAdditionalHeadersByType(requestHeaders, type, server, directories, secretId = null) {


231    const headerGetters = {


232        [TEXTGEN_TYPES.MANCER]: getMancerHeaders,


233        [TEXTGEN_TYPES.VLLM]: getVllmHeaders,


234        [TEXTGEN_TYPES.APHRODITE]: getAphroditeHeaders,


235        [TEXTGEN_TYPES.TABBY]: getTabbyHeaders,


236        [TEXTGEN_TYPES.TOGETHERAI]: getTogetherAIHeaders,


237        [TEXTGEN_TYPES.OOBA]: getOobaHeaders,


238        [TEXTGEN_TYPES.INFERMATICAI]: getInfermaticAIHeaders,


239        [TEXTGEN_TYPES.DREAMGEN]: getDreamGenHeaders,


240        [TEXTGEN_TYPES.OPENROUTER]: getOpenRouterHeaders,


241        [TEXTGEN_TYPES.KOBOLDCPP]: getKoboldCppHeaders,


242        [TEXTGEN_TYPES.LLAMACPP]: getLlamaCppHeaders,


243        [TEXTGEN_TYPES.FEATHERLESS]: getFeatherlessHeaders,


244        [TEXTGEN_TYPES.HUGGINGFACE]: getHuggingFaceHeaders,


245        [TEXTGEN_TYPES.GENERIC]: getGenericHeaders,


246    };


247


248    const getHeaders = headerGetters[type];


249    const headers = getHeaders ? getHeaders(directories, secretId) : {};


250


251    if (typeof server === 'string' && server.length > 0) {


252        try {


253            const url = new URL(server);


254            const overrideHeaders = getOverrideHeaders(url.host);


255


256            if (overrideHeaders && Object.keys(overrideHeaders).length > 0) {


257                Object.assign(headers, overrideHeaders);


258            }


259        } catch {


260            // Do nothing


261        }


262    }


263


264    Object.assign(requestHeaders, headers);


265}