File size: 4,491 Bytes
a8b3f00 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 |
'use client'
import { createContext, useContext, useContextSelector } from 'use-context-selector'
import useSWR from 'swr'
import { useEffect, useState } from 'react'
import {
fetchModelList,
fetchModelProviders,
fetchSupportRetrievalMethods,
} from '@/service/common'
import {
ModelStatusEnum,
ModelTypeEnum,
} from '@/app/components/header/account-setting/model-provider-page/declarations'
import type { Model, ModelProvider } from '@/app/components/header/account-setting/model-provider-page/declarations'
import type { RETRIEVE_METHOD } from '@/types/app'
import { Plan, type UsagePlanInfo } from '@/app/components/billing/type'
import { fetchCurrentPlanInfo } from '@/service/billing'
import { parseCurrentPlan } from '@/app/components/billing/utils'
import { defaultPlan } from '@/app/components/billing/config'
type ProviderContextState = {
modelProviders: ModelProvider[]
textGenerationModelList: Model[]
supportRetrievalMethods: RETRIEVE_METHOD[]
isAPIKeySet: boolean
plan: {
type: Plan
usage: UsagePlanInfo
total: UsagePlanInfo
}
isFetchedPlan: boolean
enableBilling: boolean
onPlanInfoChanged: () => void
enableReplaceWebAppLogo: boolean
modelLoadBalancingEnabled: boolean
datasetOperatorEnabled: boolean
}
const ProviderContext = createContext<ProviderContextState>({
modelProviders: [],
textGenerationModelList: [],
supportRetrievalMethods: [],
isAPIKeySet: true,
plan: {
type: Plan.sandbox,
usage: {
vectorSpace: 32,
buildApps: 12,
teamMembers: 1,
annotatedResponse: 1,
documentsUploadQuota: 50,
},
total: {
vectorSpace: 200,
buildApps: 50,
teamMembers: 1,
annotatedResponse: 10,
documentsUploadQuota: 500,
},
},
isFetchedPlan: false,
enableBilling: false,
onPlanInfoChanged: () => { },
enableReplaceWebAppLogo: false,
modelLoadBalancingEnabled: false,
datasetOperatorEnabled: 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
}
export const ProviderContextProvider = ({
children,
}: ProviderContextProviderProps) => {
const { data: providersData } = useSWR('/workspaces/current/model-providers', fetchModelProviders)
const fetchModelListUrlPrefix = '/workspaces/current/models/model-types/'
const { data: textGenerationModelList } = useSWR(`${fetchModelListUrlPrefix}${ModelTypeEnum.textGeneration}`, fetchModelList)
const { data: supportRetrievalMethods } = useSWR('/datasets/retrieval-setting', fetchSupportRetrievalMethods)
const [plan, setPlan] = useState(defaultPlan)
const [isFetchedPlan, setIsFetchedPlan] = useState(false)
const [enableBilling, setEnableBilling] = useState(true)
const [enableReplaceWebAppLogo, setEnableReplaceWebAppLogo] = useState(false)
const [modelLoadBalancingEnabled, setModelLoadBalancingEnabled] = useState(false)
const [datasetOperatorEnabled, setDatasetOperatorEnabled] = useState(false)
const fetchPlan = async () => {
const data = await fetchCurrentPlanInfo()
const enabled = data.billing.enabled
setEnableBilling(enabled)
setEnableReplaceWebAppLogo(data.can_replace_logo)
if (enabled) {
setPlan(parseCurrentPlan(data))
setIsFetchedPlan(true)
}
if (data.model_load_balancing_enabled)
setModelLoadBalancingEnabled(true)
if (data.dataset_operator_enabled)
setDatasetOperatorEnabled(true)
}
useEffect(() => {
fetchPlan()
}, [])
return (
<ProviderContext.Provider value={{
modelProviders: providersData?.data || [],
textGenerationModelList: textGenerationModelList?.data || [],
isAPIKeySet: !!textGenerationModelList?.data.some(model => model.status === ModelStatusEnum.active),
supportRetrievalMethods: supportRetrievalMethods?.retrieval_method || [],
plan,
isFetchedPlan,
enableBilling,
onPlanInfoChanged: fetchPlan,
enableReplaceWebAppLogo,
modelLoadBalancingEnabled,
datasetOperatorEnabled,
}}>
{children}
</ProviderContext.Provider>
)
}
export default ProviderContext
|