diff --git a/libs/labelbox/CHANGELOG.md b/libs/labelbox/CHANGELOG.md index 95df1750c..e5be4efac 100644 --- a/libs/labelbox/CHANGELOG.md +++ b/libs/labelbox/CHANGELOG.md @@ -1,4 +1,8 @@ # Changelog +# Version 7.9.0 (Unreleased) +## Added +* Add optional `batch_ids` parameter to `Project.get_overview()` to return workflow state counts scoped to one or more batches ([#TBD](https://github.com/Labelbox/labelbox-python/pull/TBD)) + # Version 7.8.0 (2026-06-11) ## Added * Add `ModelRun.total_cost` and `ModelRun.total_data_rows` properties to retrieve inference cost and data row count for a model run ([#2057](https://github.com/Labelbox/labelbox-python/pull/2057)) diff --git a/libs/labelbox/src/labelbox/schema/project.py b/libs/labelbox/src/labelbox/schema/project.py index f39677a95..8edde794c 100644 --- a/libs/labelbox/src/labelbox/schema/project.py +++ b/libs/labelbox/src/labelbox/schema/project.py @@ -86,6 +86,20 @@ pass +_MAX_BATCH_IDS = 1000 + + +def _validate_batch_ids(batch_ids: List[str]) -> None: + if not isinstance(batch_ids, list): + raise ValueError("batch_ids filter expects a list.") + if len(batch_ids) == 0: + raise ValueError("batch_ids filter expects a non-empty list.") + if len(batch_ids) > _MAX_BATCH_IDS: + raise ValueError( + f"batch_ids filter only supports a max of {_MAX_BATCH_IDS} items." + ) + + DataRowPriority = int LabelingParameterOverrideInput = Tuple[DataRowIdentifier, DataRowPriority] @@ -1737,7 +1751,9 @@ def __check_data_rows_have_been_processed( ] def get_overview( - self, details=False + self, + details: bool = False, + batch_ids: Optional[List[str]] = None, ) -> Union[ProjectOverview, ProjectOverviewDetailed]: """Return the overview of a project. @@ -1745,30 +1761,45 @@ def get_overview( which is equivalent to the Overview tab of a project. Args: - details (bool, optional): Whether to include detailed queue information for review and rework queues. - Defaults to False. + details (bool, optional): Whether to include detailed queue information for + review and rework queues. Defaults to False. + batch_ids (Optional[List[str]], optional): When provided, limits counts to data + rows in the given batch(es). Multiple batch IDs return combined counts, not + per-batch breakdowns. Unknown or foreign batch IDs return zero counts. + Defaults to None (project-wide counts). Returns: - Union[ProjectOverview, ProjectOverviewDetailed]: An object representing the project overview. - If `details` is False, returns a `ProjectOverview` object. - If `details` is True, returns a `ProjectOverviewDetailed` object. + Union[ProjectOverview, ProjectOverviewDetailed]: An object representing the + project overview. If `details` is False, returns a `ProjectOverview` object. + If `details` is True, returns a `ProjectOverviewDetailed` object. When + `batch_ids` is set, `issues` is None because issue counts are not + batch-scoped. Raises: + ValueError: If `batch_ids` is an empty list or exceeds the maximum allowed size. Exception: If there is an error executing the query. """ - query = """query ProjectGetOverviewPyApi($projectId: ID!) { - project(where: { id: $projectId }) { - workstreamStateCounts { + if batch_ids is not None: + _validate_batch_ids(batch_ids) + + query = """query ProjectGetOverviewPyApi( + $projectId: ID!, + $batchIds: [String!], + $countInput: DataRowCountQueryInput, + $includeIssues: Boolean! + ) { + project(where: { id: $projectId }) { + workstreamStateCounts(batchIds: $batchIds) { state count } taskQueues { queueType name - dataRowCount + dataRowCount(input: $countInput) } - issues { + issues @include(if: $includeIssues) { totalCount } completedDataRowCount @@ -1776,9 +1807,26 @@ def get_overview( } """ + variables: Dict[str, Any] = { + "projectId": self.uid, + "batchIds": batch_ids, + "includeIssues": batch_ids is None, + } + if batch_ids is not None: + variables["countInput"] = { + "searchQuery": { + "scope": {"projectId": self.uid}, + "query": [ + {"ids": batch_ids, "operator": "is", "type": "batch"} + ], + } + } + else: + variables["countInput"] = None + # Must use experimental to access "issues" result = self.client.execute( - query, {"projectId": self.uid}, experimental=True + query, variables, experimental=True )["project"] # Reformat category names @@ -1788,7 +1836,10 @@ def get_overview( if st["state"] != "NotInTaskQueue" } - overview["issues"] = result.get("issues", {}).get("totalCount") + if batch_ids is None: + overview["issues"] = result.get("issues", {}).get("totalCount") + else: + overview["issues"] = None # Rename categories overview["to_label"] = overview.pop("unlabeled") diff --git a/libs/labelbox/src/labelbox/schema/project_overview.py b/libs/labelbox/src/labelbox/schema/project_overview.py index 8be20fbda..6002f0689 100644 --- a/libs/labelbox/src/labelbox/schema/project_overview.py +++ b/libs/labelbox/src/labelbox/schema/project_overview.py @@ -1,4 +1,4 @@ -from typing import Dict, List +from typing import Dict, List, Optional from pydantic import BaseModel from typing_extensions import TypedDict @@ -16,6 +16,7 @@ class ProjectOverview(BaseModel): The `skipped` attribute represents the number of data rows that have been skipped (Skipped). The `done` attribute represents the number of data rows that have been marked as Done (Done). The `issues` attribute represents the number of data rows with associated issues (Issues). + When the overview is scoped to one or more batches, `issues` is `None`. The following don't appear in the UI The `labeled` attribute represents the number of data rows that have been labeled. @@ -27,7 +28,7 @@ class ProjectOverview(BaseModel): in_rework: int skipped: int done: int - issues: int + issues: Optional[int] = None labeled: int total_data_rows: int @@ -60,6 +61,7 @@ class ProjectOverviewDetailed(BaseModel): The `skipped` attribute represents the number of data rows that have been skipped (Skipped). The `done` attribute represents the number of data rows that have been marked as Done (Done). The `issues` attribute represents the number of data rows with associated issues (Issues). + When the overview is scoped to one or more batches, `issues` is `None`. The following don't appear in the UI The `labeled` attribute represents the number of data rows that have been labeled. @@ -71,6 +73,6 @@ class ProjectOverviewDetailed(BaseModel): in_rework: _QueueDetail skipped: int done: int - issues: int + issues: Optional[int] = None labeled: int total_data_rows: int diff --git a/libs/labelbox/tests/integration/test_batch.py b/libs/labelbox/tests/integration/test_batch.py index f6b1b1091..963cb6d71 100644 --- a/libs/labelbox/tests/integration/test_batch.py +++ b/libs/labelbox/tests/integration/test_batch.py @@ -2,6 +2,7 @@ from uuid import uuid4 import pytest +import time from lbox.exceptions import ( LabelboxError, MalformedQueryException, @@ -139,6 +140,43 @@ def test_create_batch_with_data_row_class( assert batch.size == len(data_rows) +def test_get_overview_batch_scoped(project: Project, small_dataset: Dataset): + export_task = small_dataset.export() + export_task.wait_till_done() + stream = export_task.get_buffered_stream() + data_rows = [dr.json["data_row"]["id"] for dr in stream] + + batch_a = project.create_batch("batch-a-overview", [data_rows[0]]) + batch_b = project.create_batch("batch-b-overview", [data_rows[1]]) + + timeout_seconds = 60 + sleep_time = 2 + overview_a = None + overview_b = None + while timeout_seconds > 0: + overview_a = project.get_overview(batch_ids=[batch_a.uid]) + overview_b = project.get_overview(batch_ids=[batch_b.uid]) + if ( + overview_a.total_data_rows == 1 + and overview_b.total_data_rows == 1 + ): + break + timeout_seconds -= sleep_time + time.sleep(sleep_time) + else: + raise AssertionError( + "Timed out waiting for batch-scoped overview counts" + ) + + assert overview_a.issues is None + assert overview_b.issues is None + assert overview_a.total_data_rows == 1 + assert overview_b.total_data_rows == 1 + + combined = project.get_overview(batch_ids=[batch_a.uid, batch_b.uid]) + assert combined.total_data_rows == 2 + + def test_archive_batch(project: Project, small_dataset: Dataset): export_task = small_dataset.export() export_task.wait_till_done() diff --git a/libs/labelbox/tests/unit/test_project.py b/libs/labelbox/tests/unit/test_project.py index 42559e2dc..8732f01d8 100644 --- a/libs/labelbox/tests/unit/test_project.py +++ b/libs/labelbox/tests/unit/test_project.py @@ -1,7 +1,7 @@ import pytest from unittest.mock import MagicMock -from labelbox.schema.project import Project +from labelbox.schema.project import Project, _validate_batch_ids from labelbox.schema.ontology_kind import EditorTaskType @@ -30,6 +30,82 @@ def project_entity(): ) +def _workstream_state_counts_response(): + return { + "workstreamStateCounts": [ + {"state": "Unlabeled", "count": 10}, + {"state": "InReview", "count": 2}, + {"state": "InRework", "count": 1}, + {"state": "Skipped", "count": 0}, + {"state": "Done", "count": 5}, + {"state": "Labeled", "count": 8}, + {"state": "All", "count": 18}, + {"state": "NotInTaskQueue", "count": 3}, + ], + "taskQueues": [], + "issues": {"totalCount": 4}, + "completedDataRowCount": 5, + } + + +def test_get_overview_project_wide(project_entity): + client = project_entity.client + client.execute.return_value = {"project": _workstream_state_counts_response()} + + overview = project_entity.get_overview() + + assert overview.to_label == 10 + assert overview.total_data_rows == 18 + assert overview.issues == 4 + + args, kwargs = client.execute.call_args + variables = args[1] + assert variables["projectId"] == "test" + assert variables["batchIds"] is None + assert variables["countInput"] is None + assert variables["includeIssues"] is True + assert kwargs["experimental"] is True + query = args[0] + assert "issues @include(if: $includeIssues)" in query + + +def test_get_overview_batch_scoped(project_entity): + client = project_entity.client + client.execute.return_value = {"project": _workstream_state_counts_response()} + + overview = project_entity.get_overview(batch_ids=["batch-1"]) + + assert overview.issues is None + + args, kwargs = client.execute.call_args + variables = args[1] + assert variables["batchIds"] == ["batch-1"] + assert variables["includeIssues"] is False + assert variables["countInput"] == { + "searchQuery": { + "scope": {"projectId": "test"}, + "query": [{"ids": ["batch-1"], "operator": "is", "type": "batch"}], + } + } + + +@pytest.mark.parametrize( + "batch_ids,expected_message", + [ + ([], "batch_ids filter expects a non-empty list."), + (["batch-1"] * 1001, "batch_ids filter only supports a max of 1000 items."), + ], +) +def test_validate_batch_ids_rejects_invalid(batch_ids, expected_message): + with pytest.raises(ValueError, match=expected_message): + _validate_batch_ids(batch_ids) + + +def test_get_overview_rejects_empty_batch_ids(project_entity): + with pytest.raises(ValueError, match="batch_ids filter expects a non-empty list."): + project_entity.get_overview(batch_ids=[]) + + @pytest.mark.parametrize( "api_editor_task_type, expected_editor_task_type", [