model.ts 2.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105
  1. import { LLMModel } from "../client/api";
  2. const customProvider = (modelName: string) => ({
  3. id: modelName,
  4. providerName: "",
  5. providerType: "custom",
  6. });
  7. export function collectModelTable(
  8. models: readonly LLMModel[],
  9. customModels: string,
  10. ) {
  11. const modelTable: Record<
  12. string,
  13. {
  14. available: boolean;
  15. name: string;
  16. displayName: string;
  17. provider?: LLMModel["provider"]; // Marked as optional
  18. isDefault?: boolean;
  19. }
  20. > = {};
  21. // default models
  22. models.forEach((m) => {
  23. modelTable[m.name] = {
  24. ...m,
  25. displayName: m.name, // 'provider' is copied over if it exists
  26. };
  27. });
  28. // server custom models
  29. customModels
  30. .split(",")
  31. .filter((v) => !!v && v.length > 0)
  32. .forEach((m) => {
  33. const available = !m.startsWith("-");
  34. const nameConfig =
  35. m.startsWith("+") || m.startsWith("-") ? m.slice(1) : m;
  36. const [name, displayName] = nameConfig.split("=");
  37. // enable or disable all models
  38. if (name === "all") {
  39. Object.values(modelTable).forEach(
  40. (model) => (model.available = available),
  41. );
  42. } else {
  43. modelTable[name] = {
  44. name,
  45. displayName: displayName || name,
  46. available,
  47. provider: modelTable[name]?.provider ?? customProvider(name), // Use optional chaining
  48. };
  49. }
  50. });
  51. return modelTable;
  52. }
  53. export function collectModelTableWithDefaultModel(
  54. models: readonly LLMModel[],
  55. customModels: string,
  56. defaultModel: string,
  57. ) {
  58. let modelTable = collectModelTable(models, customModels);
  59. if (defaultModel && defaultModel !== "") {
  60. delete modelTable[defaultModel];
  61. modelTable[defaultModel] = {
  62. name: defaultModel,
  63. displayName: defaultModel,
  64. available: true,
  65. provider:
  66. modelTable[defaultModel]?.provider ?? customProvider(defaultModel),
  67. isDefault: true,
  68. };
  69. }
  70. return modelTable;
  71. }
  72. /**
  73. * Generate full model table.
  74. */
  75. export function collectModels(
  76. models: readonly LLMModel[],
  77. customModels: string,
  78. ) {
  79. const modelTable = collectModelTable(models, customModels);
  80. const allModels = Object.values(modelTable);
  81. return allModels;
  82. }
  83. export function collectModelsWithDefaultModel(
  84. models: readonly LLMModel[],
  85. customModels: string,
  86. defaultModel: string,
  87. ) {
  88. const modelTable = collectModelTableWithDefaultModel(
  89. models,
  90. customModels,
  91. defaultModel,
  92. );
  93. const allModels = Object.values(modelTable);
  94. return allModels;
  95. }