Adds task + label & multi label mutations (#93)

Co-authored-by: JWittmeyer <[email protected]>

Author: JWittmeyer

controller/labeling_task_label/manager.py

@@ -16,6 +16,26 @@
 from typing import List, Any, Dict
 import re
 
+LABEL_COLORS = [
+    "red",
+    "orange",
+    "amber",
+    "yellow",
+    "lime",
+    "green",
+    "emerald",
+    "teal",
+    "cyan",
+    "sky",
+    "blue",
+    "indigo",
+    "violet",
+    "purple",
+    "fuchsia",
+    "pink",
+    "rose",
+]
+
 
 def get_label(project_id: str, label_id: str) -> LabelingTaskLabel:
     return labeling_task_label.get(project_id, label_id)
@@ -73,6 +93,25 @@ def create_label(
     return label
 
 
+def create_labels(
+    project_id: str, labeling_task_id: str, names: List[str]
+) -> List[LabelingTaskLabel]:
+    task = labeling_task.get(project_id, labeling_task_id)
+    labels = []
+
+    available_colors = len(LABEL_COLORS)
+    for idx, name in enumerate(names):
+        labels.append(
+            labeling_task_label.create(
+                project_id, name, labeling_task_id, LABEL_COLORS[idx % available_colors]
+            )
+        )
+        if task.task_type == LabelingTaskType.INFORMATION_EXTRACTION.value:
+            create_knowledge_base_if_not_existing(name, project_id)
+    general.commit()
+    return labels
+
+
 def delete_label(project_id: str, label_id: str) -> None:
     labeling_task_label.delete(project_id, label_id, with_commit=True)
 

graphql_api/mutation/labeling_task.py

@@ -1,7 +1,8 @@
-from typing import Optional
+from typing import Optional, List
 
 from controller.auth import manager as auth
 from controller.labeling_task import manager
+from controller.labeling_task_label import manager as label_manager
 from controller.project import manager as project_manager
 from submodules.model import events
 from util import doc_ock, notification
@@ -46,6 +47,51 @@ def mutate(
         return CreateLabelingTask(ok=True)
 
 
+class CreateTaskAndLabels(graphene.Mutation):
+    class Arguments:
+        project_id = graphene.ID(required=True)
+        labeling_task_name = graphene.String(required=True)
+        labeling_task_type = graphene.String(required=True)
+        labeling_task_target_id = graphene.ID(required=False)
+        labels = graphene.List(graphene.String, required=False)
+
+    ok = graphene.Boolean()
+    task_id = graphene.ID()
+
+    def mutate(
+        self,
+        info,
+        project_id: str,
+        labeling_task_name: str,
+        labeling_task_type: str,
+        labeling_task_target_id: Optional[str] = None,
+        labels: Optional[List[str]] = None,
+    ):
+        auth.check_demo_access(info)
+        auth.check_project_access(info, project_id)
+        user = auth.get_user_by_info(info)
+        project = project_manager.get_project(project_id)
+        item = manager.create_labeling_task(
+            project_id, labeling_task_name, labeling_task_type, labeling_task_target_id
+        )
+
+        if labels is not None:
+            label_manager.create_labels(project_id, str(item.id), labels)
+
+        doc_ock.post_event(
+            user,
+            events.AddLabelingTask(
+                ProjectName=f"{project.name}-{project.id}",
+                Name=labeling_task_name,
+                Type=labeling_task_type,
+            ),
+        )
+        notification.send_organization_update(
+            project_id, f"labeling_task_created:{str(item.id)}"
+        )
+        return CreateTaskAndLabels(ok=True, task_id=str(item.id))
+
+
 class UpdateLabelingTask(graphene.Mutation):
     class Arguments:
         project_id = graphene.ID(required=True)
@@ -103,3 +149,4 @@ class LabelingTaskMutation(graphene.ObjectType):
     create_labeling_task = CreateLabelingTask.Field()
     update_labeling_task = UpdateLabelingTask.Field()
     delete_labeling_task = DeleteLabelingTask.Field()
+    create_task_and_labels = CreateTaskAndLabels.Field()

graphql_api/mutation/labeling_task_label.py

@@ -7,7 +7,7 @@
 from controller.auth import manager as auth
 from controller.project import manager as project_manager
 from util import doc_ock, notification
-from typing import Dict
+from typing import Dict, List
 
 
 class CreateLabelingTaskLabel(graphene.Mutation):
@@ -50,6 +50,37 @@ def mutate(
         return CreateLabelingTaskLabel(label=label)
 
 
+class CreateLabelingTaskLabels(graphene.Mutation):
+    class Arguments:
+        project_id = graphene.ID(required=True)
+        labeling_task_id = graphene.ID(required=True)
+        labels = graphene.List(graphene.String, required=True)
+
+    ok = graphene.Boolean()
+
+    def mutate(self, info, project_id: str, labeling_task_id: str, labels: List[str]):
+        auth.check_demo_access(info)
+        auth.check_project_access(info, project_id)
+        user = auth.get_user_by_info(info)
+        created_labels = manager.create_labels(project_id, labeling_task_id, labels)
+        task = task_manager.get_labeling_task(project_id, labeling_task_id)
+        project = project_manager.get_project(project_id)
+        for label in created_labels:
+            doc_ock.post_event(
+                user,
+                events.AddLabel(
+                    ProjectName=f"{project.name}-{project_id}",
+                    Name=label.name,
+                    LabelingTaskName=task.name,
+                ),
+            )
+            notification.send_organization_update(
+                project_id, f"label_created:{label.id}:labeling_task:{labeling_task_id}"
+            )
+
+        return CreateLabelingTaskLabels(ok=True)
+
+
 class UpdateLabelingTaskLabelColor(graphene.Mutation):
     class Arguments:
         project_id = graphene.ID(required=True)
@@ -134,6 +165,7 @@ def mutate(self, info, project_id: str, label_id: str):
 
 class LabelingTaskLabelMutation(graphene.ObjectType):
     create_label = CreateLabelingTaskLabel.Field()
+    create_labels = CreateLabelingTaskLabels.Field()
     delete_label = DeleteLabelingTaskLabel.Field()
     update_label_color = UpdateLabelingTaskLabelColor.Field()
     update_label_hotkey = UpdateLabelingTaskLabelHotkey.Field()