dify/api/core/sandbox/factory.py

from collections.abc import Mapping, Sequence
from enum import StrEnum
from typing import TYPE_CHECKING, Any

from core.virtual_environment.__base.virtual_environment import VirtualEnvironment

if TYPE_CHECKING:
    from core.sandbox.initializer import SandboxInitializer


class VMType(StrEnum):
    DOCKER = "docker"
    E2B = "e2b"
    LOCAL = "local"


class VMFactory:
    @classmethod
    def create(
        cls,
        tenant_id: str,
        vm_type: VMType,
        options: Mapping[str, Any] | None = None,
        environments: Mapping[str, str] | None = None,
        user_id: str | None = None,
        initializers: Sequence["SandboxInitializer"] | None = None,
    ) -> VirtualEnvironment:
        options = options or {}
        environments = environments or {}

        vm_class = cls._get_vm_class(vm_type)
        vm = vm_class(tenant_id=tenant_id, options=options, environments=environments, user_id=user_id)

        if initializers:
            for initializer in initializers:
                initializer.initialize(vm)

        return vm

    @classmethod
    def _get_vm_class(cls, vm_type: VMType) -> type[VirtualEnvironment]:
        match vm_type:
            case VMType.DOCKER:
                from core.virtual_environment.providers.docker_daemon_sandbox import DockerDaemonEnvironment

                return DockerDaemonEnvironment
            case VMType.E2B:
                from core.virtual_environment.providers.e2b_sandbox import E2BEnvironment

                return E2BEnvironment
            case VMType.LOCAL:
                from core.virtual_environment.providers.local_without_isolation import LocalVirtualEnvironment

                return LocalVirtualEnvironment
            case _:
                raise ValueError(f"Unsupported VM type: {vm_type}")

    @classmethod
    def validate(cls, vm_type: VMType, options: Mapping[str, Any]) -> None:
        vm_class = cls._get_vm_class(vm_type)
        vm_class.validate(options)