From 2bc5e15adc7fa22788c67d5d53030338bc4bee0d Mon Sep 17 00:00:00 2001
From: -LAN- <laipz8200@outlook.com>
Date: Wed, 5 Feb 2025 20:05:25 +0800
Subject: [PATCH] feat(quotas): Use TrialProvider and add default quotas config
 for deepseek

Signed-off-by: -LAN- <laipz8200@outlook.com>
---
 .../feature/hosted_service/__init__.py        |  4 ++++
 api/core/hosting_configuration.py             | 24 +++++++++----------
 api/core/provider_manager.py                  |  4 ++--
 api/models/provider.py                        |  4 ++--
 4 files changed, 19 insertions(+), 17 deletions(-)

diff --git a/api/configs/feature/hosted_service/__init__.py b/api/configs/feature/hosted_service/__init__.py
index a20ffa2f387809..4235d7231dd275 100644
--- a/api/configs/feature/hosted_service/__init__.py
+++ b/api/configs/feature/hosted_service/__init__.py
@@ -208,6 +208,10 @@ class HostedDeepseekConfig(BaseSettings):
         description="Comma-separated list of available models for paid access",
         default="",
     )
+    HOSTED_DEEPSEEK_QUOTA_LIMIT: NonNegativeInt = Field(
+        description="Quota limit for hosted DeepSeek service usage",
+        default=100000,
+    )
 
 
 class HostedServiceConfig(
diff --git a/api/core/hosting_configuration.py b/api/core/hosting_configuration.py
index 41dd74280d7061..48e72da8055a3f 100644
--- a/api/core/hosting_configuration.py
+++ b/api/core/hosting_configuration.py
@@ -155,21 +155,19 @@ def init_openai(self) -> HostingProvider:
         )
 
     def init_deepseek(self) -> HostingProvider:
-        quota_unit = QuotaUnit.TOKENS
-        quotas: list[HostingQuota] = []
+        if not dify_config.HOSTED_DEEPSEEK_ENABLED:
+            return HostingProvider(
+                enabled=False,
+                quota_unit=QuotaUnit.TOKENS,
+            )
 
-        if dify_config.HOSTED_DEEPSEEK_ENABLED:
-            paid_models = self.parse_restrict_models_from_env(dify_config.HOSTED_DEEPSEEK_MODELS)
-            paid_quota = PaidHostingQuota(restrict_models=paid_models)
-            quotas.append(paid_quota)
-        if len(quotas) > 0:
-            credentials = {
-                "api_key": dify_config.HOSTED_DEEPSEEK_API_KEY,
-            }
-            return HostingProvider(enabled=True, credentials=credentials, quota_unit=quota_unit, quotas=quotas)
+        trial_quota = TrialHostingQuota(quota_limit=dify_config.HOSTED_DEEPSEEK_QUOTA_LIMIT)
+        quotas: list[HostingQuota] = [trial_quota]
         return HostingProvider(
-            enabled=False,
-            quota_unit=quota_unit,
+            enabled=True,
+            credentials={"api_key": dify_config.HOSTED_DEEPSEEK_API_KEY},
+            quota_unit=QuotaUnit.TOKENS,
+            quotas=quotas,
         )
 
     def init_anthropic(self) -> HostingProvider:
diff --git a/api/core/provider_manager.py b/api/core/provider_manager.py
index 2430d598ffcb7d..73baed5109a795 100644
--- a/api/core/provider_manager.py
+++ b/api/core/provider_manager.py
@@ -484,7 +484,7 @@ def _init_trial_provider_records(
 
             provider_quota_to_provider_record_dict = {}
             for provider_record in provider_records:
-                if provider_record.provider_type != ProviderType.SYSTEM.value:
+                if provider_record.provider_type != ProviderType.SYSTEM:
                     continue
 
                 provider_quota_to_provider_record_dict[ProviderQuotaType.value_of(provider_record.quota_type)] = (
@@ -496,7 +496,7 @@ def _init_trial_provider_records(
                     # Init trial provider records if not exists
                     if ProviderQuotaType.TRIAL not in provider_quota_to_provider_record_dict:
                         try:
-                            # FIXME ignore the type errork, onyl TrialHostingQuota has limit need to change the logic
+                            # FIXME ignore the type error, onyl TrialHostingQuota has limit need to change the logic
                             provider_record = Provider(
                                 tenant_id=tenant_id,
                                 provider_name=provider_name,
diff --git a/api/models/provider.py b/api/models/provider.py
index abe673975c1ccc..f0f320fb4d823c 100644
--- a/api/models/provider.py
+++ b/api/models/provider.py
@@ -1,4 +1,4 @@
-from enum import Enum
+from enum import Enum, StrEnum
 
 from sqlalchemy import func
 
@@ -6,7 +6,7 @@
 from .types import StringUUID
 
 
-class ProviderType(Enum):
+class ProviderType(StrEnum):
     CUSTOM = "custom"
     SYSTEM = "system"