[PLT-1305] Vb/start labeling service plt 1305 (#1749)

Author: vbrodsky

docs/labelbox/labeling-service.rst

@@ -0,0 +1,6 @@
+Project Labeling Service
+===============================================================================================
+
+.. automodule:: labelbox.schema.labeling_service
+   :members:
+   :show-inheritance:

libs/labelbox/src/labelbox/__init__.py

@@ -43,3 +43,4 @@
 from labelbox.schema.identifiable import UniqueId, GlobalKey
 from labelbox.schema.ontology_kind import OntologyKind
 from labelbox.schema.project_overview import ProjectOverview, ProjectOverviewDetailed
+from labelbox.schema.labeling_service import LabelingService, LabelingServiceStatus

libs/labelbox/src/labelbox/schema/labeling_service.py

@@ -0,0 +1,86 @@
+from datetime import datetime
+from enum import Enum
+from typing import Any
+
+from labelbox.exceptions import ResourceNotFoundError
+
+from labelbox.data.annotation_types.types import Cuid
+from labelbox.pydantic_compat import BaseModel
+from labelbox.utils import _CamelCaseMixin
+
+
+class LabelingServiceStatus(Enum):
+    Accepted = 'ACCEPTED',
+    Calibration = 'CALIBRATION',
+    Complete = 'COMPLETE',
+    Production = 'PRODUCTION',
+    Requested = 'REQUESTED',
+    SetUp = 'SET_UP'
+
+
+class LabelingService(BaseModel):
+    id: Cuid
+    project_id: Cuid
+    created_at: datetime
+    updated_at: datetime
+    created_by_id: Cuid
+    status: LabelingServiceStatus
+    client: Any  # type Any to avoid circular import from client
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        if not self.client.enable_experimental:
+            raise RuntimeError(
+                "Please enable experimental in client to use LabelingService")
+
+    class Config(_CamelCaseMixin.Config):
+        ...
+
+    @classmethod
+    def start(cls, client, project_id: Cuid) -> 'LabelingService':
+        """
+        Starts the labeling service for the project. This is equivalent to a UI acction to Request Specialized Labelers
+
+        Returns:
+            LabelingService: The labeling service for the project.
+        Raises:
+            Exception: If the service fails to start.
+        """
+        query_str = """mutation CreateProjectBoostWorkforcePyApi($projectId: ID!) {
+          upsertProjectBoostWorkforce(data: { projectId: $projectId }) {
+            success
+          }
+        }"""
+        result = client.execute(query_str, {"projectId": project_id})
+        success = result["upsertProjectBoostWorkforce"]["success"]
+        if not success:
+            raise Exception("Failed to start labeling service")
+        return cls.get(client, project_id)
+
+    @classmethod
+    def get(cls, client, project_id: Cuid) -> 'LabelingService':
+        """
+        Returns the labeling service associated with the project.
+
+        Raises:
+            ResourceNotFoundError: If the project does not have a labeling service.
+        """
+        query = """
+            query GetProjectBoostWorkforcePyApi($projectId: ID!) {
+            projectBoostWorkforce(data: { projectId: $projectId }) {
+                    id
+                    projectId
+                    createdAt
+                    updatedAt
+                    createdById
+                    status
+                }
+            }
+        """
+        result = client.execute(query, {"projectId": project_id})
+        if result["projectBoostWorkforce"] is None:
+            raise ResourceNotFoundError(
+                message="The project does not have a labeling service.")
+        data = result["projectBoostWorkforce"]
+        data["client"] = client
+        return LabelingService(**data)

libs/labelbox/src/labelbox/schema/project.py

@@ -9,6 +9,7 @@
 from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Tuple, TypeVar, Union, overload
 from urllib.parse import urlparse
 
+from labelbox.schema.labeling_service import LabelingService
 import requests
 
 from labelbox import parser
@@ -149,7 +150,7 @@ def is_chat_evaluation(self) -> bool:
             True if this project is a live chat evaluation project, False otherwise
         """
         return self.media_type == MediaType.Conversational and self.editor_task_type == EditorTaskType.ModelChatEvaluation
-    
+
     def is_prompt_response(self) -> bool:
         """
         Returns:
@@ -1915,6 +1916,24 @@ def clone(self) -> "Project":
         result = self.client.execute(mutation, {"projectId": self.uid})
         return self.client.get_project(result["cloneProject"]["id"])
 
+    @experimental
+    def get_labeling_service(self) -> LabelingService:
+        """Get the labeling service for this project.
+
+        Returns:
+            LabelingService: The labeling service for this project.
+        """
+        return LabelingService.get(self.client, self.uid)  # type: ignore
+
+    @experimental
+    def request_labeling_service(self) -> LabelingService:
+        """Get the labeling service for this project.
+
+        Returns:
+            LabelingService: The labeling service for this project.
+        """
+        return LabelingService.start(self.client, self.uid)  # type: ignore
+
 
 class ProjectMember(DbObject):
     user = Relationship.ToOne("User", cache=True)

libs/labelbox/tests/integration/test_labeling_service.py

@@ -0,0 +1,20 @@
+import pytest
+
+from labelbox.exceptions import ResourceNotFoundError
+from labelbox.schema.labeling_service import LabelingServiceStatus
+
+
+def test_get_labeling_service_throws_exception(project):
+    with pytest.raises(ResourceNotFoundError):  # No labeling service by default
+        project.get_labeling_service()
+
+
+def test_start_labeling_service(project):
+    labeling_service = project.request_labeling_service()
+    assert labeling_service.status == LabelingServiceStatus.SetUp
+    assert labeling_service.project_id == project.uid
+
+    # Check that the labeling service is now available
+    labeling_service = project.get_labeling_service()
+    assert labeling_service.status == LabelingServiceStatus.SetUp
+    assert labeling_service.project_id == project.uid