dify

dify/api/core/entities/model_entities.py

@@ -0,0 +1,109 @@
+from collections.abc import Sequence
+from enum import StrEnum, auto
+
+from pydantic import BaseModel, ConfigDict
+
+from core.model_runtime.entities.common_entities import I18nObject
+from core.model_runtime.entities.model_entities import ModelType, ProviderModel
+from core.model_runtime.entities.provider_entities import ProviderEntity
+
+
+class ModelStatus(StrEnum):
+    """
+    Enum class for model status.
+    """
+
+    ACTIVE = auto()
+    NO_CONFIGURE = "no-configure"
+    QUOTA_EXCEEDED = "quota-exceeded"
+    NO_PERMISSION = "no-permission"
+    DISABLED = auto()
+    CREDENTIAL_REMOVED = "credential-removed"
+
+
+class SimpleModelProviderEntity(BaseModel):
+    """
+    Simple provider.
+    """
+
+    provider: str
+    label: I18nObject
+    icon_small: I18nObject | None = None
+    icon_large: I18nObject | None = None
+    supported_model_types: list[ModelType]
+
+    def __init__(self, provider_entity: ProviderEntity):
+        """
+        Init simple provider.
+
+        :param provider_entity: provider entity
+        """
+        super().__init__(
+            provider=provider_entity.provider,
+            label=provider_entity.label,
+            icon_small=provider_entity.icon_small,
+            icon_large=provider_entity.icon_large,
+            supported_model_types=provider_entity.supported_model_types,
+        )
+
+
+class ProviderModelWithStatusEntity(ProviderModel):
+    """
+    Model class for model response.
+    """
+
+    status: ModelStatus
+    load_balancing_enabled: bool = False
+    has_invalid_load_balancing_configs: bool = False
+
+    def raise_for_status(self):
+        """
+        Check model status and raise ValueError if not active.
+
+        :raises ValueError: When model status is not active, with a descriptive message
+        """
+        if self.status == ModelStatus.ACTIVE:
+            return
+
+        error_messages = {
+            ModelStatus.NO_CONFIGURE: "Model is not configured",
+            ModelStatus.QUOTA_EXCEEDED: "Model quota has been exceeded",
+            ModelStatus.NO_PERMISSION: "No permission to use this model",
+            ModelStatus.DISABLED: "Model is disabled",
+        }
+
+        if self.status in error_messages:
+            raise ValueError(error_messages[self.status])
+
+
+class ModelWithProviderEntity(ProviderModelWithStatusEntity):
+    """
+    Model with provider entity.
+    """
+
+    provider: SimpleModelProviderEntity
+
+
+class DefaultModelProviderEntity(BaseModel):
+    """
+    Default model provider entity.
+    """
+
+    provider: str
+    label: I18nObject
+    icon_small: I18nObject | None = None
+    icon_large: I18nObject | None = None
+    supported_model_types: Sequence[ModelType] = []
+
+
+class DefaultModelEntity(BaseModel):
+    """
+    Default model entity.
+    """
+
+    model: str
+    model_type: ModelType
+    provider: DefaultModelProviderEntity
+
+    # pydantic configs
+    model_config = ConfigDict(protected_namespaces=())