Handle out of band exceptions for Project add_model_config in order to deal with a model_setup_complete scenario

Author: Val Brodsky

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

@@ -1,3 +1,4 @@
+import re
 import json
 import logging
 from string import Template
@@ -1268,7 +1269,22 @@ def add_model_config(self, model_config_id: str) -> str:
             "projectId": self.uid,
             "modelConfigId": model_config_id,
         }
-        result = self.client.execute(query, params)
+        try:
+            result = self.client.execute(query, params)
+        except LabelboxError as e:
+            # unfortunately, this is the type of errors our client does not deal with and so the error message is not in the same format as the other errors
+            # needs custom parsing
+            error_string = e.message
+            # Regex to find the message content
+            pattern = r"'message': '([^']+)'"
+            # Search for the pattern in the error string
+            match = re.search(pattern, error_string)
+            if match:
+                error_content = match.group(1)
+            else:
+                error_content = "Unknown error"
+            raise LabelboxError(message=error_content) from e
+
         if not result:
             raise ResourceNotFoundError(ModelConfig, params)
         return result["createProjectModelConfig"]["projectModelConfigId"]
@@ -1301,18 +1317,19 @@ def set_project_model_setup_complete(self) -> bool:
 
         Returns:
             bool, indicates if the model setup is complete.
+
+        NOTE: This method should only be used for live model evaluation projects.
         """
-        query = """query ModelSetupCompletePyApi($projectId: ID!) {
-            project(where: {id: $projectId}, data: {modelSetupComplete: true}) {
+        query = """mutation SetProjectModelSetupCompletePyApi($projectId: ID!) {
+            setProjectModelSetupComplete(where: {id: $projectId}, data: {modelSetupComplete: true}) {
                 modelSetupComplete
             }
         }"""
 
         result = self.client.execute(query, {"projectId": self.uid})
-
-        self.update(
-            model_setup_complete=result["project"]["modelSetupComplete"])
-        return result["project"]["modelSetupComplete"]
+        self.model_setup_complete = result["setProjectModelSetupComplete"][
+            "modelSetupComplete"]
+        return result["setProjectModelSetupComplete"]["modelSetupComplete"]
 
     def set_labeling_parameter_overrides(
             self, data: List[LabelingParameterOverrideInput]) -> bool:

libs/labelbox/tests/integration/conftest.py

@@ -407,7 +407,7 @@ def chat_evaluation_ontology(client, rand_gen):
 
 
 @pytest.fixture
-def chat_evaluation_project_create_dataset(client, rand_gen):
+def live_chat_evaluation_project_with_new_dataset(client, rand_gen):
     project_name = f"test-model-evaluation-project-{rand_gen(str)}"
     dataset_name = f"test-model-evaluation-dataset-{rand_gen(str)}"
     project = client.create_model_evaluation_project(name=project_name,

libs/labelbox/tests/integration/test_chat_evaluation_ontology_project.py

@@ -6,7 +6,7 @@
 
 def test_create_chat_evaluation_ontology_project(
         client, chat_evaluation_ontology,
-        chat_evaluation_project_create_dataset, conversation_data_row,
+        live_chat_evaluation_project_with_new_dataset, conversation_data_row,
         rand_gen):
     ontology = chat_evaluation_ontology
 
@@ -23,7 +23,7 @@ def test_create_chat_evaluation_ontology_project(
         assert classification.schema_id
         assert classification.feature_schema_id
 
-    project = chat_evaluation_project_create_dataset
+    project = live_chat_evaluation_project_with_new_dataset
     assert project.model_setup_complete is None
 
     project.setup_editor(ontology)

libs/labelbox/tests/integration/test_project_set_model_setup_complete.py

@@ -0,0 +1,41 @@
+import pytest
+
+from labelbox.exceptions import LabelboxError, OperationNotAllowedException
+
+
+def test_live_chat_evaluation_project(
+        live_chat_evaluation_project_with_new_dataset, model_config):
+
+    project = live_chat_evaluation_project_with_new_dataset
+
+    project.set_project_model_setup_complete()
+    assert bool(project.model_setup_complete) is True
+
+    with pytest.raises(
+            expected_exception=LabelboxError,
+            match=
+            "Cannot create model config for project because model setup is complete"
+    ):
+        project.add_model_config(model_config.uid)
+
+
+def test_offline_chat_evaluation_project(offline_chat_evaluation_project,
+                                         model_config):
+
+    project = offline_chat_evaluation_project
+
+    with pytest.raises(
+            expected_exception=OperationNotAllowedException,
+            match=
+            "Only live model chat evaluation projects can complete model setup"
+    ):
+        project.set_project_model_setup_complete()
+
+
+def test_any_other_project(project, model_config):
+    with pytest.raises(
+            expected_exception=OperationNotAllowedException,
+            match=
+            "Only live model chat evaluation projects can complete model setup"
+    ):
+        project.set_project_model_setup_complete()