provider.py

import glob
import os
from collections.abc import Sequence
from enum import Enum

from pydantic import BaseModel, ConfigDict, Field, model_validator

from dify_plugin.core.documentation.schema_doc import docs
from dify_plugin.core.utils.yaml_loader import load_yaml_file
from dify_plugin.entities import I18nObject
from dify_plugin.entities.model import AIModelEntity, ModelType


@docs(
    description="Configurate method",
    name="ModelConfigurateMethod",
)
class ConfigurateMethod(Enum):
    """
    Enum class for configurate method of provider model.
    """

    PREDEFINED_MODEL = "predefined-model"
    CUSTOMIZABLE_MODEL = "customizable-model"


@docs(
    description="Model form type",
    name="ModelFormType",
)
class FormType(Enum):
    """
    Enum class for form type.
    """

    TEXT_INPUT = "text-input"
    SECRET_INPUT = "secret-input"
    SELECT = "select"
    RADIO = "radio"
    SWITCH = "switch"


@docs(
    description="Form show on",
    name="ModelFormShowOnObject",
)
class FormShowOnObject(BaseModel):
    """
    Model class for form show on.
    """

    variable: str
    value: str


@docs(
    description="Form option",
    name="ModelFormOption",
)
class FormOption(BaseModel):
    """
    Model class for form option.
    """

    label: I18nObject
    value: str
    show_on: list[FormShowOnObject] = Field(default_factory=list)

    def __init__(self, **data):
        super().__init__(**data)
        if not self.label:
            self.label = I18nObject(en_US=self.value)


@docs(
    description="Credential form schema",
    name="ModelCredentialFormSchema",
)
class CredentialFormSchema(BaseModel):
    """
    Model class for credential form schema.
    """

    variable: str
    label: I18nObject
    type: FormType
    required: bool = True
    default: str | None = None
    options: list[FormOption] | None = None
    placeholder: I18nObject | None = None
    max_length: int = 0
    show_on: list[FormShowOnObject] = Field(default_factory=list)


@docs(
    description="Model provider credential schema",
    name="ModelProviderCredentialSchema",
)
class ProviderCredentialSchema(BaseModel):
    """
    Model class for provider credential schema.
    """

    credential_form_schemas: list[CredentialFormSchema]


@docs(
    description="Field model schema",
    name="ModelFieldModelSchema",
)
class FieldModelSchema(BaseModel):
    label: I18nObject
    placeholder: I18nObject | None = None


class ModelCredentialSchema(BaseModel):
    """
    Model class for model credential schema.
    """

    model: FieldModelSchema
    credential_form_schemas: list[CredentialFormSchema]


class SimpleProviderEntity(BaseModel):
    """
    Simple model class for provider.
    """

    provider: str
    label: I18nObject
    icon_small: I18nObject | None = None
    icon_large: I18nObject | None = None
    icon_small_dark: I18nObject | None = None
    icon_large_dark: I18nObject | None = None
    supported_model_types: Sequence[ModelType]
    models: list[AIModelEntity] = []


@docs(
    description="Model provider help",
    name="ModelProviderHelp",
)
class ProviderHelpEntity(BaseModel):
    """
    Model class for provider help.
    """

    title: I18nObject
    url: I18nObject


@docs(
    description="Model position",
    name="ModelPosition",
)
class ModelPosition(BaseModel):
    """
    Model class for ai models
    """

    llm: list[str] | None = Field(
        default_factory=list, description="Sorts of llm model in ascending order, fill model name here"
    )
    text_embedding: list[str] | None = Field(
        default_factory=list, description="Sorts of text embedding model in ascending order, fill model name here"
    )
    rerank: list[str] | None = Field(
        default_factory=list, description="Sorts of rerank model in ascending order, fill model name here"
    )
    tts: list[str] | None = Field(
        default_factory=list, description="Sorts of tts model in ascending order, fill model name here"
    )
    speech2text: list[str] | None = Field(
        default_factory=list, description="Sorts of speech2text model in ascending order, fill model name here"
    )
    moderation: list[str] | None = Field(
        default_factory=list, description="Sorts of moderation model in ascending order, fill model name here"
    )


class ProviderEntity(BaseModel):
    """
    Model class for provider.
    """

    provider: str
    label: I18nObject
    description: I18nObject | None = None
    icon_small: I18nObject | None = None
    icon_large: I18nObject | None = None
    icon_small_dark: I18nObject | None = None
    icon_large_dark: I18nObject | None = None
    background: str | None = None
    help: ProviderHelpEntity | None = None
    supported_model_types: Sequence[ModelType]
    configurate_methods: list[ConfigurateMethod]
    models: list[AIModelEntity] = Field(default_factory=list)
    provider_credential_schema: ProviderCredentialSchema | None = None
    model_credential_schema: ModelCredentialSchema | None = None
    position: ModelPosition | None = None

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

    def to_simple_provider(self) -> SimpleProviderEntity:
        """
        Convert to simple provider.

        :return: simple provider
        """
        return SimpleProviderEntity(
            provider=self.provider,
            label=self.label,
            icon_small=self.icon_small,
            icon_large=self.icon_large,
            icon_small_dark=self.icon_small_dark,
            icon_large_dark=self.icon_large_dark,
            supported_model_types=self.supported_model_types,
            models=self.models,
        )

    @model_validator(mode="before")
    @classmethod
    def validate_models(cls, values) -> dict:
        value = values.get("models", {})
        if not isinstance(value, dict):
            raise ValueError("models should be a glob path list")

        cwd = os.getcwd()

        model_entities = []

        def load_models(model_type: str):
            if model_type not in value:
                return

            for path in value[model_type].get("predefined", []):
                yaml_paths = glob.glob(os.path.join(cwd, path))
                for yaml_path in yaml_paths:
                    if yaml_path.endswith("_position.yaml"):
                        if "position" not in values:
                            values["position"] = {}

                        position = load_yaml_file(yaml_path)
                        values["position"][model_type] = position
                    else:
                        model_entity = load_yaml_file(yaml_path)
                        if not model_entity:
                            raise ValueError(f"Error loading model entity: {yaml_path}")

                        provider_model = AIModelEntity(**model_entity)
                        model_entities.append(provider_model)

        load_models("llm")
        load_models("text_embedding")
        load_models("rerank")
        load_models("tts")
        load_models("speech2text")
        load_models("moderation")

        values["models"] = model_entities

        return values


@docs(
    description="Model provider configuration extra",
    name="ModelProviderExtra",
)
class ModelProviderConfigurationExtra(BaseModel):
    class Python(BaseModel):
        provider_source: str
        model_sources: list[str] = Field(default_factory=list)

        model_config = ConfigDict(protected_namespaces=())

    python: Python


@docs(
    name="ModelProvider",
    description="Model provider configuration",
    outside_reference_fields={"models": AIModelEntity},
)
class ModelProviderConfiguration(ProviderEntity):
    extra: ModelProviderConfigurationExtra


# class ProviderConfig(BaseModel):
#     """
#     Model class for provider config.
#     """

#     provider: str
#     credentials: dict