dify/api/services/entities/model_provider_entities.py

from collections.abc import Sequence
from decimal import Decimal
from enum import StrEnum
from typing import Annotated, Any

from pydantic import BaseModel, ConfigDict, Field, model_validator

from configs import dify_config
from core.entities.model_entities import (
    ModelWithProviderEntity,
    ProviderModelWithStatusEntity,
)
from core.entities.provider_entities import (
    CredentialConfiguration,
    CustomModelConfiguration,
    ProviderQuotaType,
    QuotaConfiguration,
    UnaddedModelConfiguration,
)
from graphon.model_runtime.entities.common_entities import I18nObject
from graphon.model_runtime.entities.model_entities import (
    FetchFrom,
    ModelFeature,
    ModelPropertyKey,
    ModelType,
    ParameterRule,
)
from graphon.model_runtime.entities.provider_entities import (
    ConfigurateMethod,
    ModelCredentialSchema,
    ProviderCredentialSchema,
    ProviderHelpEntity,
)
from models.provider import ProviderType

_DECIMAL_STRING_PATTERN = r"^(?![-+.]*$)[+-]?0*\d*\.?\d*$"
CodegenSafeDecimal = Annotated[Decimal, Field(json_schema_extra={"pattern": _DECIMAL_STRING_PATTERN})]


class CustomConfigurationStatus(StrEnum):
    """
    Enum class for custom configuration status.
    """

    ACTIVE = "active"
    NO_CONFIGURE = "no-configure"


class CustomConfigurationResponse(BaseModel):
    """
    Model class for provider custom configuration response.
    """

    status: CustomConfigurationStatus
    current_credential_id: str | None = None
    current_credential_name: str | None = None
    available_credentials: list[CredentialConfiguration] | None = None
    custom_models: list[CustomModelConfiguration] | None = None
    can_added_models: list[UnaddedModelConfiguration] | None = None


class SystemConfigurationResponse(BaseModel):
    """
    Model class for provider system configuration response.
    """

    enabled: bool
    current_quota_type: ProviderQuotaType | None = None
    quota_configurations: list[QuotaConfiguration] = []


class ProviderResponse(BaseModel):
    """
    Model class for provider response.
    """

    tenant_id: str
    provider: str
    label: I18nObject
    description: I18nObject | None = None
    icon_small: I18nObject | None = None
    icon_small_dark: I18nObject | None = None
    background: str | None = None
    help: ProviderHelpEntity | None = None
    supported_model_types: Sequence[ModelType]
    configurate_methods: list[ConfigurateMethod]
    provider_credential_schema: ProviderCredentialSchema | None = None
    model_credential_schema: ModelCredentialSchema | None = None
    preferred_provider_type: ProviderType
    custom_configuration: CustomConfigurationResponse
    system_configuration: SystemConfigurationResponse

    # pydantic configs
    model_config = ConfigDict(protected_namespaces=())

    @model_validator(mode="after")
    def _(self):
        url_prefix = (
            dify_config.CONSOLE_API_URL + f"/console/api/workspaces/{self.tenant_id}/model-providers/{self.provider}"
        )
        if self.icon_small is not None:
            self.icon_small = I18nObject(
                en_US=f"{url_prefix}/icon_small/en_US", zh_Hans=f"{url_prefix}/icon_small/zh_Hans"
            )
        if self.icon_small_dark is not None:
            self.icon_small_dark = I18nObject(
                en_US=f"{url_prefix}/icon_small_dark/en_US",
                zh_Hans=f"{url_prefix}/icon_small_dark/zh_Hans",
            )
        return self


class ProviderWithModelsResponse(BaseModel):
    """
    Model class for provider with models response.
    """

    tenant_id: str
    provider: str
    label: I18nObject
    icon_small: I18nObject | None = None
    icon_small_dark: I18nObject | None = None
    status: CustomConfigurationStatus
    models: list[ProviderModelWithStatusEntity]

    @model_validator(mode="after")
    def _(self):
        url_prefix = (
            dify_config.CONSOLE_API_URL + f"/console/api/workspaces/{self.tenant_id}/model-providers/{self.provider}"
        )
        if self.icon_small is not None:
            self.icon_small = I18nObject(
                en_US=f"{url_prefix}/icon_small/en_US", zh_Hans=f"{url_prefix}/icon_small/zh_Hans"
            )

        if self.icon_small_dark is not None:
            self.icon_small_dark = I18nObject(
                en_US=f"{url_prefix}/icon_small_dark/en_US", zh_Hans=f"{url_prefix}/icon_small_dark/zh_Hans"
            )
        return self


class PriceConfigResponse(BaseModel):
    """Serialized pricing info with codegen-safe decimal string patterns."""

    input: CodegenSafeDecimal
    output: CodegenSafeDecimal | None = None
    unit: CodegenSafeDecimal
    currency: str


class AIModelEntityResponse(BaseModel):
    model: str
    label: I18nObject
    model_type: ModelType
    features: list[ModelFeature] | None = None
    fetch_from: FetchFrom
    model_properties: dict[ModelPropertyKey, Any]
    deprecated: bool = False
    parameter_rules: list[ParameterRule] = []
    pricing: PriceConfigResponse | None = None


class SimpleProviderEntityResponse(BaseModel):
    """
    Simple provider entity response.
    """

    provider: str
    provider_name: str = ""
    label: I18nObject
    icon_small: I18nObject | None = None
    icon_small_dark: I18nObject | None = None
    supported_model_types: Sequence[ModelType]
    tenant_id: str
    models: list[AIModelEntityResponse] = []

    @model_validator(mode="after")
    def _(self):
        url_prefix = (
            dify_config.CONSOLE_API_URL + f"/console/api/workspaces/{self.tenant_id}/model-providers/{self.provider}"
        )
        if self.icon_small is not None:
            self.icon_small = I18nObject(
                en_US=f"{url_prefix}/icon_small/en_US", zh_Hans=f"{url_prefix}/icon_small/zh_Hans"
            )

        if self.icon_small_dark is not None:
            self.icon_small_dark = I18nObject(
                en_US=f"{url_prefix}/icon_small_dark/en_US", zh_Hans=f"{url_prefix}/icon_small_dark/zh_Hans"
            )
        return self


class DefaultModelResponse(BaseModel):
    """
    Default model entity.
    """

    model: str
    model_type: ModelType
    provider: SimpleProviderEntityResponse

    # pydantic configs
    model_config = ConfigDict(protected_namespaces=())


class ModelWithProviderEntityResponse(ProviderModelWithStatusEntity):
    """
    Model with provider entity.
    """

    provider: SimpleProviderEntityResponse

    def __init__(self, tenant_id: str, model: ModelWithProviderEntity):
        dump_model = model.model_dump()
        dump_model["provider"]["tenant_id"] = tenant_id
        super().__init__(**dump_model)