Просмотр исходного кода

feat: CUSTOM_MODELS support mapper

Yidadaa 1 год назад
Родитель
Сommit
a5a1f2e8ad
3 измененных файлов с 27 добавлено и 17 удалено
  1. 1 1
      app/api/common.ts
  2. 5 5
      app/components/chat.tsx
  3. 21 11
      app/utils/model.ts

+ 1 - 1
app/api/common.ts

@@ -81,7 +81,7 @@ export async function requestOpenai(req: NextRequest) {
       const jsonBody = JSON.parse(clonedBody) as { model?: string };
 
       // not undefined and is false
-      if (modelTable[jsonBody?.model ?? ""] === false) {
+      if (modelTable[jsonBody?.model ?? ""].available === false) {
         return NextResponse.json(
           {
             error: true,

+ 5 - 5
app/components/chat.tsx

@@ -433,7 +433,7 @@ export function ChatActions(props: {
   const currentModel = chatStore.currentSession().mask.modelConfig.model;
   const allModels = useAllModels();
   const models = useMemo(
-    () => allModels.filter((m) => m.available).map((m) => m.name),
+    () => allModels.filter((m) => m.available),
     [allModels],
   );
   const [showModelSelector, setShowModelSelector] = useState(false);
@@ -441,9 +441,9 @@ export function ChatActions(props: {
   useEffect(() => {
     // if current model is not available
     // switch to first available model
-    const isUnavaliableModel = !models.includes(currentModel);
+    const isUnavaliableModel = !models.some((m) => m.name === currentModel);
     if (isUnavaliableModel && models.length > 0) {
-      const nextModel = models[0] as ModelType;
+      const nextModel = models[0].name as ModelType;
       chatStore.updateCurrentSession(
         (session) => (session.mask.modelConfig.model = nextModel),
       );
@@ -531,8 +531,8 @@ export function ChatActions(props: {
         <Selector
           defaultSelectedValue={currentModel}
           items={models.map((m) => ({
-            title: m,
-            value: m,
+            title: m.displayName,
+            value: m.name,
           }))}
           onClose={() => setShowModelSelector(false)}
           onSelection={(s) => {

+ 21 - 11
app/utils/model.ts

@@ -4,21 +4,34 @@ export function collectModelTable(
   models: readonly LLMModel[],
   customModels: string,
 ) {
-  const modelTable: Record<string, boolean> = {};
+  const modelTable: Record<
+    string,
+    { available: boolean; name: string; displayName: string }
+  > = {};
 
   // default models
-  models.forEach((m) => (modelTable[m.name] = m.available));
+  models.forEach(
+    (m) =>
+      (modelTable[m.name] = {
+        ...m,
+        displayName: m.name,
+      }),
+  );
 
   // server custom models
   customModels
     .split(",")
     .filter((v) => !!v && v.length > 0)
     .map((m) => {
-      if (m.startsWith("+")) {
-        modelTable[m.slice(1)] = true;
-      } else if (m.startsWith("-")) {
-        modelTable[m.slice(1)] = false;
-      } else modelTable[m] = true;
+      const available = !m.startsWith("-");
+      const nameConfig =
+        m.startsWith("+") || m.startsWith("-") ? m.slice(1) : m;
+      const [name, displayName] = nameConfig.split(":");
+      modelTable[name] = {
+        name,
+        displayName: displayName || name,
+        available,
+      };
     });
   return modelTable;
 }
@@ -31,10 +44,7 @@ export function collectModels(
   customModels: string,
 ) {
   const modelTable = collectModelTable(models, customModels);
-  const allModels = Object.keys(modelTable).map((m) => ({
-    name: m,
-    available: modelTable[m],
-  }));
+  const allModels = Object.values(modelTable);
 
   return allModels;
 }