dify/web/context/provider-context.tsx

'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
  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)

  const handleOperateUtm = () => {
    let utm
    try {
      utm = JSON.parse(localStorage?.getItem('utm') || '{}')
    }
    catch (e) {
      utm = {
        utm_source: '',
        utm_medium: '',
        utm_campaign: '',
        utm_content: '',
        utm_term: '',
      }
    }
    operationUtm({ url: '/operation/utm', body: utm })
  }
  useEffect(() => {
    (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)
      }
    })()
  }, [])

  return (
    <ProviderContext.Provider value={{
      modelProviders: providersData?.data || [],
      textGenerationModelList: textGenerationModelList?.data || [],
      agentThoughtModelList,
      hasSettedApiKey: !!textGenerationModelList?.data.some(model => model.status === ModelStatusEnum.active),
      supportRetrievalMethods: supportRetrievalMethods?.retrieval_method || [],
      plan,
      isFetchedPlan,
      enableBilling,
      enableReplaceWebAppLogo,
    }}>
      {children}
    </ProviderContext.Provider>
  )
}

export default ProviderContext