'use client' import { createContext, useContext } from 'use-context-selector' import useSWR from 'swr' import { useEffect, useMemo, useState } from 'react' import { fetchModelList, fetchModelProviders, fetchSupportRetrievalMethods, operationUtm, } 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 onPlanInfoChanged: () => void 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, onPlanInfoChanged: () => { }, 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) const handleOperateUtm = () => { let utm try { utm = JSON.parse(localStorage?.getItem('utm') || '{}') } catch (e) { utm = { utm_source: '', utm_medium: '', utm_campaign: '', utm_content: '', utm_term: '', } } if (utm.utm_source || utm.utm_medium || utm.utm_campaign || utm.utm_content || utm.utm_term) operationUtm({ url: '/operation/utm', body: utm }) } const fetchPlan = async () => { const data = await fetchCurrentPlanInfo() const enabled = data.billing.enabled setEnableBilling(enabled) setEnableReplaceWebAppLogo(data.can_replace_logo) if (enabled) { setPlan(parseCurrentPlan(data)) handleOperateUtm() setIsFetchedPlan(true) } } useEffect(() => { fetchPlan() }, []) return ( model.status === ModelStatusEnum.active), supportRetrievalMethods: supportRetrievalMethods?.retrieval_method || [], plan, isFetchedPlan, enableBilling, onPlanInfoChanged: fetchPlan, enableReplaceWebAppLogo, }}> {children} ) } export default ProviderContext