From bfbbf6e94a8ac0f61acd2bf6afa64c4def83488b Mon Sep 17 00:00:00 2001
From: Val Brodsky <vbrodsky@labelbox.com>
Date: Fri, 28 Jun 2024 11:45:16 -0400
Subject: [PATCH] Prevent adding batches to live chat evaluation projects

---
 libs/labelbox/src/labelbox/schema/project.py      | 15 +++++++++++++++
 .../test_chat_evaluation_ontology_project.py      | 10 ++++++----
 2 files changed, 21 insertions(+), 4 deletions(-)

diff --git a/libs/labelbox/src/labelbox/schema/project.py b/libs/labelbox/src/labelbox/schema/project.py
index f1263e06d..c5f9ef0b0 100644
--- a/libs/labelbox/src/labelbox/schema/project.py
+++ b/libs/labelbox/src/labelbox/schema/project.py
@@ -138,8 +138,16 @@ class Project(DbObject, Updateable, Deletable):
     _wait_processing_max_seconds = 3600
 
     def is_chat_evaluation(self) -> bool:
+        """
+        Returns:
+            True if this project is a live chat evaluation project, False otherwise
+        """
         return self.media_type == MediaType.Conversational and self.editor_task_type == EditorTaskType.ModelChatEvaluation
 
+    def is_auto_data_generation(self) -> bool:
+        return self.media_type == MediaType.LLMPromptCreation or self.media_type == MediaType.LLMPromptResponseCreation or self.is_chat_evaluation(
+        )
+
     def project_model_configs(self):
         query_str = """query ProjectModelConfigsPyApi($id: ID!) {
             project(where: {id : $id}) {
@@ -861,11 +869,18 @@ def create_batch(
                 'coverage_percentage': 0.1}
 
         Returns: the created batch
+
+        Raises:
+            labelbox.exceptions.ValueError if a project is not batch mode, if the project is auto data generation, if the batch exceeds 100k data rows
         """
         # @TODO: make this automatic?
         if self.queue_mode != QueueMode.Batch:
             raise ValueError("Project must be in batch mode")
 
+        if self.is_auto_data_generation():
+            raise ValueError(
+                "Cannot create batches for auto data generation projects")
+
         dr_ids = []
         if data_rows is not None:
             for dr in data_rows:
diff --git a/libs/labelbox/tests/integration/test_chat_evaluation_ontology_project.py b/libs/labelbox/tests/integration/test_chat_evaluation_ontology_project.py
index 078f7bab1..eddc54aaf 100644
--- a/libs/labelbox/tests/integration/test_chat_evaluation_ontology_project.py
+++ b/libs/labelbox/tests/integration/test_chat_evaluation_ontology_project.py
@@ -33,15 +33,17 @@ def test_create_chat_evaluation_ontology_project(
     assert project.labeling_frontend().name == "Editor"
     assert project.ontology().name == ontology.name
 
-    with pytest.raises(MalformedQueryException,
-                       match="No valid data rows to add to project"):
+    with pytest.raises(
+            ValueError,
+            match="Cannot create batches for auto data generation projects"):
         project.create_batch(
             rand_gen(str),
             [offline_conversational_data_row.uid],  # sample of data row objects
         )
 
-    with pytest.raises(MalformedQueryException,
-                       match="No valid data rows to add to project"):
+    with pytest.raises(
+            ValueError,
+            match="Cannot create batches for auto data generation projects"):
         with patch('labelbox.schema.project.MAX_SYNC_BATCH_ROW_COUNT',
                    new=0):  # force to async