Add LabelingService start()

Author: Val Brodsky

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

@@ -1,9 +1,10 @@
 from datetime import datetime
 from enum import Enum
+from typing import Any
+
+from labelbox.exceptions import ResourceNotFoundError
 
-from ..client import Client
 from labelbox.data.annotation_types.types import Cuid
-from labelbox.orm.db_object import experimental
 from labelbox.pydantic_compat import BaseModel
 from labelbox.utils import _CamelCaseMixin
 
@@ -17,43 +18,69 @@ class LabelingServiceStatus(Enum):
     SetUp = 'SET_UP'
 
 
-@experimental
-class LabelingService(_CamelCaseMixin, BaseModel):
+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: Client, project_id: Cuid) -> 'LabelingService':
-        return cls._create(client=client, project_id=project_id)
-
-"""
-mutation CreateProjectBoostWorkforce($projectId: ID!) {
-  upsertProjectBoostWorkforce(data: { projectId: $projectId }) {
-    success
-    __typename
-  }
-}
-{
-  "projectId": "clz0b7jg901fh07zic3u67b7g"
-}
-
-
-{
-  "data": {
-    "upsertProjectBoostWorkforce": {
-      "success": true,
-      "__typename": "ProjectBoostWorkforceResult"
-    }
-  }
-}
-"""
+    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 _create(cls, client: Client, project_id: Cuid) -> 'LabelingService':
-        ...
-    
-    def status_as_string(self):
-        return self.status.value
+    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

@@ -150,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:
@@ -1918,32 +1918,13 @@ def clone(self) -> "Project":
 
     @experimental
     def get_labeling_service(self) -> LabelingService:
-        """
-        Returns the labeling service associated with the project.
+        """Get the labeling service for this project.
 
         Returns:
-            LabelingService: The labeling service associated with the project.
-
-        Raises:
-            ResourceNotFoundError: If the project does not have a labeling service.
+            LabelingService: The labeling service for this project.
         """
-        query = """
-            query GetProjectBoostWorkforcePyApi($projectId: ID!) {
-            projectBoostWorkforce(data: { projectId: $projectId }) {
-                    id
-                    projectId
-                    createdAt
-                    updatedAt
-                    createdById
-                    status
-                }
-            }
-        """
-        result = self.client.execute(query, {"projectId": self.uid})
-        if result["projectBoostWorkforce"] is None:
-            raise ResourceNotFoundError(
-                message="The project does not have a labeling service.")
-        return LabelingService(**result["projectBoostWorkforce"])
+        return LabelingService.get(self.client, self.uid)  # type: ignore
+
 
 class ProjectMember(DbObject):
     user = Relationship.ToOne("User", cache=True)

libs/labelbox/tests/integration/test_labeling_service.py

@@ -1,7 +1,20 @@
-from labelbox.exceptions import ResourceNotFoundError
 import pytest
 
+from labelbox.exceptions import ResourceNotFoundError
+from labelbox.schema.labeling_service import LabelingService, 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 = LabelingService.start(project.client, project.uid)
+    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