diff --git a/api/core/entities/provider_configuration.py b/api/core/entities/provider_configuration.py index 0c26a03d07b..95b8b686e05 100644 --- a/api/core/entities/provider_configuration.py +++ b/api/core/entities/provider_configuration.py @@ -1233,16 +1233,24 @@ class ProviderConfiguration(BaseModel): available_credentials_count = session.execute(count_stmt).scalar() or 0 session.delete(credential_record) + model_credentials_cache_identity_id: str | None = None + if provider_model_record and ( + available_credentials_count <= 1 or provider_model_record.credential_id == credential_id + ): + model_credentials_cache_identity_id = provider_model_record.id + if provider_model_record and available_credentials_count <= 1: # If all credentials are deleted, delete the custom model record session.delete(provider_model_record) elif provider_model_record and provider_model_record.credential_id == credential_id: provider_model_record.credential_id = None provider_model_record.updated_at = naive_utc_now() + + if model_credentials_cache_identity_id: provider_model_credentials_cache = ProviderCredentialsCache( tenant_id=self.tenant_id, - identity_id=provider_model_record.id, - cache_type=ProviderCredentialsCacheType.PROVIDER, + identity_id=model_credentials_cache_identity_id, + cache_type=ProviderCredentialsCacheType.MODEL, ) provider_model_credentials_cache.delete() 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 9eb32206039..b5a48918079 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 @@ -25,6 +25,7 @@ from core.entities.provider_entities import ( SystemConfiguration, SystemConfigurationStatus, ) +from core.helper.model_provider_cache import ProviderCredentialsCacheType from graphon.model_runtime.entities.common_entities import I18nObject from graphon.model_runtime.entities.model_entities import AIModelEntity, FetchFrom, ModelType from graphon.model_runtime.entities.provider_entities import ( @@ -1336,6 +1337,8 @@ def test_create_update_delete_custom_model_credential_flow() -> None: configuration.delete_custom_model_credential(ModelType.LLM, "gpt-4o", "cred-1") assert provider_model_record.credential_id is None assert mock_cache.return_value.delete.call_count == 2 + assert mock_cache.call_args_list[0].kwargs["cache_type"] == ProviderCredentialsCacheType.LOAD_BALANCING_MODEL + assert mock_cache.call_args_list[1].kwargs["cache_type"] == ProviderCredentialsCacheType.MODEL session = Mock() mismatched_credential_record = SimpleNamespace( @@ -2032,9 +2035,16 @@ def test_delete_custom_model_credential_removes_custom_model_record_when_last_cr with _patched_session(session): with patch.object(ProviderConfiguration, "_get_custom_model_record", return_value=provider_model_record): - configuration.delete_custom_model_credential(ModelType.LLM, "gpt-4o", "cred-1") + with patch("core.entities.provider_configuration.ProviderCredentialsCache") as mock_cache: + configuration.delete_custom_model_credential(ModelType.LLM, "gpt-4o", "cred-1") assert any(call.args and call.args[0] is provider_model_record for call in session.delete.call_args_list) + mock_cache.assert_called_once_with( + tenant_id="tenant-1", + identity_id="model-1", + cache_type=ProviderCredentialsCacheType.MODEL, + ) + mock_cache.return_value.delete.assert_called_once() def test_delete_custom_model_credential_rolls_back_on_error() -> None: