datasource_manifest.py

from enum import Enum, StrEnum
from typing import Any, Union

from pydantic import BaseModel, Field, ValidationInfo, field_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,
    ParameterAutoGenerate,
    ParameterOption,
    ParameterTemplate,
)
from dify_plugin.entities.oauth import OAuthSchema
from dify_plugin.entities.provider_config import CommonParameterType, ProviderConfig


@docs(
    description="The label of the datasource",
)
class DatasourceLabelEnum(Enum):
    SEARCH = "search"
    IMAGE = "image"
    VIDEOS = "videos"
    WEATHER = "weather"
    FINANCE = "finance"
    DESIGN = "design"
    TRAVEL = "travel"
    SOCIAL = "social"
    NEWS = "news"
    MEDICAL = "medical"
    PRODUCTIVITY = "productivity"
    EDUCATION = "education"
    BUSINESS = "business"
    ENTERTAINMENT = "entertainment"
    UTILITIES = "utilities"
    OTHER = "other"


@docs(
    name="DatasourceProviderType",
    description="Type of datasource provider",
)
class DatasourceProviderType(StrEnum):
    """
    Enum class for datasource provider
    """

    ONLINE_DOCUMENT = "online_document"
    WEBSITE_CRAWL = "website_crawl"
    ONLINE_DRIVE = "online_drive"


@docs(
    name="DatasourceParameter",
    description="Parameter of datasource entity",
)
class DatasourceParameter(BaseModel):
    """
    Overrides type
    """

    @docs(
        description="Datasource parameter type (aliases CommonParameterType values)",
    )
    class DatasourceParameterType(StrEnum):
        STRING = CommonParameterType.STRING.value
        NUMBER = CommonParameterType.NUMBER.value
        BOOLEAN = CommonParameterType.BOOLEAN.value
        SELECT = CommonParameterType.SELECT.value
        SECRET_INPUT = CommonParameterType.SECRET_INPUT.value
        FILE = CommonParameterType.FILE.value
        FILES = CommonParameterType.FILES.value

    name: str = Field(..., description="The name of the parameter")
    label: I18nObject = Field(..., description="The label presented to the user")
    placeholder: I18nObject | None = Field(default=None, description="The placeholder presented to the user")
    scope: str | None = None
    auto_generate: ParameterAutoGenerate | None = None
    template: ParameterTemplate | None = None
    required: bool = False
    default: Union[float, int, str] | None = None
    min: Union[float, int] | None = None
    max: Union[float, int] | None = None
    precision: int | None = None
    options: list[ParameterOption] = Field(default_factory=list)
    type: DatasourceParameterType = Field(..., description="The type of the parameter")
    description: I18nObject = Field(..., description="The description of the parameter")


@docs(
    name="DatasourceIdentity",
    description="Identity of datasource entity",
)
class DatasourceIdentity(BaseModel):
    author: str = Field(..., description="The author of the datasource")
    name: str = Field(..., description="The name of the datasource")
    label: I18nObject = Field(..., description="The label of the datasource")
    icon: str | None = None


@docs(
    name="DatasourceEntityExtra",
    description="The extra of the datasource entity",
)
class DatasourceEntityExtra(BaseModel):
    class Python(BaseModel):
        source: str

    python: Python


@docs(
    name="Datasource",
    description="Datasource entity",
)
class DatasourceEntity(BaseModel):
    identity: DatasourceIdentity
    parameters: list[DatasourceParameter] = Field(default_factory=list)
    description: I18nObject = Field(..., description="The label of the datasource")
    output_schema: dict[str, Any] = Field(default_factory=dict, description="Output schema definition")
    extra: DatasourceEntityExtra

    @field_validator("parameters", mode="before")
    @classmethod
    def set_parameters(cls, v, validation_info: ValidationInfo) -> list[DatasourceParameter]:
        return v or []


@docs(
    description="Identity of datasource provider",
)
class DatasourceProviderIdentity(BaseModel):
    author: str = Field(..., description="The author of the datasource")
    name: str = Field(..., description="The name of the datasource")
    description: I18nObject = Field(..., description="The description of the datasource")
    icon: str = Field(..., description="The icon of the datasource")
    label: I18nObject = Field(..., description="The label of the datasource")
    tags: list[DatasourceLabelEnum] | None = Field(
        default=[],
        description="The tags of the datasource",
    )


@docs(
    name="DatasourceProviderExtra",
    description="The extra of the datasource provider",
)
class DatasourceProviderConfigurationExtra(BaseModel):
    class Python(BaseModel):
        source: str

    python: Python


@docs(
    name="DatasourceProvider",
    description="Manifest of datasource providers",
    outside_reference_fields={"datasources": DatasourceEntity},
)
class DatasourceProviderManifest(BaseModel):
    """
    Datasource provider entity
    """

    identity: DatasourceProviderIdentity = Field(..., description="The identity of the datasource provider")
    credentials_schema: list[ProviderConfig] = Field(
        default_factory=list, description="The credentials schema of the datasource provider"
    )
    oauth_schema: OAuthSchema | None = Field(
        default=None, description="The OAuth schema of the datasource provider if OAuth is supported"
    )
    provider_type: DatasourceProviderType = Field(..., description="The type of the datasource provider")
    datasources: list[DatasourceEntity] = Field(
        default_factory=list, description="The datasources of the datasource provider"
    )
    extra: DatasourceProviderConfigurationExtra = Field(..., description="The extra of the datasource provider")

    @field_validator("datasources", mode="before")
    @classmethod
    def validate_datasources(cls, value) -> list[DatasourceEntity]:
        if not isinstance(value, list):
            raise ValueError("datasources should be a list")

        datasources: list[DatasourceEntity] = []

        for datasource in value:
            if not isinstance(datasource, str):
                raise ValueError("datasource path should be a string")
            try:
                file = load_yaml_file(datasource)
                datasources.append(DatasourceEntity(**file))
            except Exception as e:
                raise ValueError(f"Error loading datasource configuration: {e!s}") from e

        return datasources