import { memo, useCallback, useEffect, useMemo, useState } from 'react' import { useTranslation } from 'react-i18next' import type { Credential, ModelItem, ModelLoadBalancingConfig, ModelLoadBalancingConfigEntry, ModelProvider, } from '../declarations' import { ConfigurationMethodEnum, FormTypeEnum, } from '../declarations' import ModelIcon from '../model-icon' import ModelName from '../model-name' import ModelLoadBalancingConfigs from './model-load-balancing-configs' import classNames from '@/utils/classnames' import Modal from '@/app/components/base/modal' import Button from '@/app/components/base/button' import Loading from '@/app/components/base/loading' import { useToastContext } from '@/app/components/base/toast' import { SwitchCredentialInLoadBalancing } from '@/app/components/header/account-setting/model-provider-page/model-auth' import { useGetModelCredential, useUpdateModelLoadBalancingConfig, } from '@/service/use-models' export type ModelLoadBalancingModalProps = { provider: ModelProvider configurateMethod: ConfigurationMethodEnum model: ModelItem credential?: Credential open?: boolean onClose?: () => void onSave?: (provider: string) => void } // model balancing config modal const ModelLoadBalancingModal = ({ provider, configurateMethod, model, credential, open = false, onClose, onSave, }: ModelLoadBalancingModalProps) => { const { t } = useTranslation() const { notify } = useToastContext() const [loading, setLoading] = useState(false) const providerFormSchemaPredefined = configurateMethod === ConfigurationMethodEnum.predefinedModel const configFrom = providerFormSchemaPredefined ? 'predefined-model' : 'custom-model' const { isLoading, data, refetch, } = useGetModelCredential(true, provider.provider, credential?.credential_id, model.model, model.model_type, configFrom) const modelCredential = data const { load_balancing, current_credential_id, available_credentials, current_credential_name, } = modelCredential ?? {} const originalConfig = load_balancing const [draftConfig, setDraftConfig] = useState() const originalConfigMap = useMemo(() => { if (!originalConfig) return {} return originalConfig?.configs.reduce((prev, config) => { if (config.id) prev[config.id] = config return prev }, {} as Record) }, [originalConfig]) useEffect(() => { if (originalConfig) setDraftConfig(originalConfig) }, [originalConfig]) const toggleModalBalancing = useCallback((enabled: boolean) => { if (draftConfig) { setDraftConfig({ ...draftConfig, enabled, }) } }, [draftConfig]) const extendedSecretFormSchemas = useMemo( () => { if (providerFormSchemaPredefined) { return provider?.provider_credential_schema?.credential_form_schemas?.filter( ({ type }) => type === FormTypeEnum.secretInput, ) ?? [] } return provider?.model_credential_schema?.credential_form_schemas?.filter( ({ type }) => type === FormTypeEnum.secretInput, ) ?? [] }, [provider?.model_credential_schema?.credential_form_schemas, provider?.provider_credential_schema?.credential_form_schemas, providerFormSchemaPredefined], ) const encodeConfigEntrySecretValues = useCallback((entry: ModelLoadBalancingConfigEntry) => { const result = { ...entry } extendedSecretFormSchemas.forEach(({ variable }) => { if (entry.id && result.credentials[variable] === originalConfigMap[entry.id]?.credentials?.[variable]) result.credentials[variable] = '[__HIDDEN__]' }) return result }, [extendedSecretFormSchemas, originalConfigMap]) const { mutateAsync: updateModelLoadBalancingConfig } = useUpdateModelLoadBalancingConfig(provider.provider) const initialCustomModelCredential = useMemo(() => { if (!current_credential_id) return undefined return { credential_id: current_credential_id, credential_name: current_credential_name, } }, [current_credential_id, current_credential_name]) const [customModelCredential, setCustomModelCredential] = useState(initialCustomModelCredential) const handleSave = async () => { try { setLoading(true) const res = await updateModelLoadBalancingConfig( { credential_id: customModelCredential?.credential_id || current_credential_id, config_from: configFrom, model: model.model, model_type: model.model_type, load_balancing: { ...draftConfig, configs: draftConfig!.configs.map(encodeConfigEntrySecretValues), enabled: Boolean(draftConfig?.enabled), }, }, ) if (res.result === 'success') { notify({ type: 'success', message: t('common.actionMsg.modifiedSuccessfully') }) onSave?.(provider.provider) onClose?.() } } finally { setLoading(false) } } return (
{ draftConfig?.enabled ? t('common.modelProvider.auth.configLoadBalancing') : t('common.modelProvider.auth.configModel') }
{Boolean(model) && (
)} } > {!draftConfig ? : ( <>
toggleModalBalancing(false) : undefined} >
{Boolean(model) && ( )}
{ providerFormSchemaPredefined ? t('common.modelProvider.auth.providerManaged') : t('common.modelProvider.auth.specifyModelCredential') }
{ providerFormSchemaPredefined ? t('common.modelProvider.auth.providerManagedTip') : t('common.modelProvider.auth.specifyModelCredentialTip') }
{ !providerFormSchemaPredefined && ( ) }
{ modelCredential && ( ) }
) }
) } export default memo(ModelLoadBalancingModal)