diff --git a/api/core/entities/provider_configuration.py b/api/core/entities/provider_configuration.py index 25774a7054b..0c26a03d07b 100644 --- a/api/core/entities/provider_configuration.py +++ b/api/core/entities/provider_configuration.py @@ -1188,7 +1188,17 @@ class ProviderConfiguration(BaseModel): ) credential_record = session.execute(stmt).scalar_one_or_none() if not credential_record: - raise ValueError("Credential record not found.") + fallback_stmt = select(ProviderModelCredential).where( + ProviderModelCredential.id == credential_id, + ProviderModelCredential.tenant_id == self.tenant_id, + ProviderModelCredential.provider_name.in_(self._get_provider_names()), + ) + credential_record = session.execute(fallback_stmt).scalar_one_or_none() + if not credential_record: + raise ValueError("Credential record not found.") + + model = credential_record.model_name + model_type = ModelType(credential_record.model_type) lb_stmt = select(LoadBalancingModelConfig).where( LoadBalancingModelConfig.tenant_id == self.tenant_id, diff --git a/api/tests/unit_tests/core/entities/test_entities_provider_configuration.py b/api/tests/unit_tests/core/entities/test_entities_provider_configuration.py index 23d8fd28df0..9eb32206039 100644 --- a/api/tests/unit_tests/core/entities/test_entities_provider_configuration.py +++ b/api/tests/unit_tests/core/entities/test_entities_provider_configuration.py @@ -1337,6 +1337,29 @@ def test_create_update_delete_custom_model_credential_flow() -> None: assert provider_model_record.credential_id is None assert mock_cache.return_value.delete.call_count == 2 + session = Mock() + mismatched_credential_record = SimpleNamespace( + id="cred-2", + model_name="stored-model", + model_type=ModelType.TEXT_EMBEDDING, + ) + provider_model_record = SimpleNamespace(id="model-2", credential_id="cred-2", updated_at=None) + session.execute.side_effect = [ + _exec_result(scalar_one_or_none=None), + _exec_result(scalar_one_or_none=mismatched_credential_record), + _exec_result(scalars_all=[]), + _exec_result(scalar=1), + ] + with _patched_session(session): + with patch.object( + ProviderConfiguration, + "_get_custom_model_record", + return_value=provider_model_record, + ) as mock_get_model: + configuration.delete_custom_model_credential(ModelType.LLM, "request-model", "cred-2") + mock_get_model.assert_called_once_with(ModelType.TEXT_EMBEDDING, "stored-model", session=session) + session.delete.assert_any_call(mismatched_credential_record) + def test_add_model_credential_to_model_and_switch_custom_model_credential() -> None: configuration = _build_provider_configuration() diff --git a/web/app/components/header/account-setting/model-provider-page/declarations.ts b/web/app/components/header/account-setting/model-provider-page/declarations.ts index 3ad756a4e07..36a05ba8a91 100644 --- a/web/app/components/header/account-setting/model-provider-page/declarations.ts +++ b/web/app/components/header/account-setting/model-provider-page/declarations.ts @@ -306,6 +306,14 @@ export type ModelCredential = { current_credential_name?: string } +export type ModelCredentialPayload = { + credentials: Record + model: string + model_type: ModelTypeEnum + name?: string + credential_id?: string +} + export enum ModelModalModeEnum { configProviderCredential = 'config-provider-credential', configCustomModel = 'config-custom-model', diff --git a/web/app/components/header/account-setting/model-provider-page/model-auth/hooks/__tests__/use-auth.spec.tsx b/web/app/components/header/account-setting/model-provider-page/model-auth/hooks/__tests__/use-auth.spec.tsx index bb682cf8a4d..2f87bf9a629 100644 --- a/web/app/components/header/account-setting/model-provider-page/model-auth/hooks/__tests__/use-auth.spec.tsx +++ b/web/app/components/header/account-setting/model-provider-page/model-auth/hooks/__tests__/use-auth.spec.tsx @@ -168,8 +168,6 @@ describe('useAuth', () => { expect(mockDeleteProviderCredential).toHaveBeenCalledWith({ credential_id: 'cred-1', - model: 'gpt-4', - model_type: ModelTypeEnum.textGeneration, }) expect(mockDeleteModelService).not.toHaveBeenCalled() expect(onRemove).toHaveBeenCalledWith('cred-1') diff --git a/web/app/components/header/account-setting/model-provider-page/model-auth/hooks/use-auth.ts b/web/app/components/header/account-setting/model-provider-page/model-auth/hooks/use-auth.ts index 1cdb28523e3..f02e49463d7 100644 --- a/web/app/components/header/account-setting/model-provider-page/model-auth/hooks/use-auth.ts +++ b/web/app/components/header/account-setting/model-provider-page/model-auth/hooks/use-auth.ts @@ -6,6 +6,16 @@ import { useModelModalHandler, useRefreshModel } from '@/app/components/header/a import { useDeleteModel } from '@/service/use-models' import { useAuthService } from './use-auth-service' +type ProviderCredentialOperationPayload = { + credential_id: string +} + +type ModelCredentialOperationPayload = { + credential_id: string + model: string + model_type: CustomModel['model_type'] +} + export const useAuth = (provider: ModelProvider, configurationMethod: ConfigurationMethodEnum, currentCustomConfigurationModelFixedFields?: CustomConfigurationModelFixedFields, extra: { isModelCredential?: boolean onUpdate?: (newPayload?: any, formValues?: Record) => void @@ -30,6 +40,18 @@ export const useAuth = (provider: ModelProvider, configurationMethod: Configurat setDeleteModel(model) pendingOperationModel.current = model }, []) + const resolveModelContext = useCallback((model?: CustomModel | null): CustomModel | undefined => { + if (model) + return model + + if (!currentCustomConfigurationModelFixedFields) + return undefined + + return { + model: currentCustomConfigurationModelFixedFields.__model_name, + model_type: currentCustomConfigurationModelFixedFields.__model_type, + } + }, [currentCustomConfigurationModelFixedFields]) const openConfirmDelete = useCallback((credential?: Credential, model?: CustomModel) => { if (credential) handleSetDeleteCredentialId(credential.credential_id) @@ -51,18 +73,32 @@ export const useAuth = (provider: ModelProvider, configurationMethod: Configurat return try { handleSetDoingAction(true) - await getActiveCredentialService(!!model)({ - credential_id: credential.credential_id, - model: model?.model, - model_type: model?.model_type, - }) + const modelContext = model ?? (isModelCredential ? resolveModelContext() : undefined) + if (modelContext) { + const activeModelCredential = getActiveCredentialService(true) as ( + payload: ModelCredentialOperationPayload, + ) => Promise + await activeModelCredential({ + credential_id: credential.credential_id, + model: modelContext.model, + model_type: modelContext.model_type, + }) + } + else { + const activeProviderCredential = getActiveCredentialService(false) as ( + payload: ProviderCredentialOperationPayload, + ) => Promise + await activeProviderCredential({ + credential_id: credential.credential_id, + }) + } toast.success(t('api.actionSuccess', { ns: 'common' })) handleRefreshModel(provider, undefined, true) } finally { handleSetDoingAction(false) } - }, [getActiveCredentialService, t, handleSetDoingAction]) + }, [getActiveCredentialService, isModelCredential, resolveModelContext, t, handleSetDoingAction]) const handleConfirmDelete = useCallback(async () => { if (doingActionRef.current) return @@ -74,12 +110,30 @@ export const useAuth = (provider: ModelProvider, configurationMethod: Configurat handleSetDoingAction(true) let payload: any = {} if (pendingOperationCredentialId.current) { - payload = { - credential_id: pendingOperationCredentialId.current, - model: pendingOperationModel.current?.model, - model_type: pendingOperationModel.current?.model_type, + if (isModelCredential) { + const modelContext = resolveModelContext(pendingOperationModel.current) + if (!modelContext) + return + + payload = { + credential_id: pendingOperationCredentialId.current, + model: modelContext.model, + model_type: modelContext.model_type, + } + const deleteModelCredential = getDeleteCredentialService(true) as ( + payload: ModelCredentialOperationPayload, + ) => Promise + await deleteModelCredential(payload) + } + else { + payload = { + credential_id: pendingOperationCredentialId.current, + } + const deleteProviderCredential = getDeleteCredentialService(false) as ( + payload: ProviderCredentialOperationPayload, + ) => Promise + await deleteProviderCredential(payload) } - await getDeleteCredentialService(!!isModelCredential)(payload) } if (!pendingOperationCredentialId.current && pendingOperationModel.current) { payload = { @@ -96,7 +150,7 @@ export const useAuth = (provider: ModelProvider, configurationMethod: Configurat finally { handleSetDoingAction(false) } - }, [t, handleSetDoingAction, getDeleteCredentialService, isModelCredential, closeConfirmDelete, handleRefreshModel, provider, configurationMethod, deleteModelService]) + }, [t, handleSetDoingAction, getDeleteCredentialService, isModelCredential, closeConfirmDelete, handleRefreshModel, provider, configurationMethod, deleteModelService, resolveModelContext]) const handleSaveCredential = useCallback(async (payload: Record) => { if (doingActionRef.current) return diff --git a/web/app/components/header/account-setting/model-provider-page/model-modal/__tests__/index.spec.tsx b/web/app/components/header/account-setting/model-provider-page/model-modal/__tests__/index.spec.tsx index b3dad54e25e..0fb0feaef02 100644 --- a/web/app/components/header/account-setting/model-provider-page/model-modal/__tests__/index.spec.tsx +++ b/web/app/components/header/account-setting/model-provider-page/model-modal/__tests__/index.spec.tsx @@ -353,4 +353,44 @@ describe('ModelModal', () => { expect(mockHandlers.openConfirmDelete).toHaveBeenCalledWith({ credential_id: 'remove-1' }, undefined) removable.unmount() }) + + it('should use fixed model context when saving a model credential without model prop', async () => { + mockState.formSchemas = [{ variable: 'api_key', type: 'secret-input' } as unknown as CredentialFormSchema] + mockFormState.responses = [ + { isCheckValidated: true, values: { __authorization_name__: 'Xinference Auth', api_key: 'secret' } }, + ] + + renderModal({ + mode: ModelModalModeEnum.configModelCredential, + currentCustomConfigurationModelFixedFields: { + __model_name: 'bge-m3', + __model_type: ModelTypeEnum.textEmbedding, + }, + }) + fireEvent.click(screen.getByRole('button', { name: 'common.operation.save' })) + + await waitFor(() => { + expect(mockHandlers.handleSaveCredential).toHaveBeenCalledWith({ + credential_id: undefined, + credentials: { api_key: 'secret' }, + name: 'Xinference Auth', + model: 'bge-m3', + model_type: ModelTypeEnum.textEmbedding, + }) + }) + }) + + it('should not submit model credential payload when model context is missing', async () => { + mockState.formSchemas = [{ variable: 'api_key', type: 'secret-input' } as unknown as CredentialFormSchema] + mockFormState.responses = [ + { isCheckValidated: true, values: { __authorization_name__: 'Missing Model Auth', api_key: 'secret' } }, + ] + + renderModal({ mode: ModelModalModeEnum.configModelCredential }) + fireEvent.click(screen.getByRole('button', { name: 'common.operation.save' })) + + await waitFor(() => { + expect(mockHandlers.handleSaveCredential).not.toHaveBeenCalled() + }) + }) }) diff --git a/web/app/components/header/account-setting/model-provider-page/model-modal/index.tsx b/web/app/components/header/account-setting/model-provider-page/model-modal/index.tsx index 245554c5c1b..17ddb93efe8 100644 --- a/web/app/components/header/account-setting/model-provider-page/model-modal/index.tsx +++ b/web/app/components/header/account-setting/model-provider-page/model-modal/index.tsx @@ -147,17 +147,22 @@ const ModelModal: FC = ({ modelNameAndTypeValues = formResult.values } - if (mode === ModelModalModeEnum.configModelCredential && model) { - modelNameAndTypeValues = { - __model_name: model.model, - __model_type: model.model_type, - } - } + if ( + mode === ModelModalModeEnum.configModelCredential + || (mode === ModelModalModeEnum.addCustomModelToModelList && selectedCredential?.addNewCredential) + ) { + const modelContext = model ?? (currentCustomConfigurationModelFixedFields + ? { + model: currentCustomConfigurationModelFixedFields.__model_name, + model_type: currentCustomConfigurationModelFixedFields.__model_type, + } + : undefined) + if (!modelContext) + return - if (mode === ModelModalModeEnum.addCustomModelToModelList && selectedCredential?.addNewCredential && model) { modelNameAndTypeValues = { - __model_name: model.model, - __model_type: model.model_type, + __model_name: modelContext.model, + __model_type: modelContext.model_type, } } const { @@ -178,7 +183,13 @@ const ModelModal: FC = ({ __authorization_name__, ...rest } = values - if (__model_name && __model_type) { + const shouldSaveModelCredential = mode === ModelModalModeEnum.configCustomModel + || mode === ModelModalModeEnum.configModelCredential + || (mode === ModelModalModeEnum.addCustomModelToModelList && selectedCredential?.addNewCredential) + if (shouldSaveModelCredential) { + if (!__model_name || !__model_type) + return + await handleSaveCredential({ credential_id: credential?.credential_id, credentials: rest, @@ -195,7 +206,7 @@ const ModelModal: FC = ({ }) } onSave(values) - }, [mode, selectedCredential, model, canUseCredential, canCreateCredential, canManageCredential, onSave, handleActiveCredential, onCancel, handleSaveCredential, credential]) + }, [mode, selectedCredential, model, currentCustomConfigurationModelFixedFields, canUseCredential, canCreateCredential, canManageCredential, onSave, handleActiveCredential, onCancel, handleSaveCredential, credential]) const modalTitle = useMemo(() => { let label = t('modelProvider.auth.apiKeyModal.title', { ns: 'common' }) diff --git a/web/service/use-models.ts b/web/service/use-models.ts index 05582b41052..5b984fc32ce 100644 --- a/web/service/use-models.ts +++ b/web/service/use-models.ts @@ -1,5 +1,6 @@ import type { ModelCredential, + ModelCredentialPayload, ModelItem, ModelLoadBalancingConfig, ModelTypeEnum, @@ -91,7 +92,7 @@ export const useGetModelCredential = ( export const useAddModelCredential = (provider: string) => { return useMutation({ - mutationFn: (data: ModelCredential) => post<{ result: string }>(`/workspaces/current/model-providers/${provider}/models/credentials`, { + mutationFn: (data: ModelCredentialPayload) => post<{ result: string }>(`/workspaces/current/model-providers/${provider}/models/credentials`, { body: data, }), }) @@ -99,7 +100,7 @@ export const useAddModelCredential = (provider: string) => { export const useEditModelCredential = (provider: string) => { return useMutation({ - mutationFn: (data: ModelCredential) => put<{ result: string }>(`/workspaces/current/model-providers/${provider}/models/credentials`, { + mutationFn: (data: ModelCredentialPayload) => put<{ result: string }>(`/workspaces/current/model-providers/${provider}/models/credentials`, { body: data, }), }) @@ -109,8 +110,8 @@ export const useDeleteModelCredential = (provider: string) => { return useMutation({ mutationFn: (data: { credential_id: string - model?: string - model_type?: ModelTypeEnum + model: string + model_type: ModelTypeEnum }) => del<{ result: string }>(`/workspaces/current/model-providers/${provider}/models/credentials`, { body: data, }), @@ -132,8 +133,8 @@ export const useActiveModelCredential = (provider: string) => { return useMutation({ mutationFn: (data: { credential_id: string - model?: string - model_type?: ModelTypeEnum + model: string + model_type: ModelTypeEnum }) => post<{ result: string }>(`/workspaces/current/model-providers/${provider}/models/credentials/switch`, { body: data, }),