taskweaver/plugin/context.py

import contextlib
from abc import ABC, abstractmethod
from typing import Any, Dict, List, Literal, Optional, Tuple

LogErrorLevel = Literal["info", "warning", "error"]
ArtifactType = Literal["chart", "image", "df", "file", "txt", "svg", "html"]


class PluginContext(ABC):
    """
    interface for API to interact with execution environment of plugin

    The runtime will provide an implementation of this interface to the plugin.
    Plugin could use the API provded withotu need to implement this interface.
    """

    @property
    @abstractmethod
    def env_id(self) -> str:
        """get the environment id of the plugin"""
        ...

    @property
    @abstractmethod
    def session_id(self) -> str:
        """get the session id of the plugin"""
        ...

    @property
    @abstractmethod
    def execution_id(self) -> str:
        """get the execution id of the plugin"""
        ...

    @abstractmethod
    def add_artifact(
        self,
        name: str,
        file_name: str,
        type: ArtifactType,
        val: Any,
        desc: Optional[str] = None,
    ) -> str:
        """
        add an artifact to the execution context

        :param name: the name of the artifact
        :param file_name: the name of the file
        :param type: the type of the artifact
        :param val: the value of the artifact
        :param desc: the description of the artifact

        :return: the id of the artifact
        """
        ...

    @abstractmethod
    def create_artifact_path(
        self,
        name: str,
        file_name: str,
        type: ArtifactType,
        desc: str,
    ) -> Tuple[str, str]:
        """
        create a path for an artifact and the plugin can use this path to save the artifact.
        This methods is provided for the plugin to save the artifact by itself rather than saving by the runtime,
        for general cases when the file content could be passed directly, the plugin should use add_artifact instead.

        :param name: the name of the artifact
        :param file_name: the name of the file
        :param type: the type of the artifact
        :param desc: the description of the artifact

        :return: the id and the path of the artifact
        """
        ...

    @abstractmethod
    def get_session_var(
        self,
        variable_name: str,
        default: Optional[str],
    ) -> Optional[str]:
        """
        get a session variable from the context

        :param variable_name: the name of the variable
        :param default: the default value of the variable

        :return: the value of the variable
        """
        ...

    @abstractmethod
    def log(self, level: LogErrorLevel, tag: str, message: str) -> None:
        """log a message from the plugin"""

    @abstractmethod
    def get_env(self, plugin_name: str, variable_name: str) -> str:
        """get an environment variable from the context"""


class TestPluginContxt(PluginContext):
    """
    This plugin context is used for testing purpose.
    """

    def __init__(self, temp_dir: str) -> None:
        self._session_id = "test"
        self._env_id = "test"
        self._execution_id = "test"
        self._logs: List[Tuple[LogErrorLevel, str, str]] = []
        self._env: Dict[str, str] = {}
        self._session_var: Dict[str, str] = {}
        self._temp_dir = temp_dir
        self._artifacts: List[Dict[str, str]] = []

    @property
    def env_id(self) -> str:
        return "test"

    @property
    def session_id(self) -> str:
        return "test"

    @property
    def execution_id(self) -> str:
        return "test"

    def add_artifact(
        self,
        name: str,
        file_name: str,
        type: ArtifactType,
        val: Any,
        desc: Optional[str] = None,
    ) -> str:
        id = f"test_artifact_id_{len(self._artifacts)}"
        self._artifacts.append(
            {
                "id": id,
                "name": name,
                "file_name": file_name,
                "type": type,
                "desc": desc or "",
            },
        )
        return id

    def create_artifact_path(
        self,
        name: str,
        file_name: str,
        type: ArtifactType,
        desc: str,
    ) -> Tuple[str, str]:
        id = f"test_artifact_id_{len(self._artifacts)}"
        self._artifacts.append(
            {
                "id": id,
                "name": name,
                "file_name": file_name,
                "type": type,
                "desc": desc or "",
            },
        )
        return id, self._temp_dir + "/" + file_name

    def log(self, level: LogErrorLevel, tag: str, message: str) -> None:
        return self._logs.append((level, tag, message))

    def get_env(self, plugin_name: str, variable_name: str) -> str:
        return self._env[plugin_name + "_" + variable_name]

    def get_session_var(
        self,
        variable_name: str,
        default: Optional[str],
    ) -> Optional[str]:
        return self._session_var.get(variable_name, default)


@contextlib.contextmanager
def temp_context(workspace_dir: Optional[str] = None):
    import os
    import shutil
    import tempfile
    import uuid

    if workspace_dir is None:
        workspace_dir = tempfile.mkdtemp()
    else:
        workspace_dir = os.path.join(workspace_dir, str(uuid.uuid4()))
        os.makedirs(workspace_dir)

    try:
        yield TestPluginContxt(workspace_dir)
    finally:
        shutil.rmtree(workspace_dir)