fix(provider_manager): fix custom provider (#14340)

Signed-off-by: -LAN- <[email protected]>
langgenius · Feb 25, 2025 · 57b60dd · 57b60dd
1 parent ff911d0
commit 57b60dd
Showing 1 changed file with 5 additions and 8 deletions.
diff --git a/api/core/provider_manager.py b/api/core/provider_manager.py
@@ -6,7 +6,6 @@
 from sqlalchemy.exc import IntegrityError
 
 from configs import dify_config
-from core.entities import DEFAULT_PLUGIN_ID
 from core.entities.model_entities import DefaultModelEntity, DefaultModelProviderEntity
 from core.entities.provider_configuration import ProviderConfiguration, ProviderConfigurations, ProviderModelBundle
 from core.entities.provider_entities import (
@@ -370,7 +369,8 @@ def _get_all_providers(tenant_id: str) -> dict[str, list[Provider]]:
 
         provider_name_to_provider_records_dict = defaultdict(list)
         for provider in providers:
-            provider_name_to_provider_records_dict[provider.provider_name].append(provider)
+            # TODO: Use provider name with prefix after the data migration
+            provider_name_to_provider_records_dict[str(ModelProviderID(provider.provider_name))].append(provider)
 
         return provider_name_to_provider_records_dict
 
@@ -505,14 +505,12 @@ def _init_trial_provider_records(
                 if quota.quota_type == ProviderQuotaType.TRIAL:
                     # Init trial provider records if not exists
                     if ProviderQuotaType.TRIAL not in provider_quota_to_provider_record_dict:
-                        if not provider_name.startswith(DEFAULT_PLUGIN_ID):
-                            continue
-                        hosting_provider_name = provider_name.split("/")[-1]
                         try:
                             # FIXME ignore the type errork, onyl TrialHostingQuota has limit need to change the logic
                             provider_record = Provider(
                                 tenant_id=tenant_id,
-                                provider_name=hosting_provider_name,
+                                # TODO: Use provider name with prefix after the data migration.
+                                provider_name=ModelProviderID(provider_name).provider_name,
                                 provider_type=ProviderType.SYSTEM.value,
                                 quota_type=ProviderQuotaType.TRIAL.value,
                                 quota_limit=quota.quota_limit,  # type: ignore
@@ -527,13 +525,12 @@ def _init_trial_provider_records(
                                 db.session.query(Provider)
                                 .filter(
                                     Provider.tenant_id == tenant_id,
-                                    Provider.provider_name == hosting_provider_name,
+                                    Provider.provider_name == ModelProviderID(provider_name).provider_name,
                                     Provider.provider_type == ProviderType.SYSTEM.value,
                                     Provider.quota_type == ProviderQuotaType.TRIAL.value,
                                 )
                                 .first()
                             )
-
                             if provider_record and not provider_record.is_valid:
                                 provider_record.is_valid = True
                                 db.session.commit()