From 17467cdfba72f0fcbbdaa2dc1fc2f0eb64d560ad Mon Sep 17 00:00:00 2001 From: Val Brodsky Date: Wed, 9 Oct 2024 05:38:31 -0700 Subject: [PATCH] Address QA feedback --- .../src/labelbox/schema/data_row_payload_templates.py | 9 +++++---- libs/labelbox/tests/integration/test_mmc_data_rows.py | 6 ++++-- 2 files changed, 9 insertions(+), 6 deletions(-) diff --git a/libs/labelbox/src/labelbox/schema/data_row_payload_templates.py b/libs/labelbox/src/labelbox/schema/data_row_payload_templates.py index bf64e055f..2e2728daa 100644 --- a/libs/labelbox/src/labelbox/schema/data_row_payload_templates.py +++ b/libs/labelbox/src/labelbox/schema/data_row_payload_templates.py @@ -1,11 +1,11 @@ -from typing import Dict, List +from typing import Dict, List, Optional from pydantic import BaseModel, Field from labelbox.schema.data_row import DataRowMetadataField -class ModelEvalutationTemlateRowData(BaseModel): +class ModelEvalutationTemplateRowData(BaseModel): type: str = Field( default="application/vnd.labelbox.conversational.model-chat-evaluation", frozen=True, @@ -15,6 +15,7 @@ class ModelEvalutationTemlateRowData(BaseModel): actors: Dict = Field(default={}) version: int = Field(default=2, frozen=True) messages: Dict = Field(default={}) + global_key: Optional[str] = None class ModelEvaluationTemplate(BaseModel): @@ -32,8 +33,8 @@ class ModelEvaluationTemplate(BaseModel): >>> task = dataset.create_data_rows([content]) """ - row_data: ModelEvalutationTemlateRowData = Field( - default=ModelEvalutationTemlateRowData() + row_data: ModelEvalutationTemplateRowData = Field( + default=ModelEvalutationTemplateRowData() ) attachments: List[Dict] = Field(default=[]) embeddings: List[Dict] = Field(default=[]) diff --git a/libs/labelbox/tests/integration/test_mmc_data_rows.py b/libs/labelbox/tests/integration/test_mmc_data_rows.py index 3b4f95530..2fa7bdd1b 100644 --- a/libs/labelbox/tests/integration/test_mmc_data_rows.py +++ b/libs/labelbox/tests/integration/test_mmc_data_rows.py @@ -10,7 +10,7 @@ def mmc_data_row(dataset): data = ModelEvaluationTemplate() - content_all = data.model_dump() + content_all = data.model_dump(exclude_none=True) task = dataset.create_data_rows([content_all]) task.wait_till_done() assert task.status == "COMPLETE" @@ -26,12 +26,13 @@ def mmc_data_row(dataset): def mmc_data_row_all(dataset, make_metadata_fields, embedding): data = ModelEvaluationTemplate() data.row_data.rootMessageIds = ["root1"] + data.row_data.global_key = "global_key" vector = [random.uniform(1.0, 2.0) for _ in range(embedding.dims)] data.embeddings = [{"embedding_id": embedding.id, "vector": vector}] data.metadata_fields = make_metadata_fields data.attachments = [{"type": "RAW_TEXT", "value": "attachment value"}] - content_all = data.model_dump() + content_all = data.model_dump(exclude_none=True) task = dataset.create_data_rows([content_all]) task.wait_till_done() assert task.status == "COMPLETE" @@ -64,6 +65,7 @@ def test_mmc_all(mmc_data_row_all, embedding, constants): "actors": {}, "messages": {}, "version": 2, + "globalKey": "global_key", } metadata_fields = data_row.metadata_fields