PLT-1016 - Add methods for External Workforce management

docs/labelbox/external-workforce.rst

@@ -0,0 +1,6 @@
+ExternalWorkforce
+===============================================================================================
+
+.. automodule:: labelbox.schema.external_workforce
+   :members:
+   :show-inheritance:

libs/labelbox/src/labelbox/__init__.py

@@ -41,3 +41,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.external_workforce import ExternalWorkforce

libs/labelbox/src/labelbox/schema/__init__.py

@@ -25,4 +25,5 @@
 import labelbox.schema.identifiable
 import labelbox.schema.catalog
 import labelbox.schema.ontology_kind
+import labelbox.schema.external_workforce
 import labelbox.schema.project_overview

libs/labelbox/src/labelbox/schema/external_workforce.py

@@ -0,0 +1,12 @@
+from labelbox.pydantic_compat import BaseModel
+
+class ExternalWorkforce(BaseModel):
+    """
+    Represents an external workforce used in the Labelbox system.
+
+    Attributes:
+        id (str): The unique identifier of the external workforce.
+        name (str): The name of the external workforce.
+    """
+    id: str
+    name: str

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

@@ -27,6 +27,7 @@
 from labelbox.schema.export_filters import ProjectExportFilters, validate_datetime, build_filters
 from labelbox.schema.export_params import ProjectExportParams
 from labelbox.schema.export_task import ExportTask
+from labelbox.schema.external_workforce import ExternalWorkforce
 from labelbox.schema.id_type import IdType
 from labelbox.schema.identifiable import DataRowIdentifier, GlobalKey, UniqueId
 from labelbox.schema.identifiables import DataRowIdentifiers, UniqueIds
@@ -40,6 +41,7 @@
 from labelbox.schema.ontology_kind import (EditorTaskType, OntologyKind)
 from labelbox.schema.project_overview import ProjectOverview, ProjectOverviewDetailed
 
+
 if TYPE_CHECKING:
     from labelbox import BulkImportRequest
 
@@ -1806,7 +1808,7 @@ def get_overview(
 
         """
         query = """query ProjectGetOverviewPyApi($projectId: ID!) {
-            project(where: { id: $projectId }) {      
+            project(where: { id: $projectId }) {
             workstreamStateCounts {
                 state
                 count
@@ -1878,6 +1880,115 @@ def clone(self) -> "Project":
         result = self.client.execute(mutation, {"projectId": self.uid})
         return self.client.get_project(result["cloneProject"]["id"])
 
+    def add_external_workforce(self, workforce_id: Union[str, ExternalWorkforce]) -> List[ExternalWorkforce]:
+
+        """
+        Add an external workforce (organization) to a project.
+
+        Args:
+            workforce_id (Union[str, ExternalWorkforce]): Organization id of the external workforce.
+
+        Returns:
+            List[ExternalWorkforce]: The remaining external workforces or an empty list if no external workforces are left.
+
+        Raises:
+            LabelboxError: If the external workforce with the given ID cannot be found.
+
+        Note:
+            This method adds an external workforce (organization) to the current project.
+            The `workforce_id` parameter can be either a string representing the organization ID or an instance of the `ExternalWorkforce` class.
+        """
+        workforce_id = workforce_id.uid if isinstance(workforce_id, ExternalWorkforce) else workforce_id
+
+        mutation = """
+            mutation ShareProjectWithExternalOrganizationPyApi($projectId: ID!, $organizationId: ID!) {
+                shareProjectWithExternalOrganization(
+                    data: { projectId: $projectId, organizationId: $organizationId }
+                ) {
+                    id
+                    sharedWithOrganizations {
+                        id
+                        name
+                    }
+                    }
+            }
+        """
+
+        result = self.client.execute(mutation, {"projectId": self.uid, "organizationId": workforce_id})
+
+        if not result:
+            raise ResourceNotFoundError(ExternalWorkforce, {"id": workforce_id})
+
+        return [ExternalWorkforce(**workforce_dic)
+                    for workforce_dic
+                    in result["shareProjectWithExternalOrganization"]["sharedWithOrganizations"]]
+
+
+    def remove_external_workforce(self, workforce_id: Union[str, ExternalWorkforce]) -> List[ExternalWorkforce]:
+        """Remove an external workforce (organization) from a project.
+
+        Args:
+            workforce_id (Union[str, ExternalWorkforce]): Organization id of the external workforce
+            or an instance of the ExternalWorkforce class.
+
+        Returns:
+            List[ExternalWorkforce]: The remaining external workforces or an empty list if no external workforces are left.
+
+        Raises:
+            LabelboxError: If the external workforce cannot be found.
+        """
+        workforce_id = workforce_id.uid if isinstance(workforce_id, ExternalWorkforce) else workforce_id
+
+        mutation = """
+            mutation UnshareProjectWithExternalOrganizationPyApi($projectId: ID!, $organizationId: ID!) {
+                unshareProjectWithExternalOrganization(
+                    data: { projectId: $projectId, organizationId: $organizationId }
+                ) {
+                    id
+                    sharedWithOrganizations {
+                        id
+                        name
+                    }
+                }
+            }
+            """
+
+        result = self.client.execute(mutation, {"projectId": self.uid, "organizationId": workforce_id})
+
+        if not result:
+            raise ResourceNotFoundError(ExternalWorkforce, {"id": workforce_id})
+
+        return [ExternalWorkforce(**workforce_dic)
+                for workforce_dic
+                in result["unshareProjectWithExternalOrganization"]["sharedWithOrganizations"]]
+
+
+    def get_external_workforces(self) -> List[ExternalWorkforce]:
+        """List the external workforces (organizations) attached to a project
+
+        Args:
+            project_id: Id of the project to check Project to check
+
+        Returns:
+            list of dictionaries with id and name for each external workforce (aka organization)
+            or empty list if no external workforces were found
+        """
+
+        query = """
+            query GetProjectExternalOganizationsPyApi($projectId: ID!) {
+                project(where: { id: $projectId }) {
+                    id
+                    sharedWithOrganizations {
+                        id
+                        name
+                    }
+                }
+            }
+            """
+
+        result = self.client.execute(query, {"projectId": self.uid})["project"]["sharedWithOrganizations"]
+        return [ExternalWorkforce(**workforce_dic) for workforce_dic in result]
+
 
 class ProjectMember(DbObject):
     user = Relationship.ToOne("User", cache=True)

libs/labelbox/tests/integration/test_external_workforce.py

@@ -0,0 +1,56 @@
+from labelbox import Project, ExternalWorkforce
+import os
+import sys
+
+def _get_workforce_id_for_test() -> str:
+    # For this test, we need an additional organization id to connect the project to.
+    # Given that the SDK is tested against multiple versions of Python, and different environments,
+    # we decided to choose the organization id from 
+    # https://labelbox.atlassian.net/wiki/spaces/PLT/pages/2110816271/How+to+labelbox-python+SDK+CI+Tests
+    #
+    # In particular, we have:
+    # clum46jsb00jp07z338dh1gvb: for Python 3.8 in staging
+    # cltp16p2v04dr07ywfgqxf23u: for Python 3.12 in staging    
+    # ckcz6bubudyfi0855o1dt1g9s: for Python 3.8 in production
+    # cltp1fral01dh07009zsng3zs: for Python 3.12 in production
+    #
+    # The idea is that when depending on the environment, if the current version of Python is 3.8,
+    # then we use the organization id for Python 3.12. This way, We always us a different organization id
+    # from the current one as an external workforce to avoid conflict.
+
+    # Note: A better approach would be to create a new organization dynamically for testing purpose.
+    # This is not currently possible.
+    current_version = sys.version_info[:2]
+
+    if os.environ['LABELBOX_TEST_ENVIRON'] == "staging":
+        return "cltp16p2v04dr07ywfgqxf23u" if current_version == (3, 8) else "clum46jsb00jp07z338dh1gvb"
+    else:
+        return "cltp1fral01dh07009zsng3zs" if current_version == (3, 8) else "ckcz6bubudyfi0855o1dt1g9s"
+
+
+def test_add_external_workforce(project: Project):
+    workforce_id = _get_workforce_id_for_test()
+
+    external_workforces = project.add_external_workforce(workforce_id)
+    assert len(external_workforces) == 1
+    assert isinstance(external_workforces[0], ExternalWorkforce)
+
+
+def test_get_external_workforces(project: Project):
+    workforce_id = _get_workforce_id_for_test()
+
+    external_workforces = project.add_external_workforce(workforce_id)
+
+    external_workforces = project.get_external_workforces()
+    assert len(external_workforces) == 1
+    assert isinstance(external_workforces[0], ExternalWorkforce)
+
+
+def test_remove_external_workforce(project: Project):
+    workforce_id = _get_workforce_id_for_test()
+    external_workforces = project.add_external_workforce(workforce_id)
+
+    external_workforces = project.remove_external_workforce(workforce_id)
+    assert len(external_workforces) == 0
+
+