feat(entities): add grouped count queries

packages/nemo_platform_plugin/src/nemo_platform_plugin/entities.py

@@ -233,6 +233,14 @@ async def list(
         page: int = 1,
         page_size: int = 100,
     ) -> ListResponse[EntityT]: ...
+    async def count_by(
+        self,
+        entity_type: EntityTypeLike,
+        field: str,
+        *,
+        workspace: str = DEFAULT_WORKSPACE,
+        filter_obj: dict[str, Any] | None = None,
+    ) -> dict[str, int]: ...
     async def get(self, entity_type: EntityTypeLike, name: str, *, workspace: Optional[str] = None) -> EntityT: ...
     async def get_by_id(self, entity_type: EntityTypeLike, entity_id: str) -> EntityT: ...
     async def update(self, entity: EntityT, *, original_name: str | None = None) -> EntityT: ...
@@ -493,6 +501,34 @@ async def list(
 
         return ListResponse(data=entities, pagination=pagination)
 
+    async def count_by(
+        self,
+        entity_type: EntityTypeLike,
+        field: str,
+        *,
+        workspace: str = DEFAULT_WORKSPACE,
+        filter_obj: dict[str, Any] | None = None,
+    ) -> dict[str, int]:
+        """Return the number of matching entities grouped by ``field``."""
+        if not field.isidentifier():
+            raise ValueError(f"Field '{field}' is not a direct entity data field")
+
+        filter_dict = _convert_filter_obj_to_filter_str(filter_obj) if filter_obj else {}
+        effective_filter = json.dumps(filter_dict) if filter_dict else omit
+
+        response = await self.entities_api.list(
+            _get_entity_type(entity_type),
+            workspace=workspace,
+            filter=effective_filter,
+            page=1,
+            page_size=1,
+            extra_query={"count_by": f"data.{field}"},
+        )
+        group_counts = getattr(response, "group_counts", None)
+        if group_counts is None:
+            raise EntityStoreError("Grouped counts not found in response")
+        return TypeAdapter(dict[str, int]).validate_python(group_counts)
+
     async def create(self, entity: EntityT) -> EntityT:
         """Create a new entity.
 

packages/nemo_platform_plugin/tests/test_entity_client.py

@@ -0,0 +1,62 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from unittest.mock import AsyncMock, Mock
+
+import pytest
+from nemo_platform.types.entities import EntitiesPage
+from nemo_platform.types.shared.pagination_data import PaginationData
+from nemo_platform_plugin.entities import EntityClient, EntityStoreError
+
+
+class ExperimentGroup:
+    __entity_type__ = "experiment_group"
+
+
+def _entities_page(group_counts: dict[str, int] | None = None) -> EntitiesPage:
+    return EntitiesPage.model_construct(
+        data=[],
+        pagination=PaginationData(page=1, page_size=1, current_page_size=0, total_pages=0, total_results=0),
+        group_counts=group_counts,
+    )
+
+
+@pytest.mark.asyncio
+async def test_count_by_returns_grouped_counts_for_shorthand_filter() -> None:
+    mock_api = Mock()
+    mock_api.list = AsyncMock(return_value=_entities_page(group_counts={"insight-a": 2}))
+    client = EntityClient(mock_api)
+
+    counts = await client.count_by(
+        ExperimentGroup,
+        "insight_id",
+        filter_obj={
+            "insight_id": {"$in": ["insight-a"]},
+            "is_deleted": False,
+        },
+    )
+
+    assert counts == {"insight-a": 2}
+    assert mock_api.list.await_args.kwargs["filter"] == (
+        '{"data.insight_id": {"$in": ["insight-a"]}, "data.is_deleted": false}'
+    )
+    assert mock_api.list.await_args.kwargs["extra_query"] == {"count_by": "data.insight_id"}
+
+
+@pytest.mark.asyncio
+async def test_count_by_rejects_response_without_grouped_counts() -> None:
+    mock_api = Mock()
+    mock_api.list = AsyncMock(return_value=_entities_page())
+    client = EntityClient(mock_api)
+
+    with pytest.raises(EntityStoreError, match="Grouped counts not found"):
+        await client.count_by(ExperimentGroup, "insight_id")
+
+
+@pytest.mark.asyncio
+async def test_count_by_rejects_non_direct_field() -> None:
+    mock_api = Mock()
+    client = EntityClient(mock_api)
+
+    with pytest.raises(ValueError, match="direct entity data field"):
+        await client.count_by(ExperimentGroup, "data.insight_id")