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

[maciejtatol]

Ticket: LINK

Description

Problem

Customers such as Pinterest need to poll batch completion status programmatically (To Label / In Review / In Rework / Done counts per batch). Today the only option is running multiple project exports with batch_ids + workflow_status filters - each export takes minutes. For hourly polling across many batches, this is expensive and slow.

Solution

Extend Project.get_overview() with an optional batch_ids parameter that returns the same workflow-state counts as the project Overview tab, scoped to one or more batches - without exporting label data.

overview = project.get_overview(batch_ids=[batch.uid])
print(overview.to_label, overview.in_review, overview.in_rework, overview.done)
print(overview.total_data_rows)

# Multiple batches → combined counts (not per-batch breakdown)
overview = project.get_overview(batch_ids=[batch_a.uid, batch_b.uid])

# Per-queue breakdown, also batch-scoped
overview = project.get_overview(batch_ids=[batch.uid], details=True)

Query: ProjectGetOverviewPyApi (still uses experimental=True / /_gql for issues on project-wide calls)

Validation:

batch_ids must be a non-empty list (empty list raises ValueError - backend treats [] as no filter, which would silently return project-wide counts)
Max 1000 batch IDs per call
Return type change: ProjectOverview.issues and ProjectOverviewDetailed.issues are now Optional[int] = None (always populated for project-wide calls; None when batch-scoped)

Test on Local

Project:

• 2 batches:
• 1 (the one that is requested): 3 total, 1 in review, 1 in rework, 1 in done
• 2: 4 total, all in To Label

UI view:

Test script:

#!/usr/bin/env python3
"""Manual smoke test for Project.get_overview(batch_ids=...) against local lb-api.

Prerequisites:
  - intelligence API running on http://localhost:8080 (pm2 start api)
  - Docker infra up (yarn dev:up) for MySQL + Elasticsearch
  - LABELBOX_API_KEY exported (create one in local app: http://localhost:3333)

Usage:
  export LABELBOX_API_KEY="your-local-key"
  python scripts/test_batch_overview_local.py

Optional:
  python scripts/test_batch_overview_local.py --project-id <id> --batch-id <id>
"""

from __future__ import annotations

import argparse
import os
import sys
import time
import uuid

import labelbox as lb
from labelbox.schema.media_type import MediaType

LOCAL_GRAPHQL = "http://localhost:8080/graphql"
LOCAL_REST = "http://localhost:8080/api/v1"
TEST_IMAGE = (
    "https://storage.googleapis.com/lb-artifacts-testing-public/"
    "sdk_integration_test/potato.jpeg"
)
POLL_TIMEOUT_SECONDS = 120
POLL_INTERVAL_SECONDS = 3


def _client() -> lb.Client:
    api_key = os.environ.get("LABELBOX_API_KEY")
    if not api_key:
        print(
            "ERROR: Set LABELBOX_API_KEY (create one in http://localhost:3333)",
            file=sys.stderr,
        )
        sys.exit(1)
    return lb.Client(
        api_key=api_key,
        endpoint=LOCAL_GRAPHQL,
        rest_endpoint=LOCAL_REST,
    )


def _wait_for_batch_counts(project, batch_a_id: str, batch_b_id: str) -> None:
    deadline = time.time() + POLL_TIMEOUT_SECONDS
    while time.time() < deadline:
        overview_a = project.get_overview(batch_ids=[batch_a_id])
        overview_b = project.get_overview(batch_ids=[batch_b_id])
        if (
            overview_a.total_data_rows == 1
            and overview_b.total_data_rows == 1
        ):
            return
        print(
            f"  waiting for ES indexing... "
            f"batch_a={overview_a.total_data_rows}, "
            f"batch_b={overview_b.total_data_rows}"
        )
        time.sleep(POLL_INTERVAL_SECONDS)
    raise TimeoutError(
        "Timed out waiting for batch-scoped overview counts "
        f"(>{POLL_TIMEOUT_SECONDS}s). Is catalog ES indexing healthy?"
    )


def _run_existing(project_id: str, batch_id: str) -> None:
    client = _client()
    project = client.get_project(project_id)
    print(f"Project: {project.uid} ({project.name})")

    project_overview = project.get_overview()
    batch_overview = project.get_overview(batch_ids=[batch_id])

    print("\nProject-wide overview:")
    print(f"  total_data_rows={project_overview.total_data_rows}")
    print(f"  to_label={project_overview.to_label}")
    print(f"  in_review={project_overview.in_review}")
    print(f"  in_rework={project_overview.in_rework}")
    print(f"  done={project_overview.done}")
    print(f"  issues={project_overview.issues}")

    print("\nBatch-scoped overview:")
    print(f"  total_data_rows={batch_overview.total_data_rows}")
    print(f"  to_label={batch_overview.to_label}")
    print(f"  in_review={batch_overview.in_review}")
    print(f"  in_rework={batch_overview.in_rework}")
    print(f"  done={batch_overview.done}")
    print(f"  issues={batch_overview.issues} (expected None)")

    if batch_overview.issues is not None:
        raise AssertionError("Expected issues=None for batch-scoped overview")


def _run_full_smoke() -> None:
    client = _client()
    suffix = uuid.uuid4().hex[:8]
    project = client.create_project(
        name=f"batch-overview-local-{suffix}",
        media_type=MediaType.Image,
    )
    dataset = client.create_dataset(name=f"batch-overview-dataset-{suffix}")

    print(f"Created project {project.uid}")
    print(f"Created dataset {dataset.uid}")

    try:
        task = dataset.create_data_rows(
            [
                {"row_data": TEST_IMAGE, "external_id": f"dr-a-{suffix}"},
                {"row_data": TEST_IMAGE, "external_id": f"dr-b-{suffix}"},
            ]
        )
        task.wait_till_done()
        if task.errors:
            raise RuntimeError(f"Data row creation errors: {task.errors}")

        export_task = dataset.export()
        export_task.wait_till_done()
        data_row_ids = [
            row.json["data_row"]["id"]
            for row in export_task.get_buffered_stream()
        ]
        if len(data_row_ids) < 2:
            raise RuntimeError(
                f"Expected 2 data rows, got {len(data_row_ids)}"
            )

        batch_a = project.create_batch(f"batch-a-{suffix}", [data_row_ids[0]])
        batch_b = project.create_batch(f"batch-b-{suffix}", [data_row_ids[1]])
        print(f"Created batches: {batch_a.uid}, {batch_b.uid}")

        _wait_for_batch_counts(project, batch_a.uid, batch_b.uid)

        overview_a = project.get_overview(batch_ids=[batch_a.uid])
        overview_b = project.get_overview(batch_ids=[batch_b.uid])
        combined = project.get_overview(batch_ids=[batch_a.uid, batch_b.uid])
        project_overview = project.get_overview()

        print("\nBatch A:")
        print(f"  total_data_rows={overview_a.total_data_rows}, issues={overview_a.issues}")
        print("Batch B:")
        print(f"  total_data_rows={overview_b.total_data_rows}, issues={overview_b.issues}")
        print("Combined:")
        print(f"  total_data_rows={combined.total_data_rows}, issues={combined.issues}")
        print("Project-wide:")
        print(f"  total_data_rows={project_overview.total_data_rows}, issues={project_overview.issues}")

        assert overview_a.total_data_rows == 1
        assert overview_b.total_data_rows == 1
        assert combined.total_data_rows == 2
        assert overview_a.issues is None
        assert overview_b.issues is None
        assert combined.issues is None
        assert project_overview.issues is not None

        print("\nSUCCESS: batch-scoped get_overview() works against local lb-api.")
    finally:
        print("\nCleaning up...")
        project.delete()
        dataset.delete()


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Smoke test Project.get_overview(batch_ids=...) locally"
    )
    parser.add_argument("--project-id", help="Existing project ID (skip setup)")
    parser.add_argument("--batch-id", help="Existing batch ID (requires --project-id)")
    args = parser.parse_args()

    if args.project_id or args.batch_id:
        if not args.project_id or not args.batch_id:
            parser.error("--project-id and --batch-id must be used together")
        _run_existing(args.project_id, args.batch_id)
        print("\nSUCCESS")
        return

    _run_full_smoke()


if __name__ == "__main__":
    main()

Execute script:

Type of change

Please delete options that are not relevant.

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Document change (fix typo or modifying any markdown files, code comments or anything in the examples folder only)

All Submissions

• Have you followed the guidelines in our Contributing document?
• Have you provided a description?
• Are your changes properly formatted?

New Feature Submissions

• Does your submission pass tests?
• Have you added thorough tests for your new feature?
• Have you commented your code, particularly in hard-to-understand areas?
• Have you added a Docstring?

Changes to Core Features

• Have you written new tests for your core changes, as applicable?
• Have you successfully run tests with your changes locally?
• Have you updated any code comments, as applicable?

---

Note

^{Cursor Bugbot is generating a summary for commit 77cc539. Configure here.}

[mrobers1982]

Is there a similar server-side enforcement?

[maciejtatol]

I'll check it, good call-out