import { memo, useCallback, useEffect, useMemo, useState } from 'react' import { useTranslation } from 'react-i18next' import useSWR from 'swr' import type { ModelItem, ModelLoadBalancingConfig, ModelLoadBalancingConfigEntry, ModelProvider } from '../declarations' import { FormTypeEnum } from '../declarations' import ModelIcon from '../model-icon' import ModelName from '../model-name' import { savePredefinedLoadBalancingConfig } from '../utils' 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 { fetchModelLoadBalancingConfig } from '@/service/common' import Loading from '@/app/components/base/loading' import { useToastContext } from '@/app/components/base/toast' export type ModelLoadBalancingModalProps = { provider: ModelProvider model: ModelItem open?: boolean onClose?: () => void onSave?: (provider: string) => void } // model balancing config modal const ModelLoadBalancingModal = ({ provider, model, open = false, onClose, onSave }: ModelLoadBalancingModalProps) => { const { t } = useTranslation() const { notify } = useToastContext() const [loading, setLoading] = useState(false) const { data, mutate } = useSWR( `/workspaces/current/model-providers/${provider.provider}/models/credentials?model=${model.model}&model_type=${model.model_type}`, fetchModelLoadBalancingConfig, ) const originalConfig = data?.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( () => provider.provider_credential_schema.credential_form_schemas.filter( ({ type }) => type === FormTypeEnum.secretInput, ), [provider.provider_credential_schema.credential_form_schemas], ) 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 handleSave = async () => { try { setLoading(true) const res = await savePredefinedLoadBalancingConfig( provider.provider, ({ ...(data?.credentials ?? {}), __model_type: model.model_type, __model_name: model.model, }), { ...draftConfig, enabled: Boolean(draftConfig?.enabled), configs: draftConfig!.configs.map(encodeConfigEntrySecretValues), }, ) if (res.result === 'success') { notify({ type: 'success', message: t('common.actionMsg.modifiedSuccessfully') }) mutate() onSave?.(provider.provider) onClose?.() } } finally { setLoading(false) } } return (
{t('common.modelProvider.configLoadBalancing')}
{Boolean(model) && (
)} } > {!draftConfig ? : ( <>
toggleModalBalancing(false) : undefined} >
{Boolean(model) && ( )}
{t('common.modelProvider.providerManaged')}
{t('common.modelProvider.providerManagedDescription')}
) }
) } export default memo(ModelLoadBalancingModal)