Converted TypedDict to BaseModel for client.send_to_annotate_from_catelog method (#1608)

Author: Gabefire

libs/labelbox/src/labelbox/client.py

@@ -2002,7 +2002,7 @@ def send_to_annotate_from_catalog(self, destination_project_id: str,
                                       task_queue_id: Optional[str],
                                       batch_name: str,
                                       data_rows: Union[DataRowIds, GlobalKeys],
-                                      params: SendToAnnotateFromCatalogParams):
+                                      params: Dict[str, Any]):
         """
         Sends data rows from catalog to a specified project for annotation.
 
@@ -2033,6 +2033,8 @@ def send_to_annotate_from_catalog(self, destination_project_id: str,
 
         """
 
+        validated_params = SendToAnnotateFromCatalogParams(**params)
+
         mutation_str = """mutation SendToAnnotateFromCatalogPyApi($input: SendToAnnotateFromCatalogInput!) {
                             sendToAnnotateFromCatalog(input: $input) {
                               taskId
@@ -2044,26 +2046,14 @@ def send_to_annotate_from_catalog(self, destination_project_id: str,
             task_queue_id)
         data_rows_query = self.build_catalog_query(data_rows)
 
-        source_model_run_id = params.get("source_model_run_id", None)
-        predictions_ontology_mapping = params.get(
-            "predictions_ontology_mapping", None)
         predictions_input = build_predictions_input(
-            predictions_ontology_mapping,
-            source_model_run_id) if source_model_run_id else None
+            validated_params.predictions_ontology_mapping,
+            validated_params.source_model_run_id
+        ) if validated_params.source_model_run_id else None
 
-        source_project_id = params.get("source_project_id", None)
-        annotations_ontology_mapping = params.get(
-            "annotations_ontology_mapping", None)
         annotations_input = build_annotations_input(
-            annotations_ontology_mapping,
-            source_project_id) if source_project_id else None
-
-        batch_priority = params.get("batch_priority", 5)
-        exclude_data_rows_in_project = params.get(
-            "exclude_data_rows_in_project", False)
-        override_existing_annotations_rule = params.get(
-            "override_existing_annotations_rule",
-            ConflictResolutionStrategy.KeepExisting)
+            validated_params.annotations_ontology_mapping, validated_params.
+            source_project_id) if validated_params.source_project_id else None
 
         res = self.execute(
             mutation_str, {
@@ -2072,18 +2062,18 @@ def send_to_annotate_from_catalog(self, destination_project_id: str,
                         destination_project_id,
                     "batchInput": {
                         "batchName": batch_name,
-                        "batchPriority": batch_priority
+                        "batchPriority": validated_params.batch_priority
                     },
                     "destinationTaskQueue":
                         destination_task_queue,
                     "excludeDataRowsInProject":
-                        exclude_data_rows_in_project,
+                        validated_params.exclude_data_rows_in_project,
                     "annotationsInput":
                         annotations_input,
                     "predictionsInput":
                         predictions_input,
                     "conflictLabelsResolutionStrategy":
-                        override_existing_annotations_rule,
+                        validated_params.override_existing_annotations_rule,
                     "searchQuery": {
                         "scope": None,
                         "query": [data_rows_query]

libs/labelbox/src/labelbox/schema/send_to_annotate_params.py

@@ -3,14 +3,15 @@
 from typing import Optional, Dict
 
 from labelbox.schema.conflict_resolution_strategy import ConflictResolutionStrategy
+from labelbox import pydantic_compat
 
 if sys.version_info >= (3, 8):
     from typing import TypedDict
 else:
     from typing_extensions import TypedDict
 
 
-class SendToAnnotateFromCatalogParams(TypedDict):
+class SendToAnnotateFromCatalogParams(pydantic_compat.BaseModel):
     """
     Extra parameters for sending data rows to a project through catalog. At least one of source_model_run_id or
     source_project_id must be provided.
@@ -30,14 +31,26 @@ class SendToAnnotateFromCatalogParams(TypedDict):
     :param batch_priority: Optional[int] - The priority of the batch. Defaults to 5.
     """
 
-    source_model_run_id: Optional[str]
-    predictions_ontology_mapping: Optional[Dict[str, str]]
-    source_project_id: Optional[str]
-    annotations_ontology_mapping: Optional[Dict[str, str]]
-    exclude_data_rows_in_project: Optional[bool]
-    override_existing_annotations_rule: Optional[ConflictResolutionStrategy]
-    batch_priority: Optional[int]
-
+    source_model_run_id: Optional[str] = None
+    source_project_id: Optional[str] = None
+    predictions_ontology_mapping: Optional[Dict[str, str]] = {}
+    annotations_ontology_mapping: Optional[Dict[str, str]] = {}
+    exclude_data_rows_in_project: Optional[bool] = False
+    override_existing_annotations_rule: Optional[
+        ConflictResolutionStrategy] = ConflictResolutionStrategy.KeepExisting
+    batch_priority: Optional[int] = 5
+
+    @pydantic_compat.root_validator
+    def check_project_id_or_model_run_id(cls, values):
+        if not values.get("source_model_run_id") and not values.get("source_project_id"):
+            raise ValueError(
+                'Either source_project_id or source_model_id are required'
+            )
+        if values.get("source_model_run_id") and values.get("source_project_id"):
+            raise ValueError(
+                'Provide only a source_project_id or source_model_id not both'
+            ) 
+        return values
 
 class SendToAnnotateFromModelParams(TypedDict):
     """