feat(credits): Allow to configure model-credit mapping (#13274)

Signed-off-by: -LAN- <laipz8200@outlook.com>

api/configs/feature/hosted_service/__init__.py

@@ -1,9 +1,40 @@
 from typing import Optional
 
-from pydantic import Field, NonNegativeInt
+from pydantic import Field, NonNegativeInt, computed_field
 from pydantic_settings import BaseSettings
 
 
+class HostedCreditConfig(BaseSettings):
+    HOSTED_MODEL_CREDIT_CONFIG: str = Field(
+        description="Model credit configuration in format 'model:credits,model:credits', e.g., 'gpt-4:20,gpt-4o:10'",
+        default="",
+    )
+
+    def get_model_credits(self, model_name: str) -> int:
+        """
+        Get credit value for a specific model name.
+        Returns 1 if model is not found in configuration (default credit).
+
+        :param model_name: The name of the model to search for
+        :return: The credit value for the model
+        """
+        if not self.HOSTED_MODEL_CREDIT_CONFIG:
+            return 1
+
+        try:
+            credit_map = dict(
+                item.strip().split(":", 1) for item in self.HOSTED_MODEL_CREDIT_CONFIG.split(",") if ":" in item
+            )
+
+            # Search for matching model pattern
+            for pattern, credit in credit_map.items():
+                if pattern.strip() in model_name:
+                    return int(credit)
+            return 1  # Default quota if no match found
+        except (ValueError, AttributeError):
+            return 1  # Return default quota if parsing fails
+
+
 class HostedOpenAiConfig(BaseSettings):
     """
     Configuration for hosted OpenAI service
@@ -202,5 +233,7 @@ class HostedServiceConfig(
     HostedZhipuAIConfig,
     # moderation
     HostedModerationConfig,
+    # credit config
+    HostedCreditConfig,
 ):
     pass

api/core/workflow/nodes/llm/node.py

@@ -3,6 +3,7 @@ import logging
 from collections.abc import Generator, Mapping, Sequence
 from typing import TYPE_CHECKING, Any, Optional, cast
 
+from configs import dify_config
 from core.app.entities.app_invoke_entities import ModelConfigWithCredentialsEntity
 from core.entities.model_entities import ModelStatus
 from core.entities.provider_entities import QuotaUnit
@@ -732,10 +733,7 @@ class LLMNode(BaseNode[LLMNodeData]):
             if quota_unit == QuotaUnit.TOKENS:
                 used_quota = usage.total_tokens
             elif quota_unit == QuotaUnit.CREDITS:
-                used_quota = 1
-
-                if "gpt-4" in model_instance.model:
-                    used_quota = 20
+                used_quota = dify_config.get_model_credits(model_instance.model)
             else:
                 used_quota = 1
 

api/events/event_handlers/deduct_quota_when_message_created.py

@@ -1,3 +1,4 @@
+from configs import dify_config
 from core.app.entities.app_invoke_entities import AgentChatAppGenerateEntity, ChatAppGenerateEntity
 from core.entities.provider_entities import QuotaUnit
 from events.message_event import message_was_created
@@ -37,10 +38,7 @@ def handle(sender, **kwargs):
         if quota_unit == QuotaUnit.TOKENS:
             used_quota = message.message_tokens + message.answer_tokens
         elif quota_unit == QuotaUnit.CREDITS:
-            used_quota = 1
-
-            if "gpt-4" in model_config.model:
-                used_quota = 20
+            used_quota = dify_config.get_model_credits(model_config.model)
         else:
             used_quota = 1