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

services/core/entities/src/nmp/core/entities/api/v2/entities/endpoints.py

@@ -47,7 +47,8 @@
 from sqlalchemy.exc import IntegrityError
 
 
-class EntitiesPage(Page[Entity]): ...
+class EntitiesPage(Page[Entity]):
+    group_counts: dict[str, int] | None = None
 
 
 router = APIRouter()
@@ -336,23 +337,18 @@ async def list_entities(
         description="Sort field",
         examples=["-created_at", "created_at", "-updated_at", "updated_at", "-name", "name"],
     ),
+    count_by: str | None = Query(
+        default=None,
+        description="Optional direct string data field whose matching values should be counted.",
+    ),
 ) -> EntitiesPage:
     """List entities with filtering, supporting cross-workspace queries."""
     accessible_workspaces = await get_accessible_workspaces(repository)
     # Handle cross-workspace query (workspace = "*")
     if workspace == ALL_WORKSPACES:
         # Build combined filter for workspace access and user's filter
-        combined_filter = add_workspace_filtering(accessible_workspaces, filter, field="workspace")
-
-        entities, total = await repository.list_entities(
-            workspace=ALL_WORKSPACES,  # Don't filter by single workspace
-            entity_type=entity_type,
-            page=page,
-            page_size=page_size,
-            sort=sort,
-            filter_op=combined_filter,
-            relationship_child_workspaces=accessible_workspaces,
-        )
+        query_workspace = ALL_WORKSPACES
+        effective_filter = add_workspace_filtering(accessible_workspaces, filter, field="workspace")
     else:
         raise_if_workspace_inaccessible(
             accessible_workspaces,
@@ -362,16 +358,30 @@ async def list_entities(
         # Check if workspace is being deleted (404 for user requests)
         await validate_workspace_not_deleting(workspace_repository, auth_client, workspace)
 
-        # Standard single-workspace query
-        entities, total = await repository.list_entities(
-            workspace=workspace,
-            entity_type=entity_type,
-            page=page,
-            page_size=page_size,
-            sort=sort,
-            filter_op=filter,
-            relationship_child_workspaces=accessible_workspaces,
-        )
+        query_workspace = workspace
+        effective_filter = filter
+
+    entities, total = await repository.list_entities(
+        workspace=query_workspace,
+        entity_type=entity_type,
+        page=page,
+        page_size=page_size,
+        sort=sort,
+        filter_op=effective_filter,
+        relationship_child_workspaces=accessible_workspaces,
+    )
+    group_counts = None
+    if count_by is not None:
+        try:
+            group_counts = await repository.count_entities_by(
+                workspace=query_workspace,
+                entity_type=entity_type,
+                group_by=count_by,
+                filter_op=effective_filter,
+                relationship_child_workspaces=accessible_workspaces,
+            )
+        except ValueError as e:
+            raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(e)) from e
 
     return EntitiesPage(
         data=entities,
@@ -384,6 +394,7 @@ async def list_entities(
         ),
         sort=sort,
         filter=filter.to_dict() if filter else None,
+        group_counts=group_counts,
     )
 
 

services/core/entities/src/nmp/core/entities/app/repository/entity.py

@@ -111,6 +111,20 @@ async def list_entities(
         """
         pass
 
+    @abstractmethod
+    async def count_entities_by(
+        self,
+        *,
+        workspace: str,
+        entity_type: str,
+        group_by: str,
+        filter_op: FilterOperation | None = None,
+        relationship_child_workspaces: set[str] | None = None,
+        session: AsyncSession | None = None,
+    ) -> dict[str, int]:
+        """Count filtered entities grouped by a direct string data field."""
+        pass
+
     @abstractmethod
     async def update_entity(
         self,

services/core/entities/src/nmp/core/entities/app/repository/sqlalchemy/entity.py

@@ -15,10 +15,12 @@
 from nmp.core.entities.app.repository.sqlalchemy.models import DBEntity
 from nmp.core.entities.entities import Entity
 from nmp.core.entities.utils.identifiers import generate_entity_id
-from sqlalchemy import func, select
+from sqlalchemy import String, cast, func, select
 from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
 from sqlalchemy.orm.exc import StaleDataError
 
+MAX_GROUP_COUNT_ROWS = 1000
+
 
 class SQLAlchemyEntityRepository(EntityRepositoryInterface):
     """SQLAlchemy implementation of Entity repository."""
@@ -197,6 +199,53 @@ async def list_entities(
 
             return entities, total
 
+    async def count_entities_by(
+        self,
+        *,
+        workspace: str,
+        entity_type: str,
+        group_by: str,
+        filter_op: FilterOperation | None = None,
+        relationship_child_workspaces: set[str] | None = None,
+        session: AsyncSession | None = None,
+    ) -> dict[str, int]:
+        """Count filtered entities grouped by a direct string data field."""
+        async with self._get_session(session) as sess:
+            parts = group_by.split(".")
+            if len(parts) != 2 or parts[0] != "data" or not parts[1].isidentifier():
+                raise ValueError(f"Field '{group_by}' is not a direct string data field")
+
+            field = parts[1]
+            raw_group_column = DBEntity.data[field]
+            group_column = func.trim(cast(raw_group_column, String), '"')
+            if self._is_sqlite(sess):
+                json_type = func.json_type(DBEntity.data, f"$.{field}")
+                string_type = "text"
+            else:
+                json_type = func.json_typeof(raw_group_column)
+                string_type = "string"
+
+            filter_repo = SQLAlchemyFilterRepository(
+                DBEntity, relationship_child_workspaces=relationship_child_workspaces
+            )
+            query = select(group_column, func.count()).select_from(DBEntity).where(DBEntity.entity_type == entity_type)
+
+            if workspace != ALL_WORKSPACES:
+                query = query.where(DBEntity.workspace == workspace)
+
+            if filter_op is not None:
+                query = query.where(filter_op.apply(filter_repo))
+
+            query = query.where(json_type == string_type)
+
+            rows = (await sess.execute(query.group_by(group_column).limit(MAX_GROUP_COUNT_ROWS + 1))).all()
+            if len(rows) > MAX_GROUP_COUNT_ROWS:
+                raise ValueError(
+                    f"Grouped count has more than {MAX_GROUP_COUNT_ROWS} distinct values; "
+                    "narrow the filter or choose a lower-cardinality field."
+                )
+            return {str(key): int(count) for key, count in rows}
+
     async def update_entity(
         self,
         *,