mephisto.abstractions.providers.mturk_sandbox.sandbox_mturk_provider

#!/usr/bin/env python3

# Copyright (c) Meta Platforms and its affiliates.
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

from mephisto.abstractions.providers.mturk_sandbox.provider_type import PROVIDER_TYPE
from mephisto.abstractions.providers.mturk.mturk_provider import (
    MTurkProvider,
    MTurkProviderArgs,
)
from mephisto.abstractions.providers.mturk_sandbox.sandbox_mturk_agent import (
    SandboxMTurkAgent,
)
from mephisto.abstractions.providers.mturk_sandbox.sandbox_mturk_requester import (
    SandboxMTurkRequester,
)
from mephisto.abstractions.providers.mturk_sandbox.sandbox_mturk_unit import (
    SandboxMTurkUnit,
)
from mephisto.abstractions.providers.mturk_sandbox.sandbox_mturk_worker import (
    SandboxMTurkWorker,
)
from mephisto.operations.registry import register_mephisto_abstraction

import os
from dataclasses import dataclass

from typing import Any, ClassVar, Type, List, TYPE_CHECKING

if TYPE_CHECKING:
    from mephisto.data_model.unit import Unit
    from mephisto.data_model.worker import Worker
    from mephisto.data_model.requester import Requester
    from mephisto.data_model.agent import Agent


@dataclass
class SandboxMTurkProviderArgs(MTurkProviderArgs):
    """Provider args for a sandbox MTurk provider"""

    _provider_type: str = PROVIDER_TYPE


@register_mephisto_abstraction()
class SandboxMTurkProvider(MTurkProvider):
    """
    Mock implementation of a CrowdProvider that stores everything
    in a local state in the class for use in tests.
    """

    # Ensure inherited methods use this level's provider type
    PROVIDER_TYPE = PROVIDER_TYPE

    UnitClass: ClassVar[Type["Unit"]] = SandboxMTurkUnit

    RequesterClass: ClassVar[Type["Requester"]] = SandboxMTurkRequester

    WorkerClass: ClassVar[Type["Worker"]] = SandboxMTurkWorker

    AgentClass: ClassVar[Type["Agent"]] = SandboxMTurkAgent

    ArgsClass = SandboxMTurkProviderArgs

    def _get_client(self, requester_name: str) -> Any:
        """
        Get an mturk client for usage with mturk_utils
        """
        return self.datastore.get_sandbox_client_for_requester(requester_name)

    @classmethod
    def get_wrapper_js_path(cls):
        """
        Return the path to the `wrap_crowd_source.js` file for this
        provider to be deployed to the server
        """
        return os.path.join(os.path.dirname(__file__), "wrap_crowd_source.js")

    def cleanup_qualification(self, qualification_name: str) -> None:
        """Remove the qualification from the sandbox server"""
        return super().cleanup_qualification(f"{qualification_name}_sandbox")

class SandboxMTurkProviderArgs(MTurkProviderArgs):
    """Provider args for a sandbox MTurk provider"""

    _provider_type: str = PROVIDER_TYPE

class SandboxMTurkProvider(MTurkProvider):
    """
    Mock implementation of a CrowdProvider that stores everything
    in a local state in the class for use in tests.
    """

    # Ensure inherited methods use this level's provider type
    PROVIDER_TYPE = PROVIDER_TYPE

    UnitClass: ClassVar[Type["Unit"]] = SandboxMTurkUnit

    RequesterClass: ClassVar[Type["Requester"]] = SandboxMTurkRequester

    WorkerClass: ClassVar[Type["Worker"]] = SandboxMTurkWorker

    AgentClass: ClassVar[Type["Agent"]] = SandboxMTurkAgent

    ArgsClass = SandboxMTurkProviderArgs

    def _get_client(self, requester_name: str) -> Any:
        """
        Get an mturk client for usage with mturk_utils
        """
        return self.datastore.get_sandbox_client_for_requester(requester_name)

    @classmethod
    def get_wrapper_js_path(cls):
        """
        Return the path to the `wrap_crowd_source.js` file for this
        provider to be deployed to the server
        """
        return os.path.join(os.path.dirname(__file__), "wrap_crowd_source.js")

    def cleanup_qualification(self, qualification_name: str) -> None:
        """Remove the qualification from the sandbox server"""
        return super().cleanup_qualification(f"{qualification_name}_sandbox")

class SandboxMTurkUnit(MTurkUnit):
    """
    This class tracks the status of an individual worker's contribution to a
    higher level assignment. It is the smallest 'unit' of work to complete
    the assignment, and this class is only responsible for checking
    the status of that work itself being done.
    """

    # Ensure inherited methods use this level's provider type
    PROVIDER_TYPE = PROVIDER_TYPE

    def _get_client(self, requester_name: str) -> Any:
        """
        Get an mturk client for usage with mturk_utils
        """
        return self.datastore.get_sandbox_client_for_requester(requester_name)

    @staticmethod
    def new(
        db: "MephistoDB", assignment: "Assignment", index: int, pay_amount: float
    ) -> "Unit":
        """Create a Unit for the given assignment"""
        return SandboxMTurkUnit._register_unit(
            db, assignment, index, pay_amount, PROVIDER_TYPE
        )

class SandboxMTurkRequester(MTurkRequester):
    """Wrapper around regular requester that handles removing the appended "sandbox" name"""

    # Ensure inherited methods use this level's provider type
    PROVIDER_TYPE = PROVIDER_TYPE

    def __init__(
        self,
        db: "MephistoDB",
        db_id: str,
        row: Optional[Mapping[str, Any]] = None,
        _used_new_call: bool = False,
    ):
        super().__init__(db, db_id, row=row, _used_new_call=_used_new_call)
        self.datastore: "MTurkDatastore" = self.db.get_datastore_for_provider(
            self.PROVIDER_TYPE
        )
        # Use _requester_name to preserve sandbox behavior which
        # utilizes a different requester_name
        assert self.requester_name.endswith(
            "_sandbox"
        ), f"{self.requester_name} is not a sandbox requester"
        self._requester_name = self.requester_name[:-8]

    def _get_client(self, requester_name: str) -> Any:
        """
        Get an mturk client for usage with mturk_utils
        """
        return self.datastore.get_sandbox_client_for_requester(requester_name)

    @classmethod
    def is_sandbox(cls) -> bool:
        """
        Determine if this is a requester on sandbox
        """
        return True

    # Required functions for a Requester implementation

    @staticmethod
    def new(db: "MephistoDB", requester_name: str) -> "Requester":
        if not requester_name.endswith("_sandbox"):
            requester_name += "_sandbox"
        return SandboxMTurkRequester._register_requester(
            db, requester_name, PROVIDER_TYPE
        )

class SandboxMTurkWorker(MTurkWorker):
    """
    This class represents an individual - namely a person. It maintains components of ongoing identity for a user.
    """

    # Ensure inherited methods use this level's provider type
    PROVIDER_TYPE = PROVIDER_TYPE

    def __init__(
        self,
        db: "MephistoDB",
        db_id: str,
        row: Optional[Mapping[str, Any]] = None,
        _used_new_call: bool = False,
    ):
        super().__init__(db, db_id, row=row, _used_new_call=_used_new_call)
        self.datastore: "MTurkDatastore" = self.db.get_datastore_for_provider(
            self.PROVIDER_TYPE
        )
        # sandbox workers use a different name
        self._worker_name = self.worker_name[:-8]

    def grant_crowd_qualification(
        self, qualification_name: str, value: int = 1
    ) -> None:
        """
        Grant a qualification by the given name to this worker. Check the local
        MTurk db to find the matching MTurk qualification to grant, and pass
        that. If no qualification exists, try to create one.
        """
        return super().grant_crowd_qualification(qualification_name + "_sandbox", value)

    def revoke_crowd_qualification(self, qualification_name: str) -> None:
        """
        Revoke the qualification by the given name from this worker. Check the local
        MTurk db to find the matching MTurk qualification to revoke, pass if
        no such qualification exists.
        """
        return super().revoke_crowd_qualification(qualification_name + "_sandbox")

    def _get_client(self, requester_name: str) -> Any:
        """
        Get an mturk client for usage with mturk_utils
        """
        return self.datastore.get_sandbox_client_for_requester(requester_name)

    @staticmethod
    def new(db: "MephistoDB", worker_id: str) -> "Worker":
        return MTurkWorker._register_worker(db, worker_id + "_sandbox", PROVIDER_TYPE)

class SandboxMTurkAgent(MTurkAgent):
    """
    Wrapper for a regular MTurk agent that will only communicate with sandbox
    """

    # Ensure inherited methods use this level's provider type
    PROVIDER_TYPE = PROVIDER_TYPE

    def _get_client(self) -> Any:
        """
        Get an mturk client for usage with mturk_utils for this agent
        """
        unit = self.get_unit()
        requester: "SandboxMTurkRequester" = cast(
            "SandboxMTurkRequester", unit.get_requester()
        )
        return self.datastore.get_sandbox_client_for_requester(
            requester._requester_name
        )

    @staticmethod
    def new(db: "MephistoDB", worker: "Worker", unit: "Unit") -> "Agent":
        """Create an agent for this worker to be used for work on the given Unit."""
        return SandboxMTurkAgent._register_agent(db, worker, unit, PROVIDER_TYPE)

class SandboxMTurkProviderArgs(MTurkProviderArgs):
    """Provider args for a sandbox MTurk provider"""

    _provider_type: str = PROVIDER_TYPE