nodekit 0.2.0__py3-none-any.whl

nodekit/events/__init__.py

@@ -0,0 +1,30 @@
+__all__ = [
+    "EventTypeEnum",
+    "Event",
+    # Concrete classes:
+    "TraceStartedEvent",
+    "TraceEndedEvent",
+    "PageSuspendedEvent",
+    "PageResumedEvent",
+    "NodeStartedEvent",
+    "ActionTakenEvent",
+    "NodeEndedEvent",
+    "BrowserContextSampledEvent",
+    "KeySampledEvent",
+    "PointerSampledEvent",
+]
+
+from nodekit._internal.types.events.events import (
+    Event,
+    EventTypeEnum,
+    TraceStartedEvent,
+    TraceEndedEvent,
+    NodeStartedEvent,
+    ActionTakenEvent,
+    NodeEndedEvent,
+    KeySampledEvent,
+    PointerSampledEvent,
+    PageSuspendedEvent,
+    PageResumedEvent,
+    BrowserContextSampledEvent,
+)

nodekit/experimental/recruitment_services/base.py

@@ -0,0 +1,77 @@
+from abc import ABC, abstractmethod
+from decimal import Decimal
+from typing import List, Iterable, Literal
+
+import pydantic
+
+
+# %%
+class RecruiterCredentialsError(Exception):
+    """Exception raised for errors in the recruiter credentials."""
+
+    ...
+
+
+# %%
+class ListAssignmentsItem(pydantic.BaseModel):
+    hit_id: str
+    worker_id: str
+    assignment_id: str
+    status: Literal["Submitted", "Approved", "Rejected"]
+    submission_payload: str
+
+
+class CreateHitRequest(pydantic.BaseModel):
+    entrypoint_url: str
+    title: str
+    description: str
+    keywords: List[str]
+    num_assignments: int
+    duration_sec: int
+    completion_reward_usd: Decimal
+    allowed_participant_ids: List[str]
+    unique_request_token: str
+
+
+class CreateHitResponse(pydantic.BaseModel):
+    hit_id: str
+
+
+class SendBonusPaymentRequest(pydantic.BaseModel):
+    worker_id: str
+    assignment_id: str
+    amount_usd: Decimal = pydantic.Field(decimal_places=2)
+
+
+# %%
+class RecruiterServiceClient(ABC):
+    @abstractmethod
+    def get_recruiter_service_name(self) -> str: ...
+
+    @abstractmethod
+    def create_hit(
+        self,
+        request: CreateHitRequest,
+    ) -> CreateHitResponse: ...
+
+    @abstractmethod
+    def send_bonus_payment(
+        self,
+        request: SendBonusPaymentRequest,
+    ) -> None: ...
+
+    @abstractmethod
+    def iter_assignments(
+        self,
+        hit_id: str,
+    ) -> Iterable[ListAssignmentsItem]:
+        raise NotImplementedError
+
+    @abstractmethod
+    def cleanup_hit(self, hit_id: str) -> None: ...
+
+    @abstractmethod
+    def approve_assignment(
+        self,
+        assignment_id: str,
+    ) -> None: ...

nodekit/experimental/recruitment_services/mechanical_turk/client.py

@@ -0,0 +1,359 @@
+from decimal import Decimal
+from typing import List, Iterable
+
+from boto3.session import Session
+
+
+from nodekit.experimental.recruitment_services.base import (
+    RecruiterServiceClient,
+    CreateHitRequest,
+    CreateHitResponse,
+    SendBonusPaymentRequest,
+    ListAssignmentsItem,
+    RecruiterCredentialsError,
+)
+import nodekit.experimental.recruitment_services.mechanical_turk.models as boto3_models
+from uuid import uuid4
+
+import datetime
+import re
+
+
+def extract_trace(xml: str) -> str:
+    # pull the contents of the <FreeText> element
+    m = re.search(r"<FreeText>(.*?)</FreeText>", xml, re.DOTALL)
+    if not m:
+        raise ValueError("No <FreeText> found")
+    raw = m.group(1)
+
+    # MTurk sometimes form-encodes it (+ for space, %xx etc.)
+    return raw
+
+
+# %%
+class MturkClient(RecruiterServiceClient):
+    def __init__(
+        self,
+        aws_access_key_id: str,
+        aws_secret_access_key: str,
+        sandbox: bool,
+    ):
+        # Initialize MTurk client
+        if sandbox:
+            endpoint_url = "https://mturk-requester-sandbox.us-east-1.amazonaws.com"
+        else:
+            endpoint_url = "https://mturk-requester.us-east-1.amazonaws.com"
+
+        session = Session(
+            aws_access_key_id=aws_access_key_id,
+            aws_secret_access_key=aws_secret_access_key,
+        )
+
+        self.boto3_client = session.client(
+            service_name="mturk", endpoint_url=endpoint_url, region_name="us-east-1"
+        )
+        self.sandbox = sandbox
+
+        # Try verifying the credentials; throw if invalid:
+        try:
+            self.boto3_client.get_account_balance()
+        except Exception as e:
+            raise RecruiterCredentialsError from e
+
+    def get_recruiter_service_name(self) -> str:
+        if self.sandbox:
+            return "MTurkSandbox"
+        else:
+            return "MTurk"
+
+    def create_hit(
+        self,
+        request: CreateHitRequest,
+    ) -> CreateHitResponse:
+        # Unpack:
+        entrypoint_url = request.entrypoint_url
+        title = request.title
+        description = request.description
+        keywords = request.keywords
+        num_assignments = request.num_assignments
+        duration_sec = request.duration_sec
+        completion_reward_usd = request.completion_reward_usd
+        allowed_participant_ids = request.allowed_participant_ids
+
+        # Calculate minimum approval cost
+        min_approval_cost = (
+            completion_reward_usd * Decimal("1.2") * Decimal(num_assignments)
+        )  # Turk fees are 20% of the base completion reward.
+        current_balance = Decimal(
+            self.boto3_client.get_account_balance()["AvailableBalance"]
+        )
+        if current_balance < min_approval_cost:
+            raise RuntimeError(
+                f"Insufficient balance to create HIT. Minimum required: ${min_approval_cost:.2f}, current balance: ${current_balance:.2f}"
+            )
+
+        qualification_requirements: List[boto3_models.QualificationRequirement] = []
+        if len(allowed_participant_ids) > 0:
+            # Create qualification type for this TaskRequest:
+            qual_type = self.create_qualification_type(unique_name=f"psy:{uuid4()}")
+
+            # Grant each worker ID the qualification:
+            for worker_id in allowed_participant_ids:
+                self.grant_qualification(
+                    qualification_type_id=qual_type.QualificationTypeId,
+                    worker_id=worker_id,
+                )
+
+            # Create qualification requirement for this HIT:
+            qual_requirement = self.package_qualification_exists_requirement(
+                qual_type=qual_type
+            )
+            qualification_requirements.append(qual_requirement)
+
+        hit_type_response = self.boto3_client.create_hit_type(
+            AutoApprovalDelayInSeconds=1,
+            AssignmentDurationInSeconds=duration_sec,
+            Reward=str(completion_reward_usd),
+            Title=title,
+            Keywords=",".join(keywords),
+            Description=description,
+            QualificationRequirements=[
+                qr.model_dump(mode="json") for qr in qualification_requirements
+            ],
+        )
+
+        q = (
+            f'<ExternalQuestion xmlns="http://mechanicalturk.amazonaws.com/AWSMechanicalTurkDataSchemas/2006-07-14/ExternalQuestion.xsd">'
+            f"<ExternalURL>{entrypoint_url}</ExternalURL>"
+            f"<FrameHeight>{0}</FrameHeight></ExternalQuestion>"
+        )
+
+        hit_info = self.boto3_client.create_hit_with_hit_type(
+            HITTypeId=hit_type_response["HITTypeId"],
+            Question=q,
+            MaxAssignments=num_assignments,
+            LifetimeInSeconds=1209600,  # 1209600 seconds = 2 weeks is max allowed by MTurk
+            UniqueRequestToken=request.unique_request_token,
+            RequesterAnnotation=request.unique_request_token,
+        )["HIT"]
+        hit = boto3_models.HIT.model_validate(obj=hit_info)
+
+        # Idempotency error: botocore.errorfactory.RequestError: An error occurred (RequestError) when calling the CreateHIT operation: The HIT with ID "{hit_id}" already exists.
+        return CreateHitResponse(hit_id=hit.HITId)
+
+    def send_bonus_payment(
+        self,
+        request: SendBonusPaymentRequest,
+    ) -> None:
+        unique_request_token = f"{request.worker_id}:{request.assignment_id}"
+
+        try:
+            self.boto3_client.send_bonus(
+                WorkerId=request.worker_id,
+                BonusAmount=str(request.amount_usd),
+                AssignmentId=request.assignment_id,
+                Reason="Assignment-based bonus.",
+                UniqueRequestToken=unique_request_token,  # For idempotency
+            )
+        except Exception as e:
+            if "idempotency" in str(e):
+                # Expected error for duplicate request
+                pass
+
+    def approve_assignment(self, assignment_id: str) -> None:
+        self.boto3_client.approve_assignment(
+            AssignmentId=assignment_id,
+            OverrideRejection=False,
+        )
+
+    def iter_assignments(
+        self,
+        hit_id: str,
+    ) -> Iterable[ListAssignmentsItem]:
+        AssignmentStatuses = ["Submitted", "Approved", "Rejected"]
+
+        request_kwargs = dict(
+            HITId=hit_id,
+            MaxResults=100,
+            AssignmentStatuses=AssignmentStatuses,
+        )
+
+        # Paginate over boto3 results:
+        next_token = ""
+        while next_token is not None:
+            if next_token != "":
+                request_kwargs["NextToken"] = next_token
+            else:
+                if "NextToken" in request_kwargs:
+                    del request_kwargs["NextToken"]
+
+            call_return = self.boto3_client.list_assignments_for_hit(**request_kwargs)
+            for asn_info in call_return["Assignments"]:
+                assignment = boto3_models.Assignment.model_validate(obj=asn_info)
+                item = ListAssignmentsItem(
+                    hit_id=assignment.HITId,
+                    worker_id=assignment.WorkerId,
+                    assignment_id=assignment.AssignmentId,
+                    status=assignment.AssignmentStatus,
+                    submission_payload=extract_trace(assignment.Answer),
+                )
+
+                yield item
+
+            if "NextToken" in call_return:
+                next_token = call_return["NextToken"]
+            else:
+                # Will break
+                next_token = None
+
+    def cleanup_hit(self, hit_id: str) -> None:
+        # First, retrieve the HIT to ensure it exists
+        hit_info = self.boto3_client.get_hit(HITId=hit_id)
+        hit = boto3_models.HIT.model_validate(obj=hit_info["HIT"])
+
+        # See if this HIT has any QualificationRequirements
+        qual_reqs = hit.QualificationRequirements
+        for qual_req in qual_reqs:
+            # Get workers associated with this qualification:
+            worker_ids = self.list_workers_with_qualification_type(
+                qual_type_id=qual_req.QualificationTypeId
+            )
+            # Dissociate any qualifications from workers that were previously granted
+            for worker_id in worker_ids:
+                self.boto3_client.disassociate_qualification_from_worker(
+                    WorkerId=worker_id,
+                    QualificationTypeId=qual_req.QualificationTypeId,
+                )
+
+            # Delete the qualification type
+            self.delete_qualification_type(
+                qualification_type_id=qual_req.QualificationTypeId
+            )
+
+        # Update the expiration for the HIT to *now*
+        self.boto3_client.update_expiration_for_hit(
+            HITId=hit_id, ExpireAt=datetime.datetime.now(tz=datetime.timezone.utc)
+        )
+
+    # %% Quals:
+    def list_workers_with_qualification_type(
+        self,
+        qual_type_id: str,
+    ) -> List[str]:
+        next_token = ""
+        request_kwargs = dict(
+            QualificationTypeId=qual_type_id,
+            Status="Granted",
+            MaxResults=100,
+        )
+        worker_ids = []
+        while next_token is not None:
+            if next_token != "":
+                request_kwargs["NextToken"] = next_token
+            else:
+                if "NextToken" in request_kwargs:
+                    del request_kwargs["NextToken"]
+
+            call_return = self.boto3_client.list_workers_with_qualification_type(
+                **request_kwargs
+            )
+            for worker_info in call_return["Qualifications"]:
+                worker_ids.append(worker_info["WorkerId"])
+
+            if "NextToken" in call_return:
+                next_token = call_return["NextToken"]
+            else:
+                # Will break
+                next_token = None
+        return worker_ids
+
+    def create_qualification_type(
+        self,
+        unique_name: str,
+    ) -> boto3_models.QualificationType:
+        response = self.boto3_client.create_qualification_type(
+            Name=unique_name,
+            Description=unique_name,
+            QualificationTypeStatus="Active",
+        )
+        # Validate response:
+        return boto3_models.QualificationType.model_validate(
+            obj=response["QualificationType"]
+        )
+
+    def package_qualification_exists_requirement(
+        self,
+        qual_type: boto3_models.QualificationType,
+    ) -> boto3_models.QualificationRequirement:
+        return boto3_models.QualificationRequirement(
+            QualificationTypeId=qual_type.QualificationTypeId,
+            Comparator="Exists",
+            ActionsGuarded="DiscoverPreviewAndAccept",
+        )
+
+    def list_qualification_types(self) -> List[boto3_models.QualificationType]:
+        qualification_types = []
+
+        next_token = ""
+
+        call_kwargs = {}
+        while next_token is not None:
+            res = self.boto3_client.list_qualification_types(
+                MustBeRequestable=False,
+                MustBeOwnedByCaller=True,
+                MaxResults=100,
+                **call_kwargs,
+            )
+
+            if "NextToken" in res:
+                next_token = res["NextToken"]
+            else:
+                next_token = None
+
+            call_kwargs["NextToken"] = next_token
+
+            qreturn = res["QualificationTypes"]
+            for q in qreturn:
+                qual_type = boto3_models.QualificationType.model_validate(obj=q)
+                qualification_types.append(qual_type)
+
+        return qualification_types
+
+    def grant_qualification(
+        self,
+        qualification_type_id: str,
+        worker_id: str,
+        integer_value: int = 1,
+    ):
+        self.boto3_client.associate_qualification_with_worker(
+            QualificationTypeId=qualification_type_id,
+            WorkerId=worker_id,
+            IntegerValue=integer_value,
+            SendNotification=False,
+        )
+
+    def revoke_qualification(
+        self,
+        worker_id: str,
+        qualification_type_id: str,
+    ):
+        try:
+            self.boto3_client.disassociate_qualification_from_worker(
+                QualificationTypeId=qualification_type_id,
+                WorkerId=worker_id,
+                Reason="",
+            )
+        except Exception as e:
+            message = str(e)
+            if "RequestError" in message:
+                print(message)  # todo; already revoked?
+            else:
+                raise e
+
+    def delete_qualification_type(
+        self,
+        qualification_type_id: str,
+    ):
+        self.boto3_client.delete_qualification_type(
+            QualificationTypeId=qualification_type_id
+        )

nodekit/experimental/recruitment_services/mechanical_turk/models.py

@@ -0,0 +1,116 @@
+import datetime
+from decimal import Decimal
+from typing import List, Optional, Literal, Dict, Any
+
+import pydantic
+
+
+# %%
+class HIT(pydantic.BaseModel):
+    # https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/mturk/client/create_hit.html
+    HITId: str
+    HITTypeId: Optional[str]
+    HITGroupId: Optional[str] = None
+    CreationTime: datetime.datetime
+    Title: str
+    Description: str
+    Question: str
+    Keywords: str
+    HITStatus: Literal[
+        "Assignable", "Unassignable", "Reviewable", "Reviewing", "Disposed"
+    ]
+    MaxAssignments: int
+    Reward: Decimal = pydantic.Field(decimal_places=2)
+    AutoApprovalDelayInSeconds: int
+    Expiration: datetime.datetime
+    AssignmentDurationInSeconds: int
+    RequesterAnnotation: Optional[str] = None
+    QualificationRequirements: Optional[List["QualificationRequirement"]] = None
+    HITReviewStatus: Literal[
+        "NotReviewed", "MarkedForReview", "ReviewedAppropriate", "ReviewedInappropriate"
+    ]
+    NumberOfAssignmentsPending: int
+    NumberOfAssignmentsAvailable: int
+    NumberOfAssignmentsCompleted: int
+
+
+class Assignment(pydantic.BaseModel):
+    # https://boto3.amazonaws.com/v1/documentation/api/1.26.93/reference/services/mturk/client/list_assignments_for_hit.html
+    AssignmentId: str
+    WorkerId: str
+    HITId: str
+    AssignmentStatus: Literal["Submitted", "Approved", "Rejected"]
+    AutoApprovalTime: Optional[datetime.datetime] = None
+    AcceptTime: Optional[datetime.datetime] = None
+    SubmitTime: Optional[datetime.datetime] = None
+    ApprovalTime: Optional[datetime.datetime] = None
+    RejectionTime: Optional[datetime.datetime] = None
+    Deadline: Optional[datetime.datetime] = None
+    Answer: str
+    RequesterFeedback: Optional[str] = None
+
+
+class BonusPayment(pydantic.BaseModel):
+    # https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/mturk/client/list_bonus_payments.html
+    WorkerId: str
+    AssignmentId: str
+    BonusAmount: Decimal = pydantic.Field(decimal_places=2)
+    Reason: Optional[str] = None
+    GrantTime: Optional[datetime.datetime] = None
+
+
+# %% Qualifications
+class QualificationType(pydantic.BaseModel):
+    # https://docs.aws.amazon.com/AWSMechTurk/latest/AWSMturkAPI/ApiReference_QualificationTypeDataStructureArticle.html
+    QualificationTypeId: str
+    CreationTime: Optional[datetime.datetime] = None
+    Name: Optional[str] = None
+    Description: Optional[str] = None
+    Keywords: Optional[str] = None
+    QualificationTypeStatus: Literal["Active", "Inactive"]
+    RetryDelayInSeconds: Optional[int] = None
+    Test: Optional[str] = None
+    TestDurationInSeconds: Optional[int] = None
+    AnswerKey: Optional[str] = None
+    AutoGranted: Optional[bool] = None
+    AutoGrantedValue: Optional[int] = 1
+    IsRequestable: Optional[bool] = None
+
+
+class QualificationRequirement(pydantic.BaseModel):
+    QualificationTypeId: str
+
+    Comparator: Literal[
+        "LessThan",
+        "LessThanOrEqualTo",
+        "GreaterThan",
+        "GreaterThanOrEqualTo",
+        "EqualTo",
+        "NotEqualTo",
+        "Exists",
+        "DoesNotExist",
+        "In",
+        "NotIn",
+    ]
+    ActionsGuarded: Literal["Accept", "PreviewAndAccept", "DiscoverPreviewAndAccept"]
+    IntegerValues: List[int] | None = pydantic.Field(default=None)
+    LocaleValues: List["LocaleValue"] | None = pydantic.Field(default=None)
+
+    @pydantic.model_serializer()
+    def ensure_exclude_none(self):
+        # boto3 expects IntegerValues and LocaleValues to not be present as fields if they are None:
+        base: Dict[str, Any] = dict(
+            QualificationTypeId=self.QualificationTypeId,
+            Comparator=self.Comparator,
+            ActionsGuarded=self.ActionsGuarded,
+        )
+        if self.IntegerValues is not None:
+            base["IntegerValues"] = self.IntegerValues
+        if self.LocaleValues is not None:
+            base["LocaleValues"] = self.LocaleValues
+        return base
+
+
+class LocaleValue(pydantic.BaseModel):
+    Country: str
+    Subdivision: Optional[str]