Alignerr project creation

Author: midhun-pm

libs/labelbox/pyproject.toml

@@ -12,6 +12,7 @@ dependencies = [
     "tqdm>=4.66.2",
     "geojson>=3.1.0",
     "lbox-clients==1.1.2",
+    "PyYAML>=6.0",
 ]
 readme = "README.md"
 requires-python = ">=3.9,<3.14"

libs/labelbox/src/labelbox/__init__.py

@@ -78,6 +78,7 @@
 from labelbox.schema.ontology_kind import OntologyKind
 from labelbox.schema.organization import Organization
 from labelbox.schema.project import Project
+from labelbox.alignerr.schema.project_rate import ProjectRateV2 as ProjectRate
 from labelbox.schema.project_model_config import ProjectModelConfig
 from labelbox.schema.project_overview import (
     ProjectOverview,
@@ -98,7 +99,6 @@
     ResponseOption,
     PromptResponseClassification,
 )
-from lbox.exceptions import *
 from labelbox.schema.taskstatus import TaskStatus
 from labelbox.schema.api_key import ApiKey
 from labelbox.schema.timeunit import TimeUnit

libs/labelbox/src/labelbox/alignerr/__init__.py

@@ -0,0 +1,3 @@
+from .alignerr_project import AlignerrWorkspace
+
+__all__ = ['AlignerrWorkspace']

libs/labelbox/src/labelbox/alignerr/alignerr_project.py

@@ -0,0 +1,328 @@
+import datetime
+from enum import Enum
+from typing import TYPE_CHECKING, Optional
+import yaml
+from pathlib import Path
+
+import logging
+
+from labelbox.alignerr.schema.project_rate import BillingMode
+from labelbox.alignerr.schema.project_rate import ProjectRateInput
+from labelbox.alignerr.schema.project_rate import ProjectRateV2
+from labelbox.schema.media_type import MediaType
+
+logger = logging.getLogger(__name__)
+
+
+if TYPE_CHECKING:
+    from labelbox import Client
+    from labelbox.schema.project import Project
+
+
+class AlignerrRole(Enum):
+    Labeler = "LABELER"
+    Reviewer = "REVIEWER"
+    Admin = "ADMIN"
+
+
+class AlignerrProject:
+    def __init__(self, client: "Client", project: "Project"):
+        self.client = client
+        self.project = project
+
+    @property
+    def project(self) -> Optional["Project"]:
+        return self._project
+
+    @project.setter
+    def project(self, project: "Project"):
+        self._project = project
+
+    def get_project_rate(self) -> Optional["ProjectRateV2"]:
+        query_str = """
+        query GetAllProjectRatesPyApi($projectId: ID!) {
+            project(where: { id: $projectId }) {
+                id
+                ratesV2 {
+                    id
+                    userRole {
+                        id
+                        name
+                    }
+                    isBillRate
+                    billingMode
+                    rate
+                    effectiveSince
+                    effectiveUntil
+                    createdAt
+                    updatedAt
+                    updatedBy {
+                        id
+                        email
+                        name
+                    }
+                }
+            }
+        }
+        """
+        result = self.client.execute(query_str, {"projectId": self.project.uid})
+        rates_data = result["project"]["ratesV2"]
+
+        if not rates_data:
+            return None
+
+        # Return the first rate as a ProjectRateV2 object
+        return ProjectRateV2(self.client, rates_data[0])
+
+    def set_project_rate(self, project_rate_input):
+        mutation_str = """mutation SetProjectRateV2PyApi($input: SetProjectRateV2Input!) {
+                setProjectRateV2(input: $input) {
+                    success
+                }
+        }"""
+
+        params = {
+            "projectId": self.project.uid,
+            "input": {
+                "projectId": self.project.uid,
+                "userRoleId": project_rate_input.rateForId,
+                "isBillRate": project_rate_input.isBillRate,
+                "billingMode": project_rate_input.billingMode.value
+                if hasattr(project_rate_input.billingMode, "value")
+                else project_rate_input.billingMode,
+                "rate": project_rate_input.rate,
+                "effectiveSince": project_rate_input.effectiveSince,
+                "effectiveUntil": project_rate_input.effectiveUntil,
+            },
+        }
+
+        result = self.client.execute(mutation_str, params)
+
+        return result["setProjectRateV2"]["success"]
+
+
+class AlignerrProjectBuilder:
+    def __init__(self, client: "Client"):
+        self.client = client
+        self._alignerr_rates: dict[str, ProjectRateInput] = {}
+        self._customer_rate: ProjectRateInput = None
+        self.role_name_to_id = self._get_role_name_to_id()
+
+    def set_name(self, name: str):
+        self.project_name = name
+        return self
+
+    def set_media_type(self, media_type: "MediaType"):
+        self.project_media_type = media_type
+        return self
+
+    def set_alignerr_role_rate(
+        self,
+        *,
+        role_name: AlignerrRole,
+        rate: float,
+        billing_mode: BillingMode,
+        effective_since: datetime.datetime,
+        effective_until: Optional[datetime.datetime] = None,
+    ):
+        if role_name.value not in self.role_name_to_id:
+            raise ValueError(f"Role {role_name} not found")
+
+        role_id = self.role_name_to_id[role_name.value]
+        role_name = role_name.value
+
+        # Convert datetime objects to ISO format strings
+        effective_since_str = effective_since.isoformat() if isinstance(effective_since, datetime.datetime) else effective_since
+        effective_until_str = effective_until.isoformat() if isinstance(effective_until, datetime.datetime) else effective_until
+
+        self._alignerr_rates[role_name] = ProjectRateInput(
+            rateForId=role_id,
+            isBillRate=False,
+            billingMode=billing_mode,
+            rate=rate,
+            effectiveSince=effective_since_str,
+            effectiveUntil=effective_until_str,
+        )
+        return self
+
+    def set_customer_rate(
+        self,
+        *,
+        rate: float,
+        billing_mode: BillingMode,
+        effective_since: datetime.datetime,
+        effective_until: Optional[datetime.datetime] = None,
+    ):
+        # Convert datetime objects to ISO format strings
+        effective_since_str = effective_since.isoformat() if isinstance(effective_since, datetime.datetime) else effective_since
+        effective_until_str = effective_until.isoformat() if isinstance(effective_until, datetime.datetime) else effective_until
+
+        self._customer_rate = ProjectRateInput(
+            rateForId="",  # Empty string for customer rate
+            isBillRate=True,
+            billingMode=billing_mode,
+            rate=rate,
+            effectiveSince=effective_since_str,
+            effectiveUntil=effective_until_str,
+        )
+        return self
+
+    def create(self, skip_validation: bool = False):
+        if not skip_validation:
+            self._validate()
+        logger.info("Creating project")
+
+        project_data = {
+            "name": self.project_name,
+            "media_type": self.project_media_type,
+        }
+        labelbox_project = self.client.create_project(**project_data)
+        alignerr_project = AlignerrProject(self.client, labelbox_project)
+        for alignerr_role, project_rate in self._alignerr_rates.items():
+            logger.info(f"Setting project rate for {alignerr_role}")
+            alignerr_project.set_project_rate(project_rate)
+        return alignerr_project
+
+    def _validate_alignerr_rates(self):
+        required_role_rates = set([AlignerrRole.Labeler.value, AlignerrRole.Reviewer.value])
+
+        for role_name in self._alignerr_rates.keys():
+            required_role_rates.remove(role_name)
+        if len(required_role_rates) > 0:
+            raise ValueError(
+                f"Required role rates are not set: {required_role_rates}"
+            )
+
+    def _validate_customer_rate(self):
+        if self._customer_rate is None:
+            raise ValueError("Customer rate is not set")
+
+    def _validate(self):
+        self._validate_alignerr_rates()
+        self._validate_customer_rate()
+
+    def _get_role_name_to_id(self) -> dict[str, str]:
+        roles = self.client.get_roles()
+        return {role.name: role.uid for role in roles.values()}
+
+
+class AlignerrProjectFactory:
+    def __init__(self, client: "Client"):
+        self.client = client
+
+    def create(self, yaml_file_path: str, skip_validation: bool = False):
+        """
+        Create an AlignerrProject from a YAML configuration file.
+        
+        Args:
+            yaml_file_path: Path to the YAML configuration file
+            skip_validation: Whether to skip validation of required fields
+            
+        Returns:
+            AlignerrProject: The created project with configured rates
+            
+        Raises:
+            FileNotFoundError: If the YAML file doesn't exist
+            yaml.YAMLError: If the YAML file is invalid
+            ValueError: If required fields are missing or invalid
+        """
+        logger.info(f"Creating project from YAML file: {yaml_file_path}")
+        
+        # Load and parse YAML file
+        yaml_path = Path(yaml_file_path)
+        if not yaml_path.exists():
+            raise FileNotFoundError(f"YAML file not found: {yaml_file_path}")
+            
+        try:
+            with open(yaml_path, 'r') as file:
+                config = yaml.safe_load(file)
+        except yaml.YAMLError as e:
+            raise yaml.YAMLError(f"Invalid YAML file: {e}")
+        
+        # Validate required fields
+        if not config:
+            raise ValueError("YAML file is empty")
+            
+        required_fields = ['name', 'media_type']
+        for field in required_fields:
+            if field not in config:
+                raise ValueError(f"Required field '{field}' is missing from YAML configuration")
+        
+        # Create project builder
+        builder = AlignerrProjectBuilder(self.client)
+        
+        # Set basic project properties
+        builder.set_name(config['name'])
+        
+        # Set media type
+        media_type_str = config['media_type']
+        media_type = MediaType(media_type_str)
+        
+        # Check if the media type is supported
+        if not MediaType.is_supported(media_type):
+            supported_members = MediaType.get_supported_members()
+            raise ValueError(f"Invalid media_type '{media_type_str}'. Must be one of: {supported_members}")
+        
+        builder.set_media_type(media_type)
+        
+        # Set project rates if provided
+        if 'rates' in config:
+            rates_config = config['rates']
+            if not isinstance(rates_config, dict):
+                raise ValueError("'rates' must be a dictionary")
+                
+            for role_name, rate_config in rates_config.items():
+                try:
+                    alignerr_role = AlignerrRole(role_name.upper())
+                except ValueError:
+                    raise ValueError(f"Invalid role '{role_name}'. Must be one of: {[r.value for r in AlignerrRole]}")
+                
+                # Validate rate configuration
+                required_rate_fields = ['rate', 'billing_mode', 'effective_since']
+                for field in required_rate_fields:
+                    if field not in rate_config:
+                        raise ValueError(f"Required field '{field}' is missing for role '{role_name}'")
+                
+                # Parse billing mode
+                try:
+                    billing_mode = BillingMode(rate_config['billing_mode'])
+                except ValueError:
+                    raise ValueError(f"Invalid billing_mode '{rate_config['billing_mode']}' for role '{role_name}'. Must be one of: {[e.value for e in BillingMode]}")
+                
+                # Parse effective dates
+                try:
+                    effective_since = datetime.datetime.fromisoformat(rate_config['effective_since'])
+                except ValueError:
+                    raise ValueError(f"Invalid effective_since date format for role '{role_name}'. Use ISO format (YYYY-MM-DDTHH:MM:SS)")
+                
+                effective_until = None
+                if 'effective_until' in rate_config and rate_config['effective_until']:
+                    try:
+                        effective_until = datetime.datetime.fromisoformat(rate_config['effective_until'])
+                    except ValueError:
+                        raise ValueError(f"Invalid effective_until date format for role '{role_name}'. Use ISO format (YYYY-MM-DDTHH:MM:SS)")
+                
+                # Set the rate
+                builder.set_alignerr_role_rate(
+                    role_name=alignerr_role,
+                    rate=float(rate_config['rate']),
+                    billing_mode=billing_mode,
+                    effective_since=effective_since,
+                    effective_until=effective_until
+                )
+        
+        # Create the project
+        return builder.create(skip_validation=skip_validation)
+
+
+class AlignerrWorkspace:
+    def __init__(self, client: "Client"):
+        self.client = client
+
+    def project_builder(self):
+        return AlignerrProjectBuilder(self.client)
+
+    def project_prototype(self):
+        return AlignerrProjectFactory(self.client)
+
+