[PLT-772] SDK Ontology, project, editor support for chat evaluation (#1594)

Merging with failing tests, this is an ADV issue under investigation

Author: vbrodsky

docs/labelbox/index.rst

@@ -31,6 +31,7 @@ Labelbox Python SDK Documentation
     model
     model-run
     ontology
+    ontology_kind
     organization
     pagination
     project

docs/labelbox/ontology-kind.rst

@@ -0,0 +1,7 @@
+OntologyKind
+===============================================================================================
+
+.. automodule:: labelbox.schema.ontology_kind
+   :members:
+   :exclude-members: EditorTaskType
+   :show-inheritance:

libs/labelbox/src/labelbox/__init__.py

@@ -35,4 +35,5 @@
 from labelbox.schema.queue_mode import QueueMode
 from labelbox.schema.task_queue import TaskQueue
 from labelbox.schema.identifiables import UniqueIds, GlobalKeys, DataRowIds
-from labelbox.schema.identifiable import UniqueId, GlobalKey
+from labelbox.schema.identifiable import UniqueId, GlobalKey
+from labelbox.schema.ontology_kind import OntologyKind

libs/labelbox/src/labelbox/client.py

@@ -21,7 +21,7 @@
 from labelbox.adv_client import AdvClient
 from labelbox.orm import query
 from labelbox.orm.db_object import DbObject
-from labelbox.orm.model import Entity
+from labelbox.orm.model import Entity, Field
 from labelbox.pagination import PaginatedCollection
 from labelbox.schema import role
 from labelbox.schema.conflict_resolution_strategy import ConflictResolutionStrategy
@@ -52,6 +52,8 @@
 from labelbox.schema.slice import CatalogSlice, ModelSlice
 from labelbox.schema.task import Task
 from labelbox.schema.user import User
+from labelbox.schema.ontology_kind import (OntologyKind, EditorTaskTypeMapper,
+                                           EditorTaskType)
 
 logger = logging.getLogger(__name__)
 
@@ -563,14 +565,16 @@ def get_labeling_frontends(self, where=None) -> List[LabelingFrontend]:
         """
         return self._get_all(Entity.LabelingFrontend, where)
 
-    def _create(self, db_object_type, data):
+    def _create(self, db_object_type, data, extra_params={}):
         """ Creates an object on the server. Attribute values are
             passed as keyword arguments:
 
         Args:
             db_object_type (type): A DbObjectType subtype.
             data (dict): Keys are attributes or their names (in Python,
                 snake-case convention) and values are desired attribute values.
+            extra_params (dict): Additional parameters to pass to GraphQL.
+                These have to be Field(...): value pairs.
         Returns:
             A new object of the given DB object type.
         Raises:
@@ -585,6 +589,7 @@ def _create(self, db_object_type, data):
             for attr, value in data.items()
         }
 
+        data = {**data, **extra_params}
         query_string, params = query.create(db_object_type, data)
         res = self.execute(query_string, params)
         res = res["create%s" % db_object_type.type_name()]
@@ -699,17 +704,26 @@ def create_project(self, **kwargs) -> Project:
             )
 
         media_type = kwargs.get("media_type")
-        if media_type:
-            if MediaType.is_supported(media_type):
-                media_type = media_type.value
-            else:
-                raise TypeError(f"{media_type} is not a valid media type. Use"
-                                f" any of {MediaType.get_supported_members()}"
-                                " from MediaType. Example: MediaType.Image.")
+        if media_type and MediaType.is_supported(media_type):
+            media_type_value = media_type.value
+        elif media_type:
+            raise TypeError(f"{media_type} is not a valid media type. Use"
+                            f" any of {MediaType.get_supported_members()}"
+                            " from MediaType. Example: MediaType.Image.")
         else:
             logger.warning(
                 "Creating a project without specifying media_type"
                 " through this method will soon no longer be supported.")
+            media_type_value = None
+
+        ontology_kind = kwargs.pop("ontology_kind", None)
+        if ontology_kind and OntologyKind.is_supported(ontology_kind):
+            editor_task_type_value = EditorTaskTypeMapper.to_editor_task_type(
+                ontology_kind, media_type).value
+        elif ontology_kind:
+            raise OntologyKind.get_ontology_kind_validation_error(ontology_kind)
+        else:
+            editor_task_type_value = None
 
         quality_mode = kwargs.get("quality_mode")
         if not quality_mode:
@@ -728,12 +742,47 @@ def create_project(self, **kwargs) -> Project:
         else:
             raise ValueError(f"{quality_mode} is not a valid quality mode.")
 
-        return self._create(Entity.Project, {
-            **data,
-            **({
-                "media_type": media_type
-            } if media_type else {})
-        })
+        params = {**data}
+        if media_type_value:
+            params["media_type"] = media_type_value
+        if editor_task_type_value:
+            params["editor_task_type"] = editor_task_type_value
+
+        extra_params = {
+            Field.String("dataset_name_or_id"):
+                params.pop("dataset_name_or_id", None),
+            Field.Boolean("append_to_existing_dataset"):
+                params.pop("append_to_existing_dataset", None),
+            Field.Int("data_row_count"):
+                params.pop("data_row_count", None),
+        }
+        extra_params = {k: v for k, v in extra_params.items() if v is not None}
+
+        return self._create(Entity.Project, params, extra_params)
+
+    def create_model_evaluation_project(
+            self,
+            dataset_name_or_id: str,
+            append_to_existing_dataset: bool = False,
+            data_row_count: int = 100,
+            **kwargs) -> Project:
+        """
+        Use this method exclusively to create a chat model evaluation project.
+        Args:
+            dataset_name_or_id: The name or id of the dataset to use for the project
+            append_to_existing_dataset: If True, the project will append assets (data rows) to the existing dataset
+            data_row_count: The number of data row assets to use for the project
+            **kwargs: Additional parameters to pass to the the create_project method
+        Returns:
+            Project: The created project
+        """
+        kwargs["media_type"] = MediaType.Conversational
+        kwargs["ontology_kind"] = OntologyKind.ModelEvaluation
+        kwargs["dataset_name_or_id"] = dataset_name_or_id
+        kwargs["append_to_existing_dataset"] = append_to_existing_dataset
+        kwargs["data_row_count"] = data_row_count
+
+        return self.create_project(**kwargs)
 
     def get_roles(self) -> List[Role]:
         """
@@ -938,18 +987,22 @@ def rootSchemaPayloadToFeatureSchema(client, payload):
                                    rootSchemaPayloadToFeatureSchema,
                                    ['rootSchemaNodes', 'nextCursor'])
 
-    def create_ontology_from_feature_schemas(self,
-                                             name,
-                                             feature_schema_ids,
-                                             media_type=None) -> Ontology:
+    def create_ontology_from_feature_schemas(
+            self,
+            name,
+            feature_schema_ids,
+            media_type: MediaType = None,
+            ontology_kind: OntologyKind = None) -> Ontology:
         """
         Creates an ontology from a list of feature schema ids
 
         Args:
             name (str): Name of the ontology
             feature_schema_ids (List[str]): List of feature schema ids corresponding to
                 top level tools and classifications to include in the ontology
-            media_type (MediaType or None): Media type of a new ontology
+            media_type (MediaType or None): Media type of a new ontology. NOTE for chat evaluation, we currently foce media_type to Conversational
+            ontology_kind (OntologyKind or None): set to OntologyKind.ModelEvaluation if the ontology is for chat evaluation,
+                leave as None otherwise.
         Returns:
             The created Ontology
         """
@@ -979,7 +1032,20 @@ def create_ontology_from_feature_schemas(self,
                     "Neither `tool` or `classification` found in the normalized feature schema"
                 )
         normalized = {'tools': tools, 'classifications': classifications}
-        return self.create_ontology(name, normalized, media_type)
+
+        if ontology_kind and ontology_kind is OntologyKind.ModelEvaluation:
+            if media_type is None:
+                media_type = MediaType.Conversational
+            else:
+                if media_type is not MediaType.Conversational:
+                    raise ValueError(
+                        "For chat evaluation, media_type must be Conversational."
+                    )
+
+        return self.create_ontology(name=name,
+                                    normalized=normalized,
+                                    media_type=media_type,
+                                    ontology_kind=ontology_kind)
 
     def delete_unused_feature_schema(self, feature_schema_id: str) -> None:
         """
@@ -1166,7 +1232,11 @@ def get_unused_feature_schemas(self, after: str = None) -> List[str]:
                 "Failed to get unused feature schemas, message: " +
                 str(response.json()['message']))
 
-    def create_ontology(self, name, normalized, media_type=None) -> Ontology:
+    def create_ontology(self,
+                        name,
+                        normalized,
+                        media_type: MediaType = None,
+                        ontology_kind: OntologyKind = None) -> Ontology:
         """
         Creates an ontology from normalized data
             >>> normalized = {"tools" : [{'tool': 'polygon',  'name': 'cat', 'color': 'black'}], "classifications" : []}
@@ -1184,26 +1254,43 @@ def create_ontology(self, name, normalized, media_type=None) -> Ontology:
             name (str): Name of the ontology
             normalized (dict): A normalized ontology payload. See above for details.
             media_type (MediaType or None): Media type of a new ontology
+            ontology_kind (OntologyKind or None): set to OntologyKind.ModelEvaluation if the ontology is for chat evaluation,
+                leave as None otherwise.
+
         Returns:
             The created Ontology
+
+        NOTE caller of this method is expected to set media_type to Conversational if ontology_kind is ModelEvaluation
         """
 
+        media_type_value = None
         if media_type:
             if MediaType.is_supported(media_type):
-                media_type = media_type.value
+                media_type_value = media_type.value
             else:
                 raise get_media_type_validation_error(media_type)
 
+        if ontology_kind and OntologyKind.is_supported(ontology_kind):
+            editor_task_type_value = EditorTaskTypeMapper.to_editor_task_type(
+                ontology_kind, media_type).value
+        elif ontology_kind:
+            raise OntologyKind.get_ontology_kind_validation_error(ontology_kind)
+        else:
+            editor_task_type_value = None
+
         query_str = """mutation upsertRootSchemaNodePyApi($data:  UpsertOntologyInput!){
                            upsertOntology(data: $data){ %s }
         } """ % query.results_query_part(Entity.Ontology)
         params = {
             'data': {
                 'name': name,
                 'normalized': json.dumps(normalized),
-                'mediaType': media_type
+                'mediaType': media_type_value
             }
         }
+        if editor_task_type_value:
+            params['data']['editorTaskType'] = editor_task_type_value
+
         res = self.execute(query_str, params)
         return Entity.Ontology(self, res['upsertOntology'])
 

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

@@ -24,3 +24,4 @@
 import labelbox.schema.identifiables
 import labelbox.schema.identifiable
 import labelbox.schema.catalog
+import labelbox.schema.ontology_kind

libs/labelbox/src/labelbox/schema/ontology.py

@@ -252,6 +252,9 @@ class Type(Enum):
         LINE = "line"
         NER = "named-entity"
         RELATIONSHIP = "edge"
+        MESSAGE_SINGLE_SELECTION = 'message-single-selection'
+        MESSAGE_MULTI_SELECTION = 'message-multi-selection'
+        MESSAGE_RANKING = 'message-ranking'
 
     tool: Type
     name: str

libs/labelbox/src/labelbox/schema/ontology_kind.py

@@ -0,0 +1,56 @@
+from enum import Enum
+from typing import Optional
+
+from labelbox.schema.media_type import MediaType
+
+
+class OntologyKind(Enum):
+    """
+    OntologyKind is an enum that represents the different types of ontologies
+    At the moment it is only limited to ModelEvaluation
+    """
+    ModelEvaluation = "MODEL_EVALUATION"
+    Missing = None
+
+    @classmethod
+    def is_supported(cls, value):
+        return isinstance(value, cls)
+
+    @classmethod
+    def get_ontology_kind_validation_error(cls, ontology_kind):
+        return TypeError(f"{ontology_kind}: is not a valid ontology kind. Use"
+                         f" any of {OntologyKind.__members__.items()}"
+                         " from OntologyKind.")
+
+
+class EditorTaskType(Enum):
+    ModelChatEvaluation = "MODEL_CHAT_EVALUATION"
+    Missing = None
+
+    @classmethod
+    def is_supported(cls, value):
+        return isinstance(value, cls)
+
+
+class EditorTaskTypeMapper:
+
+    @staticmethod
+    def to_editor_task_type(ontology_kind: OntologyKind,
+                            media_type: MediaType) -> EditorTaskType:
+        if ontology_kind and OntologyKind.is_supported(
+                ontology_kind) and media_type and MediaType.is_supported(
+                    media_type):
+            editor_task_type = EditorTaskTypeMapper.map_to_editor_task_type(
+                ontology_kind, media_type)
+        else:
+            editor_task_type = EditorTaskType.Missing
+
+        return editor_task_type
+
+    @staticmethod
+    def map_to_editor_task_type(onotology_kind: OntologyKind,
+                                media_type: MediaType) -> EditorTaskType:
+        if onotology_kind == OntologyKind.ModelEvaluation and media_type == MediaType.Conversational:
+            return EditorTaskType.ModelChatEvaluation
+        else:
+            return EditorTaskType.Missing