147 lines
4.1 KiB
TypeScript
147 lines
4.1 KiB
TypeScript
import { DEFAULT_MODELS } from "../constant";
|
|
import { LLMModel } from "../client/api";
|
|
|
|
const customProvider = (providerName: string) => ({
|
|
id: providerName.toLowerCase(),
|
|
providerName: providerName,
|
|
providerType: "custom",
|
|
});
|
|
|
|
export function collectModelTable(
|
|
models: readonly LLMModel[],
|
|
customModels: string,
|
|
) {
|
|
const modelTable: Record<
|
|
string,
|
|
{
|
|
available: boolean;
|
|
name: string;
|
|
displayName: string;
|
|
provider?: LLMModel["provider"]; // Marked as optional
|
|
isDefault?: boolean;
|
|
}
|
|
> = {};
|
|
|
|
// default models
|
|
models.forEach((m) => {
|
|
// using <modelName>@<providerId> as fullName
|
|
modelTable[`${m.name}@${m?.provider?.id}`] = {
|
|
...m,
|
|
displayName: m.name, // 'provider' is copied over if it exists
|
|
};
|
|
});
|
|
|
|
// server custom models
|
|
customModels
|
|
.split(",")
|
|
.filter((v) => !!v && v.length > 0)
|
|
.forEach((m) => {
|
|
const available = !m.startsWith("-");
|
|
const nameConfig =
|
|
m.startsWith("+") || m.startsWith("-") ? m.slice(1) : m;
|
|
let [name, displayName] = nameConfig.split("=");
|
|
|
|
// enable or disable all models
|
|
if (name === "all") {
|
|
Object.values(modelTable).forEach(
|
|
(model) => (model.available = available),
|
|
);
|
|
} else {
|
|
// 1. find model by name, and set available value
|
|
const [customModelName, customProviderName] = name.split("@");
|
|
let count = 0;
|
|
for (const fullName in modelTable) {
|
|
const [modelName, providerName] = fullName.split("@");
|
|
if (
|
|
customModelName == modelName &&
|
|
(customProviderName === undefined ||
|
|
customProviderName === providerName)
|
|
) {
|
|
count += 1;
|
|
modelTable[fullName]["available"] = available;
|
|
// swap name and displayName for bytedance
|
|
if (providerName === "bytedance") {
|
|
[name, displayName] = [displayName, modelName];
|
|
modelTable[fullName]["name"] = name;
|
|
}
|
|
if (displayName) {
|
|
modelTable[fullName]["displayName"] = displayName;
|
|
}
|
|
}
|
|
}
|
|
// 2. if model not exists, create new model with available value
|
|
if (count === 0) {
|
|
let [customModelName, customProviderName] = name.split("@");
|
|
const provider = customProvider(
|
|
customProviderName || customModelName,
|
|
);
|
|
// swap name and displayName for bytedance
|
|
if (displayName && provider.providerName == "ByteDance") {
|
|
[customModelName, displayName] = [displayName, customModelName];
|
|
}
|
|
modelTable[`${customModelName}@${provider?.id}`] = {
|
|
name: customModelName,
|
|
displayName: displayName || customModelName,
|
|
available,
|
|
provider, // Use optional chaining
|
|
};
|
|
}
|
|
}
|
|
});
|
|
|
|
return modelTable;
|
|
}
|
|
|
|
export function collectModelTableWithDefaultModel(
|
|
models: readonly LLMModel[],
|
|
customModels: string,
|
|
defaultModel: string,
|
|
) {
|
|
let modelTable = collectModelTable(models, customModels);
|
|
if (defaultModel && defaultModel !== "" && defaultModel in modelTable) {
|
|
modelTable[defaultModel] = {
|
|
...modelTable[defaultModel],
|
|
available: true,
|
|
isDefault: true,
|
|
};
|
|
}
|
|
return modelTable;
|
|
}
|
|
|
|
/**
|
|
* Generate full model table.
|
|
*/
|
|
export function collectModels(
|
|
models: readonly LLMModel[],
|
|
customModels: string,
|
|
) {
|
|
const modelTable = collectModelTable(models, customModels);
|
|
const allModels = Object.values(modelTable);
|
|
|
|
return allModels;
|
|
}
|
|
|
|
export function collectModelsWithDefaultModel(
|
|
models: readonly LLMModel[],
|
|
customModels: string,
|
|
defaultModel: string,
|
|
) {
|
|
const modelTable = collectModelTableWithDefaultModel(
|
|
models,
|
|
customModels,
|
|
defaultModel,
|
|
);
|
|
const allModels = Object.values(modelTable);
|
|
return allModels;
|
|
}
|
|
|
|
export function isModelAvailableInServer(
|
|
customModels: string,
|
|
modelName: string,
|
|
providerName: string,
|
|
) {
|
|
const fullName = `${modelName}@${providerName}`;
|
|
const modelTable = collectModelTable(DEFAULT_MODELS, customModels);
|
|
return modelTable[fullName]?.available === false;
|
|
}
|