module_registry.py

# =====================================================
# Apckeyl Framework
# Version 1.0
# module_registry.py
# =====================================================

"""
Apckeyl Module Registry.

Version 1.0

Реестр Compute Modules Apckeyl Framework.

Этот файл пока НЕ выполняет сетевые запросы.
Он только хранит информацию о доступных модулях
и предоставляет безопасные операции регистрации
и поиска.
"""


# =====================================================
# Module Registry
# =====================================================

class ModuleRegistry:

    def __init__(self):

        self._modules = {}


    # =================================================
    # Register Module
    # =================================================

    def register(
        self,
        module_id,
        module_name,
        module_type,
        endpoint=None,
        capabilities=None,
        version="1.0",
    ):

        if not module_id:
            raise ValueError(
                "module_id is required"
            )

        if not module_name:
            raise ValueError(
                "module_name is required"
            )

        if not module_type:
            raise ValueError(
                "module_type is required"
            )

        if capabilities is None:
            capabilities = []

        self._modules[module_id] = {

            "module_id": module_id,

            "module_name": module_name,

            "module_type": module_type,

            "endpoint": endpoint,

            "status": "unknown",

            "capabilities": list(capabilities),

            "version": version,
        }


    # =================================================
    # Remove Module
    # =================================================

    def unregister(self, module_id):

        if module_id in self._modules:

            del self._modules[module_id]


    # =================================================
    # Get Module
    # =================================================

    def get(self, module_id):

        return self._modules.get(
            module_id
        )


    # =================================================
    # Get All Modules
    # =================================================

    def get_all(self):

        return dict(
            self._modules
        )


    # =================================================
    # Find By Type
    # =================================================

    def find_by_type(self, module_type):

        return [

            module

            for module
            in self._modules.values()

            if module["module_type"] == module_type
        ]


    # =================================================
    # Update Status
    # =================================================

    def update_status(
        self,
        module_id,
        status,
    ):

        module = self.get(
            module_id
        )

        if module is None:

            raise KeyError(
                f"Unknown module: {module_id}"
            )

        module["status"] = status


    # =================================================
    # Check Module
    # =================================================

    def exists(self, module_id):

        return module_id in self._modules


# =====================================================
# Default Registry
# =====================================================

registry = ModuleRegistry()


# =====================================================
# First Production Module
# =====================================================

registry.register(

    module_id="realesrgan",

    module_name="Apckeyl_RealESRGAN",

    module_type="image_upscaler",

    endpoint=None,

    capabilities=[
        "image_upscale",
    ],

    version="1.0",
)


# =====================================================
# Initial Status
# =====================================================

registry.update_status(

    module_id="realesrgan",

    status="available",
)