Add tasks to organization and TaskStatus to filter tasks

Author: paulnoirel

libs/labelbox/src/labelbox/__init__.py

@@ -98,3 +98,4 @@
     PromptResponseClassification,
 )
 from lbox.exceptions import *
+from labelbox.schema.taskstatus import TaskStatus

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

@@ -26,3 +26,4 @@
 import labelbox.schema.catalog
 import labelbox.schema.ontology_kind
 import labelbox.schema.project_overview
+import labelbox.schema.taskstatus

libs/labelbox/src/labelbox/schema/organization.py

@@ -52,6 +52,7 @@ def __init__(self, *args, **kwargs):
     projects = Relationship.ToMany("Project", True)
     webhooks = Relationship.ToMany("Webhook", False)
     resource_tags = Relationship.ToMany("ResourceTags", False)
+    tasks = Relationship.ToMany("Task", False, "tasks")
 
     def invite_user(
         self,

libs/labelbox/src/labelbox/schema/task.py

@@ -14,6 +14,7 @@
 from labelbox.schema.internal.datarow_upload_constants import (
     DOWNLOAD_RESULT_PAGE_SIZE,
 )
+from labelbox.schema.taskstatus import TaskStatus
 
 if TYPE_CHECKING:
     from labelbox import User
@@ -45,6 +46,9 @@ class Task(DbObject):
     created_at = Field.DateTime("created_at")
     name = Field.String("name")
     status = Field.String("status")
+    status_type = Field.Enum(
+        TaskStatus, "status_type", "status"
+    )  # additional status for filtering
     completion_percentage = Field.Float("completion_percentage")
     result_url = Field.String("result_url", "result")
     errors_url = Field.String("errors_url", "errors")

libs/labelbox/src/labelbox/schema/taskstatus.py

@@ -0,0 +1,25 @@
+from enum import Enum
+
+
+class TaskStatus(str, Enum):
+    In_Progress = "IN_PROGRESS"
+    Complete = "COMPLETE"
+    Canceling = "CANCELLING"
+    Canceled = "CANCELED"
+    Failed = "FAILED"
+    Unknown = "UNKNOWN"
+
+    @classmethod
+    def _missing_(cls, value):
+        """Handle missing or unknown task status values.
+
+        If a task status value is not found in the enum, this method returns
+        the Unknown status instead of raising an error.
+
+        Args:
+            value: The status value that doesn't match any enum member
+
+        Returns:
+            TaskStatus.Unknown: The default status for unrecognized values
+        """
+        return cls.Unknown

libs/labelbox/tests/data/export/streamable/test_export_data_rows_streamable.py

@@ -1,6 +1,6 @@
 import time
 
-from labelbox import DataRow, ExportTask, StreamType
+from labelbox import DataRow, ExportTask, StreamType, Task, TaskStatus
 
 
 class TestExportDataRow:
@@ -135,3 +135,33 @@ def test_cancel_export_task(
         # Verify the task was cancelled
         cancelled_task = client.get_task_by_id(export_task.uid)
         assert cancelled_task.status in ["CANCELING", "CANCELED"]
+
+    def test_task_filter(self, client, data_row, wait_for_data_row_processing):
+        organization = client.get_organization()
+        user = client.get_user()
+
+        export_task = DataRow.export(
+            client=client,
+            data_rows=[data_row],
+            task_name="TestExportDataRow:test_task_filter",
+        )
+
+        # Check if task is listed "in progress" in organization's tasks
+        org_tasks_in_progress = organization.tasks(
+            where=Task.status_type == TaskStatus.In_Progress
+        )
+        retrieved_task_in_progress = next(
+            (t for t in org_tasks_in_progress if t.uid == export_task.uid), ""
+        )
+        assert getattr(retrieved_task_in_progress, "uid", "") == export_task.uid
+
+        export_task.wait_till_done()
+
+        # Check if task is listed "complete" in user's created tasks
+        user_tasks_complete = user.created_tasks(
+            where=Task.status_type == TaskStatus.Complete
+        )
+        retrieved_task_complete = next(
+            (t for t in user_tasks_complete if t.uid == export_task.uid), ""
+        )
+        assert getattr(retrieved_task_complete, "uid", "") == export_task.uid