Merge branch 'feat/support-use-quota-in-deepseek' into deploy/dev

langgenius · Feb 5, 2025 · f21cc9b · f21cc9b
2 parents 9f57cd0 + 2bc5e15
commit f21cc9b
Show file tree

Hide file tree

Showing 4 changed files with 19 additions and 17 deletions.
diff --git 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
@@ -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
@@ -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
@@ -1,12 +1,12 @@
-from enum import Enum
+from enum import Enum, StrEnum
 
 from sqlalchemy import func
 
 from .engine import db
 from .types import StringUUID
 
 
-class ProviderType(Enum):
+class ProviderType(StrEnum):
     CUSTOM = "custom"
     SYSTEM = "system"