feat: model load balancing (#4926)

web/context/provider-context.tsx

@@ -1,6 +1,6 @@
 'use client'
 
-import { createContext, useContext } from 'use-context-selector'
+import { createContext, useContext, useContextSelector } from 'use-context-selector'
 import useSWR from 'swr'
 import { useEffect, useState } from 'react'
 import {
@@ -19,11 +19,11 @@ import { fetchCurrentPlanInfo } from '@/service/billing'
 import { parseCurrentPlan } from '@/app/components/billing/utils'
 import { defaultPlan } from '@/app/components/billing/config'
 
-const ProviderContext = createContext<{
+type ProviderContextState = {
   modelProviders: ModelProvider[]
   textGenerationModelList: Model[]
   supportRetrievalMethods: RETRIEVE_METHOD[]
-  hasSettedApiKey: boolean
+  isAPIKeySet: boolean
   plan: {
     type: Plan
     usage: UsagePlanInfo
@@ -33,34 +33,43 @@ const ProviderContext = createContext<{
   enableBilling: boolean
   onPlanInfoChanged: () => void
   enableReplaceWebAppLogo: boolean
-}>({
-      modelProviders: [],
-      textGenerationModelList: [],
-      supportRetrievalMethods: [],
-      hasSettedApiKey: true,
-      plan: {
-        type: Plan.sandbox,
-        usage: {
-          vectorSpace: 32,
-          buildApps: 12,
-          teamMembers: 1,
-          annotatedResponse: 1,
-        },
-        total: {
-          vectorSpace: 200,
-          buildApps: 50,
-          teamMembers: 1,
-          annotatedResponse: 10,
-        },
-      },
-      isFetchedPlan: false,
-      enableBilling: false,
-      onPlanInfoChanged: () => { },
-      enableReplaceWebAppLogo: false,
-    })
+  modelLoadBalancingEnabled: boolean
+}
+const ProviderContext = createContext<ProviderContextState>({
+  modelProviders: [],
+  textGenerationModelList: [],
+  supportRetrievalMethods: [],
+  isAPIKeySet: true,
+  plan: {
+    type: Plan.sandbox,
+    usage: {
+      vectorSpace: 32,
+      buildApps: 12,
+      teamMembers: 1,
+      annotatedResponse: 1,
+    },
+    total: {
+      vectorSpace: 200,
+      buildApps: 50,
+      teamMembers: 1,
+      annotatedResponse: 10,
+    },
+  },
+  isFetchedPlan: false,
+  enableBilling: false,
+  onPlanInfoChanged: () => { },
+  enableReplaceWebAppLogo: false,
+  modelLoadBalancingEnabled: false,
+})
 
 export const useProviderContext = () => useContext(ProviderContext)
 
+// Adding a dangling comma to avoid the generic parsing issue in tsx, see:
+// https://github.com/microsoft/TypeScript/issues/15713
+// eslint-disable-next-line @typescript-eslint/comma-dangle
+export const useProviderContextSelector = <T,>(selector: (state: ProviderContextState) => T): T =>
+  useContextSelector(ProviderContext, selector)
+
 type ProviderContextProviderProps = {
   children: React.ReactNode
 }
@@ -76,6 +85,7 @@ export const ProviderContextProvider = ({
   const [isFetchedPlan, setIsFetchedPlan] = useState(false)
   const [enableBilling, setEnableBilling] = useState(true)
   const [enableReplaceWebAppLogo, setEnableReplaceWebAppLogo] = useState(false)
+  const [modelLoadBalancingEnabled, setModelLoadBalancingEnabled] = useState(false)
 
   const fetchPlan = async () => {
     const data = await fetchCurrentPlanInfo()
@@ -86,6 +96,8 @@ export const ProviderContextProvider = ({
       setPlan(parseCurrentPlan(data))
       setIsFetchedPlan(true)
     }
+    if (data.model_load_balancing_enabled)
+      setModelLoadBalancingEnabled(true)
   }
   useEffect(() => {
     fetchPlan()
@@ -95,13 +107,14 @@ export const ProviderContextProvider = ({
     <ProviderContext.Provider value={{
       modelProviders: providersData?.data || [],
       textGenerationModelList: textGenerationModelList?.data || [],
-      hasSettedApiKey: !!textGenerationModelList?.data.some(model => model.status === ModelStatusEnum.active),
+      isAPIKeySet: !!textGenerationModelList?.data.some(model => model.status === ModelStatusEnum.active),
       supportRetrievalMethods: supportRetrievalMethods?.retrieval_method || [],
       plan,
       isFetchedPlan,
       enableBilling,
       onPlanInfoChanged: fetchPlan,
       enableReplaceWebAppLogo,
+      modelLoadBalancingEnabled,
     }}>
       {children}
     </ProviderContext.Provider>