feat: Refactor Token Rates Configuration and Introduce Custom Overrides

This commit is contained in:
Ruben Talstra 2025-02-27 10:57:42 +01:00
parent 7dfb386f5a
commit 262e6aa4c9
No known key found for this signature in database
GPG key ID: 2A5A7174A60F3BEA
4 changed files with 106 additions and 36 deletions

View file

@ -25,10 +25,7 @@ const AppService = async (app) => {
/** @type {TCustomConfig} */
const config = (await loadCustomConfig()) ?? {};
const configDefaults = getConfigDefaults();
const tokenRatesConfig = loadTokenRatesConfig(config, configDefaults);
//
// // Set the global token rates configuration so that it can be used by the tx.js functions.
// setTokenRatesConfig(tokenRatesConfig);
loadTokenRatesConfig(config, configDefaults);
const filteredTools = config.filteredTools;
const includedTools = config.includedTools;

View file

@ -1,8 +1,9 @@
const { removeNullishValues } = require('librechat-data-provider');
const { logger } = require('~/config');
const { setCustomTokenOverrides, setCustomCacheOverrides } = require('~/models/tx');
/**
* Loads custom token rates from the user's YAML config, merging with default token rates if available.
* Loads token rates from the user's configuration, merging with default token rates if available.
*
* @param {TCustomConfig | undefined} config - The loaded custom configuration.
* @param {TConfigDefaults} [configDefaults] - Optional default configuration values.
@ -13,6 +14,8 @@ function loadTokenRatesConfig(config, configDefaults) {
if (!configDefaults?.tokenRates) {
logger.info(`User tokenRates configuration:\n${JSON.stringify(userTokenRates, null, 2)}`);
// Apply custom token rates even if there are no defaults
applyCustomTokenRates(userTokenRates);
return userTokenRates;
}
@ -20,8 +23,49 @@ function loadTokenRatesConfig(config, configDefaults) {
const defaultTokenRates = removeNullishValues(configDefaults.tokenRates);
const merged = { ...defaultTokenRates, ...userTokenRates };
// Apply custom token rates configuration
applyCustomTokenRates(merged);
logger.info(`Merged tokenRates configuration:\n${JSON.stringify(merged, null, 2)}`);
return merged;
}
/**
* Processes the token rates configuration to set up custom overrides for each model.
*
* The configuration is expected to be specified per model:
*
* For each model in the tokenRates configuration, this function will call the tx.js
* override functions to apply the custom token and cache multipliers.
*
* @param {TModelTokenRates} tokenRates - The token rates configuration mapping models to token costs.
*/
function applyCustomTokenRates(tokenRates) {
// Iterate over each model in the tokenRates configuration.
Object.keys(tokenRates).forEach((model) => {
const rate = tokenRates[model];
// If token multipliers are provided, set custom token overrides.
if (rate.prompt != null || rate.completion != null) {
setCustomTokenOverrides({
[model]: {
prompt: rate.prompt,
completion: rate.completion,
},
});
}
// Check for cache overrides.
const cacheOverrides = rate.cache;
if (cacheOverrides && (cacheOverrides.write != null || cacheOverrides.read != null)) {
setCustomCacheOverrides({
[model]: {
cache: {
write: cacheOverrides.write,
read: cacheOverrides.read,
},
},
});
}
});
}
module.exports = { loadTokenRatesConfig };