'use client' import { createContext, useContext } from 'use-context-selector' import useSWR from 'swr' import { useEffect, useMemo, useState } from 'react' import { fetchModelList, fetchModelProviders, fetchSupportRetrievalMethods, } from '@/service/common' import { ModelFeatureEnum, 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' const ProviderContext = createContext<{ modelProviders: ModelProvider[] textGenerationModelList: Model[] agentThoughtModelList: Model[] supportRetrievalMethods: RETRIEVE_METHOD[] hasSettedApiKey: boolean plan: { type: Plan usage: UsagePlanInfo total: UsagePlanInfo } isFetchedPlan: boolean enableBilling: boolean enableReplaceWebAppLogo: boolean }>({ modelProviders: [], textGenerationModelList: [], agentThoughtModelList: [], 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, enableReplaceWebAppLogo: false, }) export const useProviderContext = () => useContext(ProviderContext) 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 agentThoughtModelList = useMemo(() => { const result: Model[] = [] if (textGenerationModelList?.data) { textGenerationModelList?.data.forEach((item) => { const agentThoughtModels = item.models.filter(model => model.features?.includes(ModelFeatureEnum.agentThought)) if (agentThoughtModels.length) { result.push({ ...item, models: agentThoughtModels, }) } }) return result } return [] }, [textGenerationModelList]) const [plan, setPlan] = useState(defaultPlan) const [isFetchedPlan, setIsFetchedPlan] = useState(false) const [enableBilling, setEnableBilling] = useState(true) const [enableReplaceWebAppLogo, setEnableReplaceWebAppLogo] = useState(false) useEffect(() => { (async () => { const data = await fetchCurrentPlanInfo() const enabled = data.billing.enabled setEnableBilling(enabled) setEnableReplaceWebAppLogo(data.can_replace_logo) if (enabled) { setPlan(parseCurrentPlan(data)) // setPlan(parseCurrentPlan({ // ...data, // annotation_quota_limit: { // ...data.annotation_quota_limit, // limit: 10, // }, // })) setIsFetchedPlan(true) } })() }, []) return ( model.status === ModelStatusEnum.active), supportRetrievalMethods: supportRetrievalMethods?.retrieval_method || [], plan, isFetchedPlan, enableBilling, enableReplaceWebAppLogo, }}> {children} ) } export default ProviderContext