Add helper to generate mmc data row payload

Author: Val Brodsky

libs/labelbox/src/labelbox/schema/data_row_payload_templates.py

@@ -0,0 +1,26 @@
+from typing import Dict, List
+
+from pydantic import BaseModel, Field
+
+from labelbox.schema.data_row import DataRowMetadataField
+
+
+class ModelEvalutationTemlateRowData(BaseModel):
+    type: str = Field(
+        default="application/vnd.labelbox.conversational.model-chat-evaluation",
+        frozen=True,
+    )
+    draft: bool = Field(default=True, frozen=True)
+    rootMessageIds: List[str] = Field(default=[])
+    actors: Dict = Field(default={})
+    version: int = Field(default=2, frozen=True)
+    messages: Dict = Field(default={})
+
+
+class ModelEvaluationTemplate(BaseModel):
+    row_data: ModelEvalutationTemlateRowData = Field(
+        default=ModelEvalutationTemlateRowData()
+    )
+    attachments: List[Dict] = Field(default=[])
+    embeddings: List[Dict] = Field(default=[])
+    metadata_fields: List[DataRowMetadataField] = Field(default=[])

libs/labelbox/tests/integration/test_mmc_data_rows.py

@@ -3,26 +3,35 @@
 
 import pytest
 
+from labelbox.schema.data_row_payload_templates import ModelEvaluationTemplate
+
 
 @pytest.fixture
-def mmc_data_row(dataset, make_metadata_fields, embedding):
-    row_data = {
-        "type": "application/vnd.labelbox.conversational.model-chat-evaluation",
-        "draft": True,
-        "rootMessageIds": ["root1"],
-        "actors": {},
-        "messages": {},
-    }
+def mmc_data_row(dataset):
+    data = ModelEvaluationTemplate()
+
+    content_all = data.model_dump()
+    task = dataset.create_data_rows([content_all])
+    task.wait_till_done()
+    assert task.status == "COMPLETE"
+
+    data_row = list(dataset.data_rows())[0]
+
+    yield data_row
+
+    data_row.delete()
+
 
+@pytest.fixture
+def mmc_data_row_all(dataset, make_metadata_fields, embedding):
+    data = ModelEvaluationTemplate()
+    data.row_data.rootMessageIds = ["root1"]
     vector = [random.uniform(1.0, 2.0) for _ in range(embedding.dims)]
-    embeddings = [{"embedding_id": embedding.id, "vector": vector}]
+    data.embeddings = [{"embedding_id": embedding.id, "vector": vector}]
+    data.metadata_fields = make_metadata_fields
+    data.attachments = [{"type": "RAW_TEXT", "value": "attachment value"}]
 
-    content_all = {
-        "row_data": row_data,
-        "attachments": [{"type": "RAW_TEXT", "value": "attachment value"}],
-        "metadata_fields": make_metadata_fields,
-        "embeddings": embeddings,
-    }
+    content_all = data.model_dump()
     task = dataset.create_data_rows([content_all])
     task.wait_till_done()
     assert task.status == "COMPLETE"
@@ -34,14 +43,27 @@ def mmc_data_row(dataset, make_metadata_fields, embedding):
     data_row.delete()
 
 
-def test_mmc(mmc_data_row, embedding, constants):
+def test_mmc(mmc_data_row):
     data_row = mmc_data_row
+    assert json.loads(data_row.row_data) == {
+        "type": "application/vnd.labelbox.conversational.model-chat-evaluation",
+        "draft": True,
+        "rootMessageIds": [],
+        "actors": {},
+        "messages": {},
+        "version": 2,
+    }
+
+
+def test_mmc_all(mmc_data_row_all, embedding, constants):
+    data_row = mmc_data_row_all
     assert json.loads(data_row.row_data) == {
         "type": "application/vnd.labelbox.conversational.model-chat-evaluation",
         "draft": True,
         "rootMessageIds": ["root1"],
         "actors": {},
         "messages": {},
+        "version": 2,
     }
 
     metadata_fields = data_row.metadata_fields