[PLT-4090] Add batch_ids filter to project.get_overview() (#2060)

Author: maciejtatol

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))

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,48 +1751,82 @@ 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.
 
         This method returns the number of data rows per task queue and issues of a project,
         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
             }
         }
         """
 
+        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")

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

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()

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",
     [