fix(provider_manager): fix custom provider

Signed-off-by: -LAN- <[email protected]>
langgenius · Feb 25, 2025 · 2355509 · 2355509
1 parent 76bcdc2
commit 2355509
Showing 1 changed file with 8 additions and 7 deletions.
diff --git a/api/core/provider_manager.py b/api/core/provider_manager.py
@@ -370,7 +370,11 @@ 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)
+            if provider.provider_name.startswith(DEFAULT_PLUGIN_ID):
+                provider_name = provider.provider_name
+            else:
+                provider_name = f"{DEFAULT_PLUGIN_ID}/{provider.provider_name}/{provider.provider_name}"
+            provider_name_to_provider_records_dict[provider_name].append(provider)
 
         return provider_name_to_provider_records_dict
 
@@ -505,14 +509,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 +529,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()