lloydzhou 1 year ago
parent
commit
aa08183439
4 changed files with 31 additions and 7 deletions
  1. 6 1
      app/api/anthropic/[...path]/route.ts
  2. 16 2
      app/api/common.ts
  3. 2 2
      app/store/config.ts
  4. 7 2
      app/utils/model.ts

+ 6 - 1
app/api/anthropic/[...path]/route.ts

@@ -4,6 +4,7 @@ import {
   Anthropic,
   ApiPath,
   DEFAULT_MODELS,
+  ServiceProvider,
   ModelProvider,
 } from "@/app/constant";
 import { prettyObject } from "@/app/utils/format";
@@ -143,7 +144,11 @@ async function request(req: NextRequest) {
 
       // not undefined and is false
       if (
-        isModelAvailableInServer(serverConfig.customModels, jsonBody?.model)
+        isModelAvailableInServer(
+          serverConfig.customModels,
+          jsonBody?.model,
+          ServiceProvider.Anthropic,
+        )
       ) {
         return NextResponse.json(
           {

+ 16 - 2
app/api/common.ts

@@ -1,6 +1,11 @@
 import { NextRequest, NextResponse } from "next/server";
 import { getServerSideConfig } from "../config/server";
-import { DEFAULT_MODELS, OPENAI_BASE_URL, GEMINI_BASE_URL } from "../constant";
+import {
+  DEFAULT_MODELS,
+  OPENAI_BASE_URL,
+  GEMINI_BASE_URL,
+  ServiceProvider,
+} from "../constant";
 import { isModelAvailableInServer } from "../utils/model";
 import { makeAzurePath } from "../azure";
 
@@ -90,7 +95,16 @@ export async function requestOpenai(req: NextRequest) {
 
       // not undefined and is false
       if (
-        isModelAvailableInServer(serverConfig.customModels, jsonBody?.model)
+        isModelAvailableInServer(
+          serverConfig.customModels,
+          jsonBody?.model,
+          ServiceProvider.OpenAI,
+        ) ||
+        isModelAvailableInServer(
+          serverConfig.customModels,
+          jsonBody?.model,
+          ServiceProvider.Azure,
+        )
       ) {
         return NextResponse.json(
           {

+ 2 - 2
app/store/config.ts

@@ -116,12 +116,12 @@ export const useAppConfig = createPersistStore(
 
       for (const model of oldModels) {
         model.available = false;
-        modelMap[`${model.name}@${model.provider.name}`] = model;
+        modelMap[`${model.name}@${model?.provider?.name}`] = model;
       }
 
       for (const model of newModels) {
         model.available = true;
-        modelMap[`${model.name}@${model.provider.name}`] = model;
+        modelMap[`${model.name}@${model?.provider?.name}`] = model;
       }
 
       set(() => ({

+ 7 - 2
app/utils/model.ts

@@ -118,7 +118,12 @@ export function collectModelsWithDefaultModel(
   return allModels;
 }
 
-export function isModelAvailableInServer(customModels, modelName) {
+export function isModelAvailableInServer(
+  customModels,
+  modelName,
+  providerName,
+) {
+  const fullName = `${modelName}@${providerName}`;
   const modelTable = collectModelTable(DEFAULT_MODELS, customModels);
-  return modelTable[modelName ?? ""].available === false;
+  return modelTable[fullName]?.available === false;
 }