From f5d1d4e9e8a8533f6334654c84f642272c6d19fa Mon Sep 17 00:00:00 2001 From: Aditya Pandey Date: Tue, 21 Jul 2026 16:22:14 -0600 Subject: [PATCH 01/15] feat(entities): add grouped entity counts Signed-off-by: Aditya Pandey --- .../core/entities/app/repository/entity.py | 14 +++ .../app/repository/sqlalchemy/entity.py | 34 ++++++ .../app/repository/sqlalchemy/filter.py | 6 + .../repository/test_entity_group_counts.py | 104 ++++++++++++++++++ 4 files changed, 158 insertions(+) create mode 100644 services/core/entities/tests/repository/test_entity_group_counts.py diff --git a/services/core/entities/src/nmp/core/entities/app/repository/entity.py b/services/core/entities/src/nmp/core/entities/app/repository/entity.py index bd0386498a..17a7171698 100644 --- a/services/core/entities/src/nmp/core/entities/app/repository/entity.py +++ b/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 scalar field.""" + pass + @abstractmethod async def update_entity( self, diff --git a/services/core/entities/src/nmp/core/entities/app/repository/sqlalchemy/entity.py b/services/core/entities/src/nmp/core/entities/app/repository/sqlalchemy/entity.py index ff617ed380..51dafca468 100644 --- a/services/core/entities/src/nmp/core/entities/app/repository/sqlalchemy/entity.py +++ b/services/core/entities/src/nmp/core/entities/app/repository/sqlalchemy/entity.py @@ -197,6 +197,40 @@ 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 scalar field.""" + async with self._get_session(session) as sess: + filter_repo = SQLAlchemyFilterRepository( + DBEntity, relationship_child_workspaces=relationship_child_workspaces + ) + group_column, is_json = filter_repo.get_text_column(group_by) + group_counts = ( + select(group_column, func.count()).select_from(DBEntity).where(DBEntity.entity_type == entity_type) + ) + + if workspace != ALL_WORKSPACES: + group_counts = group_counts.where(DBEntity.workspace == workspace) + + if filter_op is not None: + group_counts = group_counts.where(filter_op.apply(filter_repo)) + + group_counts = group_counts.where(group_column.is_not(None)) + if is_json: + group_counts = group_counts.where(group_column != "null") + group_counts = group_counts.group_by(group_column) + + rows = (await sess.execute(group_counts)).all() + return {str(key): int(count) for key, count in rows} + async def update_entity( self, *, diff --git a/services/core/entities/src/nmp/core/entities/app/repository/sqlalchemy/filter.py b/services/core/entities/src/nmp/core/entities/app/repository/sqlalchemy/filter.py index 40eb859d82..7bc60f69e9 100644 --- a/services/core/entities/src/nmp/core/entities/app/repository/sqlalchemy/filter.py +++ b/services/core/entities/src/nmp/core/entities/app/repository/sqlalchemy/filter.py @@ -96,6 +96,12 @@ def _cast_json_to_text(self, column: Any) -> Any: # This handles both SQLite and PostgreSQL JSON string extraction return func.trim(cast(column, String), '"') + def get_text_column(self, field: str) -> tuple[ColumnElement, bool]: + """Return a text-normalized field expression and whether it is JSON.""" + column, is_json = self._get_column(field) + text_column = self._cast_json_to_text(column) if is_json else cast(column, String) + return text_column, is_json + def _cast_json_to_numeric(self, column: Any) -> Any: """Cast a JSON column element to a float for numeric comparisons. diff --git a/services/core/entities/tests/repository/test_entity_group_counts.py b/services/core/entities/tests/repository/test_entity_group_counts.py new file mode 100644 index 0000000000..6b82bc9ed3 --- /dev/null +++ b/services/core/entities/tests/repository/test_entity_group_counts.py @@ -0,0 +1,104 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for grouped entity counts.""" + +import pytest +from nmp.common.api.filter import ComparisonOperation, FilterOperator, LogicalOperation +from nmp.core.entities.app.repository import SQLAlchemyEntityRepository + +pytestmark = pytest.mark.asyncio + + +async def test_counts_filtered_entities_grouped_by_json_field( + entity_repo: SQLAlchemyEntityRepository, setup_workspaces +): + """Count only live experiment groups for the requested insights.""" + for name, insight_id in (("a-1", "insight-a"), ("a-2", "insight-a"), ("b-1", "insight-b")): + await entity_repo.create_entity( + workspace="workspace-1", + entity_type="experiment_group", + name=name, + data={"insight_id": insight_id, "is_deleted": False}, + ) + + await entity_repo.create_entity( + workspace="workspace-1", + entity_type="experiment_group", + name="deleted", + data={"insight_id": "insight-a", "is_deleted": True}, + ) + await entity_repo.create_entity( + workspace="workspace-1", + entity_type="experiment_group", + name="unlinked", + data={"is_deleted": False}, + ) + await entity_repo.create_entity( + workspace="workspace-2", + entity_type="experiment_group", + name="other-workspace", + data={"insight_id": "insight-a", "is_deleted": False}, + ) + await entity_repo.create_entity( + workspace="workspace-1", + entity_type="other_type", + name="other-type", + data={"insight_id": "insight-a", "is_deleted": False}, + ) + + filter_op = LogicalOperation( + operator=FilterOperator.AND, + operations=[ + ComparisonOperation(field="data.insight_id", operator=FilterOperator.IN, value=["insight-a", "insight-b"]), + ComparisonOperation(field="data.is_deleted", operator=FilterOperator.EQ, value=False), + ], + ) + + counts = await entity_repo.count_entities_by( + workspace="workspace-1", + entity_type="experiment_group", + group_by="data.insight_id", + filter_op=filter_op, + ) + + assert counts == {"insight-a": 2, "insight-b": 1} + + +async def test_counts_base_field_literal_null(entity_repo: SQLAlchemyEntityRepository, setup_workspaces): + await entity_repo.create_entity( + workspace="workspace-1", + entity_type="experiment_group", + name="null", + data={}, + ) + + counts = await entity_repo.count_entities_by( + workspace="workspace-1", + entity_type="experiment_group", + group_by="name", + ) + + assert counts == {"null": 1} + + +async def test_omits_missing_and_null_json_group_keys(entity_repo: SQLAlchemyEntityRepository, setup_workspaces): + for name, data in ( + ("linked", {"insight_id": "insight-a"}), + ("missing", {}), + ("null", {"insight_id": None}), + ): + await entity_repo.create_entity( + workspace="workspace-1", + entity_type="experiment_group", + name=name, + data=data, + ) + + counts = await entity_repo.count_entities_by( + workspace="workspace-1", + entity_type="experiment_group", + group_by="data.insight_id", + ) + + assert counts == {"insight-a": 1} From 8487614c12fe70f042c88f0c6441f7cddff605cf Mon Sep 17 00:00:00 2001 From: Aditya Pandey Date: Tue, 21 Jul 2026 16:28:18 -0600 Subject: [PATCH 02/15] feat(entities): expose grouped counts in list API Signed-off-by: Aditya Pandey --- .../entities/api/v2/entities/endpoints.py | 48 +++++++++++-------- .../integration/test_generic_entities.py | 36 ++++++++++++++ 2 files changed, 65 insertions(+), 19 deletions(-) diff --git a/services/core/entities/src/nmp/core/entities/api/v2/entities/endpoints.py b/services/core/entities/src/nmp/core/entities/api/v2/entities/endpoints.py index 80d8842e3a..9654b5d2bb 100644 --- a/services/core/entities/src/nmp/core/entities/api/v2/entities/endpoints.py +++ b/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 scalar 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,29 @@ 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, + 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 = ( + await repository.count_entities_by( + workspace=query_workspace, entity_type=entity_type, - page=page, - page_size=page_size, - sort=sort, - filter_op=filter, + group_by=count_by, + filter_op=effective_filter, relationship_child_workspaces=accessible_workspaces, ) + if count_by + else None + ) return EntitiesPage( data=entities, @@ -384,6 +393,7 @@ async def list_entities( ), sort=sort, filter=filter.to_dict() if filter else None, + group_counts=group_counts, ) diff --git a/services/core/entities/tests/integration/test_generic_entities.py b/services/core/entities/tests/integration/test_generic_entities.py index 6175090916..3d1e66fac3 100644 --- a/services/core/entities/tests/integration/test_generic_entities.py +++ b/services/core/entities/tests/integration/test_generic_entities.py @@ -3,6 +3,8 @@ """Integration tests for generic entity API v2 endpoints.""" +import json + import pytest from httpx import AsyncClient @@ -152,6 +154,40 @@ async def test_list_entities_with_pagination(self, client: AsyncClient, ctx): assert result["pagination"]["total_results"] == 5 assert result["pagination"]["total_pages"] == 3 + async def test_list_entities_returns_group_counts_for_filtered_field(self, client: AsyncClient, ctx): + """Test list results include counts for a requested filtered data field.""" + entities = [ + {"name": "group-count-a-1", "data": {"insight_id": "insight-a", "is_deleted": False}}, + {"name": "group-count-a-2", "data": {"insight_id": "insight-a", "is_deleted": False}}, + {"name": "group-count-b-1", "data": {"insight_id": "insight-b", "is_deleted": False}}, + {"name": "group-count-deleted", "data": {"insight_id": "insight-a", "is_deleted": True}}, + ] + for entity in entities: + response = await client.post( + "/apis/entities/v2/workspaces/default/entities/experiment_group", + json=entity, + ) + assert response.status_code == 201 + + response = await client.get( + "/apis/entities/v2/workspaces/default/entities/experiment_group", + params={ + "count_by": "data.insight_id", + "filter": json.dumps( + { + "data.insight_id": {"$in": ["insight-a", "insight-b"]}, + "data.is_deleted": False, + } + ), + "page_size": 1, + }, + ) + + assert response.status_code == 200 + result = response.json() + assert result["group_counts"] == {"insight-a": 2, "insight-b": 1} + assert len(result["data"]) == 1 + async def test_update_entity_by_name(self, client: AsyncClient, ctx): """Test updating an entity by name.""" await client.post( From 0cc817120146b37723511dca9843974858c3ea2a Mon Sep 17 00:00:00 2001 From: Aditya Pandey Date: Tue, 21 Jul 2026 16:30:26 -0600 Subject: [PATCH 03/15] feat(entities): add client count helper Signed-off-by: Aditya Pandey --- .../src/nemo_platform_plugin/entities.py | 99 +++++++++++++------ .../tests/test_entity_client.py | 65 ++++++++++++ 2 files changed, 132 insertions(+), 32 deletions(-) create mode 100644 packages/nemo_platform_plugin/tests/test_entity_client.py diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/entities.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/entities.py index 56dde5bd12..3cac16282d 100644 --- a/packages/nemo_platform_plugin/src/nemo_platform_plugin/entities.py +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/entities.py @@ -233,6 +233,16 @@ 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_operation: FilterOperation | None = None, + filter_str: str | None = None, + 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: ... @@ -292,6 +302,13 @@ def _get_entity_type(entity_class: EntityTypeLike) -> str: return "".join(["_" + c.lower() if c.isupper() else c for c in name]).lstrip("_") +def _convert_field_to_api_field(field: str) -> str: + """Convert an entity field name to its API field name.""" + if field in BASE_FIELDS or field.startswith("data."): + return field + return f"data.{field}" + + def _convert_filter_obj_to_filter_str(filter_obj: Dict[str, Any]) -> Dict[str, Any]: """Convert a filter dict to API filter format. @@ -301,16 +318,7 @@ def _convert_filter_obj_to_filter_str(filter_obj: Dict[str, Any]) -> Dict[str, A filter_dict: Dict[str, Any] = {} for field, value in filter_obj.items(): - # Base fields don't need the data. prefix - if field in BASE_FIELDS: - api_field = field - # Already has data. prefix - don't double-prefix - elif field.startswith("data."): - api_field = field - # All other fields are stored in the data JSON column - else: - api_field = f"data.{field}" - filter_dict[api_field] = value + filter_dict[_convert_field_to_api_field(field)] = value return filter_dict @@ -320,11 +328,30 @@ def _convert_sort_to_api_sort(sort: str) -> str: For EntityBase entities, fields are stored in the data JSON column, so we prefix them with 'data.' unless they're base fields. """ - field = sort.lstrip("-") - if field not in BASE_FIELDS: - return f"{'-' if sort.startswith('-') else ''}data.{field}" + direction = "-" if sort.startswith("-") else "" + return f"{direction}{_convert_field_to_api_field(sort.lstrip('-'))}" + + +def _get_effective_filter( + filter_operation: FilterOperation | None, + filter_str: str | None, + filter_obj: dict[str, Any] | None, +) -> str | None: + """Resolve structured, string, and shorthand filters into an API filter string.""" + if filter_operation is not None and filter_str is not None: + raise ValueError( + "EntityClient.list: pass either filter_operation or filter_str, not both. " + "Combining them previously silently dropped one — merge into a single filter_operation " + "via ParsedFilter.and_with." + ) - return sort + if filter_operation is not None: + return json.dumps(filter_operation.to_dict()) + if filter_str: + return filter_str + if filter_obj: + return json.dumps(_convert_filter_obj_to_filter_str(filter_obj)) + return None class EntityClient: @@ -450,24 +477,7 @@ async def list( ListResponse containing data and pagination info """ - if filter_operation is not None and filter_str is not None: - raise ValueError( - "EntityClient.list: pass either filter_operation or filter_str, not both. " - "Combining them previously silently dropped one — merge into a single filter_operation " - "via ParsedFilter.and_with." - ) - - if filter_operation is not None: - effective_filter_str = json.dumps(filter_operation.to_dict()) - else: - effective_filter_str = filter_str - - # Build filter string from filter_obj if provided - if filter_obj and not effective_filter_str: - # Convert filter_obj to filter JSON format - filter_dict = _convert_filter_obj_to_filter_str(filter_obj) - if filter_dict: - effective_filter_str = json.dumps(filter_dict) + effective_filter_str = _get_effective_filter(filter_operation, filter_str, filter_obj) response = await self.entities_api.list( _get_entity_type(entity_type), @@ -493,6 +503,31 @@ async def list( return ListResponse(data=entities, pagination=pagination) + async def count_by( + self, + entity_type: EntityTypeLike, + field: str, + *, + workspace: str = DEFAULT_WORKSPACE, + filter_operation: FilterOperation | None = None, + filter_str: str | None = None, + filter_obj: dict[str, Any] | None = None, + ) -> dict[str, int]: + """Return the number of matching entities grouped by ``field``.""" + effective_filter_str = _get_effective_filter(filter_operation, filter_str, filter_obj) + response = await self.entities_api.list( + _get_entity_type(entity_type), + workspace=workspace, + filter=effective_filter_str if effective_filter_str else omit, + page=1, + page_size=1, + extra_query={"count_by": _convert_field_to_api_field(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. diff --git a/packages/nemo_platform_plugin/tests/test_entity_client.py b/packages/nemo_platform_plugin/tests/test_entity_client.py new file mode 100644 index 0000000000..9f1b3b42ad --- /dev/null +++ b/packages/nemo_platform_plugin/tests/test_entity_client.py @@ -0,0 +1,65 @@ +# 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(**extra_fields: object) -> EntitiesPage: + return EntitiesPage.model_construct( + data=[], + pagination=PaginationData(page=1, page_size=1, current_page_size=0, total_pages=0, total_results=0), + **extra_fields, + ) + + +@pytest.mark.asyncio +async def test_list_uses_shorthand_filter_when_filter_string_is_empty() -> None: + mock_api = Mock() + mock_api.list = AsyncMock(return_value=_entities_page()) + client = EntityClient(mock_api) + + await client.list(ExperimentGroup, filter_str="", filter_obj={"insight_id": "insight-a"}) + + assert mock_api.list.await_args.kwargs["filter"] == '{"data.insight_id": "insight-a"}' + + +@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"]}}, + ) + + assert counts == {"insight-a": 2} + mock_api.list.assert_awaited_once_with( + "experiment_group", + workspace="default", + filter='{"data.insight_id": {"$in": ["insight-a"]}}', + page=1, + page_size=1, + 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") From b5254306e7363e3aeb4a9255b72a3e35568f05e2 Mon Sep 17 00:00:00 2001 From: Aditya Pandey Date: Tue, 21 Jul 2026 16:37:05 -0600 Subject: [PATCH 04/15] refactor(entities): simplify grouped count implementation Signed-off-by: Aditya Pandey --- .../src/nemo_platform_plugin/entities.py | 6 +-- .../tests/test_entity_client.py | 4 +- .../app/repository/sqlalchemy/entity.py | 15 +++---- .../repository/test_entity_group_counts.py | 42 ++++++------------- 4 files changed, 22 insertions(+), 45 deletions(-) diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/entities.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/entities.py index 3cac16282d..f8d10d634d 100644 --- a/packages/nemo_platform_plugin/src/nemo_platform_plugin/entities.py +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/entities.py @@ -315,11 +315,7 @@ def _convert_filter_obj_to_filter_str(filter_obj: Dict[str, Any]) -> Dict[str, A For EntityBase entities, fields are stored in the data JSON column, so we prefix them with 'data.' unless they're base fields. """ - filter_dict: Dict[str, Any] = {} - - for field, value in filter_obj.items(): - filter_dict[_convert_field_to_api_field(field)] = value - return filter_dict + return {_convert_field_to_api_field(field): value for field, value in filter_obj.items()} def _convert_sort_to_api_sort(sort: str) -> str: diff --git a/packages/nemo_platform_plugin/tests/test_entity_client.py b/packages/nemo_platform_plugin/tests/test_entity_client.py index 9f1b3b42ad..f49d2dea7d 100644 --- a/packages/nemo_platform_plugin/tests/test_entity_client.py +++ b/packages/nemo_platform_plugin/tests/test_entity_client.py @@ -13,11 +13,11 @@ class ExperimentGroup: __entity_type__ = "experiment_group" -def _entities_page(**extra_fields: object) -> EntitiesPage: +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), - **extra_fields, + group_counts=group_counts, ) diff --git a/services/core/entities/src/nmp/core/entities/app/repository/sqlalchemy/entity.py b/services/core/entities/src/nmp/core/entities/app/repository/sqlalchemy/entity.py index 51dafca468..7519e2e9a5 100644 --- a/services/core/entities/src/nmp/core/entities/app/repository/sqlalchemy/entity.py +++ b/services/core/entities/src/nmp/core/entities/app/repository/sqlalchemy/entity.py @@ -213,22 +213,19 @@ async def count_entities_by( DBEntity, relationship_child_workspaces=relationship_child_workspaces ) group_column, is_json = filter_repo.get_text_column(group_by) - group_counts = ( - select(group_column, func.count()).select_from(DBEntity).where(DBEntity.entity_type == entity_type) - ) + query = select(group_column, func.count()).select_from(DBEntity).where(DBEntity.entity_type == entity_type) if workspace != ALL_WORKSPACES: - group_counts = group_counts.where(DBEntity.workspace == workspace) + query = query.where(DBEntity.workspace == workspace) if filter_op is not None: - group_counts = group_counts.where(filter_op.apply(filter_repo)) + query = query.where(filter_op.apply(filter_repo)) - group_counts = group_counts.where(group_column.is_not(None)) + query = query.where(group_column.is_not(None)) if is_json: - group_counts = group_counts.where(group_column != "null") - group_counts = group_counts.group_by(group_column) + query = query.where(group_column != "null") - rows = (await sess.execute(group_counts)).all() + rows = (await sess.execute(query.group_by(group_column))).all() return {str(key): int(count) for key, count in rows} async def update_entity( diff --git a/services/core/entities/tests/repository/test_entity_group_counts.py b/services/core/entities/tests/repository/test_entity_group_counts.py index 6b82bc9ed3..e5c90d005f 100644 --- a/services/core/entities/tests/repository/test_entity_group_counts.py +++ b/services/core/entities/tests/repository/test_entity_group_counts.py @@ -14,39 +14,23 @@ async def test_counts_filtered_entities_grouped_by_json_field( entity_repo: SQLAlchemyEntityRepository, setup_workspaces ): """Count only live experiment groups for the requested insights.""" - for name, insight_id in (("a-1", "insight-a"), ("a-2", "insight-a"), ("b-1", "insight-b")): + entities = ( + ("workspace-1", "experiment_group", "a-1", {"insight_id": "insight-a", "is_deleted": False}), + ("workspace-1", "experiment_group", "a-2", {"insight_id": "insight-a", "is_deleted": False}), + ("workspace-1", "experiment_group", "b-1", {"insight_id": "insight-b", "is_deleted": False}), + ("workspace-1", "experiment_group", "deleted", {"insight_id": "insight-a", "is_deleted": True}), + ("workspace-1", "experiment_group", "unlinked", {"is_deleted": False}), + ("workspace-2", "experiment_group", "other-workspace", {"insight_id": "insight-a", "is_deleted": False}), + ("workspace-1", "other_type", "other-type", {"insight_id": "insight-a", "is_deleted": False}), + ) + for workspace, entity_type, name, data in entities: await entity_repo.create_entity( - workspace="workspace-1", - entity_type="experiment_group", + workspace=workspace, + entity_type=entity_type, name=name, - data={"insight_id": insight_id, "is_deleted": False}, + data=data, ) - await entity_repo.create_entity( - workspace="workspace-1", - entity_type="experiment_group", - name="deleted", - data={"insight_id": "insight-a", "is_deleted": True}, - ) - await entity_repo.create_entity( - workspace="workspace-1", - entity_type="experiment_group", - name="unlinked", - data={"is_deleted": False}, - ) - await entity_repo.create_entity( - workspace="workspace-2", - entity_type="experiment_group", - name="other-workspace", - data={"insight_id": "insight-a", "is_deleted": False}, - ) - await entity_repo.create_entity( - workspace="workspace-1", - entity_type="other_type", - name="other-type", - data={"insight_id": "insight-a", "is_deleted": False}, - ) - filter_op = LogicalOperation( operator=FilterOperator.AND, operations=[ From 941a3b1744e4aa13bc0652dc53f93d6485ca5b80 Mon Sep 17 00:00:00 2001 From: Aditya Pandey Date: Tue, 21 Jul 2026 16:39:13 -0600 Subject: [PATCH 05/15] chore(openapi): refresh entity grouped counts contract Signed-off-by: Aditya Pandey --- openapi/ga/individual/platform.openapi.yaml | 13 +++++++++++++ openapi/ga/openapi.yaml | 13 +++++++++++++ openapi/openapi.yaml | 13 +++++++++++++ 3 files changed, 39 insertions(+) diff --git a/openapi/ga/individual/platform.openapi.yaml b/openapi/ga/individual/platform.openapi.yaml index f91c0ade04..1c7ad894bf 100644 --- a/openapi/ga/individual/platform.openapi.yaml +++ b/openapi/ga/individual/platform.openapi.yaml @@ -699,6 +699,14 @@ paths: default: -created_at title: Sort description: Sort field + - name: count_by + in: query + required: false + schema: + description: Optional scalar field whose matching values should be counted. + title: Count By + type: string + description: Optional scalar field whose matching values should be counted. - name: filter in: query required: false @@ -10276,6 +10284,11 @@ components: description: Filtering information. additionalProperties: true type: object + group_counts: + title: Group Counts + additionalProperties: + type: integer + type: object type: object required: - data diff --git a/openapi/ga/openapi.yaml b/openapi/ga/openapi.yaml index f91c0ade04..1c7ad894bf 100644 --- a/openapi/ga/openapi.yaml +++ b/openapi/ga/openapi.yaml @@ -699,6 +699,14 @@ paths: default: -created_at title: Sort description: Sort field + - name: count_by + in: query + required: false + schema: + description: Optional scalar field whose matching values should be counted. + title: Count By + type: string + description: Optional scalar field whose matching values should be counted. - name: filter in: query required: false @@ -10276,6 +10284,11 @@ components: description: Filtering information. additionalProperties: true type: object + group_counts: + title: Group Counts + additionalProperties: + type: integer + type: object type: object required: - data diff --git a/openapi/openapi.yaml b/openapi/openapi.yaml index f91c0ade04..1c7ad894bf 100644 --- a/openapi/openapi.yaml +++ b/openapi/openapi.yaml @@ -699,6 +699,14 @@ paths: default: -created_at title: Sort description: Sort field + - name: count_by + in: query + required: false + schema: + description: Optional scalar field whose matching values should be counted. + title: Count By + type: string + description: Optional scalar field whose matching values should be counted. - name: filter in: query required: false @@ -10276,6 +10284,11 @@ components: description: Filtering information. additionalProperties: true type: object + group_counts: + title: Group Counts + additionalProperties: + type: integer + type: object type: object required: - data From 5ad127f553862794dedc16e85f690c9ff1e78227 Mon Sep 17 00:00:00 2001 From: Aditya Pandey Date: Tue, 21 Jul 2026 16:53:33 -0600 Subject: [PATCH 06/15] fix(entities): harden grouped count results Signed-off-by: Aditya Pandey --- .../src/nemo_platform_plugin/entities.py | 14 +++++- .../tests/test_entity_client.py | 11 ++++ .../entities/api/v2/entities/endpoints.py | 23 +++++---- .../app/repository/sqlalchemy/entity.py | 19 ++++++- .../integration/test_generic_entities.py | 25 ++++++++++ .../repository/test_entity_group_counts.py | 50 +++++++++++++++++++ 6 files changed, 128 insertions(+), 14 deletions(-) diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/entities.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/entities.py index f8d10d634d..e89dfbf31e 100644 --- a/packages/nemo_platform_plugin/src/nemo_platform_plugin/entities.py +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/entities.py @@ -17,7 +17,19 @@ # Regex pattern for valid workspace names ID_PATTERN = r"^[\w\-\+.@:]+$" -BASE_FIELDS = {"id", "name", "workspace", "created_at", "updated_at", "entity_type", "project"} +BASE_FIELDS = { + "id", + "name", + "workspace", + "created_at", + "created_by", + "updated_at", + "updated_by", + "entity_type", + "parent", + "project", + "db_version", +} # Default workspace when none is specified DEFAULT_WORKSPACE = "default" diff --git a/packages/nemo_platform_plugin/tests/test_entity_client.py b/packages/nemo_platform_plugin/tests/test_entity_client.py index f49d2dea7d..a4d6ed366a 100644 --- a/packages/nemo_platform_plugin/tests/test_entity_client.py +++ b/packages/nemo_platform_plugin/tests/test_entity_client.py @@ -55,6 +55,17 @@ async def test_count_by_returns_grouped_counts_for_shorthand_filter() -> None: ) +@pytest.mark.asyncio +async def test_count_by_preserves_top_level_parent_field() -> None: + mock_api = Mock() + mock_api.list = AsyncMock(return_value=_entities_page(group_counts={"parent-id": 1})) + client = EntityClient(mock_api) + + await client.count_by(ExperimentGroup, "parent") + + assert mock_api.list.await_args.kwargs["extra_query"] == {"count_by": "parent"} + + @pytest.mark.asyncio async def test_count_by_rejects_response_without_grouped_counts() -> None: mock_api = Mock() diff --git a/services/core/entities/src/nmp/core/entities/api/v2/entities/endpoints.py b/services/core/entities/src/nmp/core/entities/api/v2/entities/endpoints.py index 9654b5d2bb..7d95deeb9d 100644 --- a/services/core/entities/src/nmp/core/entities/api/v2/entities/endpoints.py +++ b/services/core/entities/src/nmp/core/entities/api/v2/entities/endpoints.py @@ -370,17 +370,18 @@ async def list_entities( filter_op=effective_filter, relationship_child_workspaces=accessible_workspaces, ) - 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, - ) - if count_by - else None - ) + group_counts = None + if count_by: + 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, diff --git a/services/core/entities/src/nmp/core/entities/app/repository/sqlalchemy/entity.py b/services/core/entities/src/nmp/core/entities/app/repository/sqlalchemy/entity.py index 7519e2e9a5..fa398d00a9 100644 --- a/services/core/entities/src/nmp/core/entities/app/repository/sqlalchemy/entity.py +++ b/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 case, 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.""" @@ -213,6 +215,14 @@ async def count_entities_by( DBEntity, relationship_child_workspaces=relationship_child_workspaces ) group_column, is_json = filter_repo.get_text_column(group_by) + if is_json and self._is_sqlite(sess): + json_path = "$." + ".".join(group_by.split(".")[1:]) + json_type = func.json_type(DBEntity.data, json_path) + group_column = case( + (json_type == "true", "true"), + (json_type == "false", "false"), + else_=group_column, + ) query = select(group_column, func.count()).select_from(DBEntity).where(DBEntity.entity_type == entity_type) if workspace != ALL_WORKSPACES: @@ -225,7 +235,12 @@ async def count_entities_by( if is_json: query = query.where(group_column != "null") - rows = (await sess.execute(query.group_by(group_column))).all() + 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( diff --git a/services/core/entities/tests/integration/test_generic_entities.py b/services/core/entities/tests/integration/test_generic_entities.py index 3d1e66fac3..275a6b70b5 100644 --- a/services/core/entities/tests/integration/test_generic_entities.py +++ b/services/core/entities/tests/integration/test_generic_entities.py @@ -4,6 +4,7 @@ """Integration tests for generic entity API v2 endpoints.""" import json +from unittest.mock import AsyncMock import pytest from httpx import AsyncClient @@ -188,6 +189,30 @@ async def test_list_entities_returns_group_counts_for_filtered_field(self, clien assert result["group_counts"] == {"insight-a": 2, "insight-b": 1} assert len(result["data"]) == 1 + async def test_list_entities_rejects_invalid_count_field(self, client: AsyncClient, ctx): + response = await client.get( + "/apis/entities/v2/workspaces/default/entities/experiment_group", + params={"count_by": "not_a_field"}, + ) + + assert response.status_code == 400 + assert "does not exist" in response.json()["detail"] + + async def test_list_entities_translates_group_count_overflow(self, client: AsyncClient, ctx, repos, monkeypatch): + monkeypatch.setattr( + repos["entity"], + "count_entities_by", + AsyncMock(side_effect=ValueError("Grouped count has more than 1000 distinct values")), + ) + + response = await client.get( + "/apis/entities/v2/workspaces/default/entities/experiment_group", + params={"count_by": "data.value"}, + ) + + assert response.status_code == 400 + assert response.json()["detail"] == "Grouped count has more than 1000 distinct values" + async def test_update_entity_by_name(self, client: AsyncClient, ctx): """Test updating an entity by name.""" await client.post( diff --git a/services/core/entities/tests/repository/test_entity_group_counts.py b/services/core/entities/tests/repository/test_entity_group_counts.py index e5c90d005f..edaf0a2d34 100644 --- a/services/core/entities/tests/repository/test_entity_group_counts.py +++ b/services/core/entities/tests/repository/test_entity_group_counts.py @@ -6,6 +6,7 @@ import pytest from nmp.common.api.filter import ComparisonOperation, FilterOperator, LogicalOperation from nmp.core.entities.app.repository import SQLAlchemyEntityRepository +from nmp.core.entities.app.repository.sqlalchemy.models import DBEntity pytestmark = pytest.mark.asyncio @@ -86,3 +87,52 @@ async def test_omits_missing_and_null_json_group_keys(entity_repo: SQLAlchemyEnt ) assert counts == {"insight-a": 1} + + +async def test_normalizes_sqlite_json_booleans_without_changing_numbers( + entity_repo: SQLAlchemyEntityRepository, setup_workspaces +): + for name, value in ( + ("true", True), + ("false", False), + ("one", 1), + ("zero", 0), + ): + await entity_repo.create_entity( + workspace="workspace-1", + entity_type="experiment_group", + name=name, + data={"value": value}, + ) + + counts = await entity_repo.count_entities_by( + workspace="workspace-1", + entity_type="experiment_group", + group_by="data.value", + ) + + assert counts == {"true": 1, "false": 1, "1": 1, "0": 1} + + +async def test_rejects_more_than_1000_group_values( + entity_repo: SQLAlchemyEntityRepository, session_maker, setup_workspaces +): + async with session_maker() as session: + session.add_all( + DBEntity( + id=f"experiment-group-{index}", + workspace="workspace-1", + entity_type="experiment_group", + name=f"group-{index}", + data={"value": f"group-{index}"}, + ) + for index in range(1001) + ) + await session.commit() + + with pytest.raises(ValueError, match="more than 1000 distinct values"): + await entity_repo.count_entities_by( + workspace="workspace-1", + entity_type="experiment_group", + group_by="data.value", + ) From 64ed8cadee3248b949e28d12b8291c4406796c01 Mon Sep 17 00:00:00 2001 From: Aditya Pandey Date: Tue, 21 Jul 2026 16:59:45 -0600 Subject: [PATCH 07/15] fix(entities): validate grouped count fields Signed-off-by: Aditya Pandey --- .../core/entities/app/repository/sqlalchemy/entity.py | 10 +++++++++- .../tests/integration/test_generic_entities.py | 5 +++-- 2 files changed, 12 insertions(+), 3 deletions(-) diff --git a/services/core/entities/src/nmp/core/entities/app/repository/sqlalchemy/entity.py b/services/core/entities/src/nmp/core/entities/app/repository/sqlalchemy/entity.py index fa398d00a9..136d0c182a 100644 --- a/services/core/entities/src/nmp/core/entities/app/repository/sqlalchemy/entity.py +++ b/services/core/entities/src/nmp/core/entities/app/repository/sqlalchemy/entity.py @@ -15,7 +15,7 @@ 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 case, func, select +from sqlalchemy import JSON, case, func, select from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker from sqlalchemy.orm.exc import StaleDataError @@ -211,6 +211,14 @@ async def count_entities_by( ) -> dict[str, int]: """Count filtered entities grouped by a scalar field.""" async with self._get_session(session) as sess: + if group_by.startswith("data."): + if not all(group_by.split(".")[1:]): + raise ValueError(f"Field '{group_by}' does not exist on model DBEntity") + else: + column = DBEntity.__table__.columns.get(group_by) + if column is None or isinstance(column.type, JSON): + raise ValueError(f"Field '{group_by}' does not exist on model DBEntity") + filter_repo = SQLAlchemyFilterRepository( DBEntity, relationship_child_workspaces=relationship_child_workspaces ) diff --git a/services/core/entities/tests/integration/test_generic_entities.py b/services/core/entities/tests/integration/test_generic_entities.py index 275a6b70b5..76278d4cec 100644 --- a/services/core/entities/tests/integration/test_generic_entities.py +++ b/services/core/entities/tests/integration/test_generic_entities.py @@ -189,10 +189,11 @@ async def test_list_entities_returns_group_counts_for_filtered_field(self, clien assert result["group_counts"] == {"insight-a": 2, "insight-b": 1} assert len(result["data"]) == 1 - async def test_list_entities_rejects_invalid_count_field(self, client: AsyncClient, ctx): + @pytest.mark.parametrize("count_by", ["not_a_field", "metadata", "data."]) + async def test_list_entities_rejects_invalid_count_field(self, client: AsyncClient, ctx, count_by: str): response = await client.get( "/apis/entities/v2/workspaces/default/entities/experiment_group", - params={"count_by": "not_a_field"}, + params={"count_by": count_by}, ) assert response.status_code == 400 From 3f470bba1b7f4651ccee3cc4d323f08311532399 Mon Sep 17 00:00:00 2001 From: Aditya Pandey Date: Tue, 21 Jul 2026 17:10:14 -0600 Subject: [PATCH 08/15] fix(entities): preserve null string group keys Signed-off-by: Aditya Pandey --- .../app/repository/sqlalchemy/entity.py | 26 +++++++++++-------- .../app/repository/sqlalchemy/filter.py | 6 ++--- .../integration/test_generic_entities.py | 2 +- .../repository/test_entity_group_counts.py | 19 ++++++++++++-- 4 files changed, 36 insertions(+), 17 deletions(-) diff --git a/services/core/entities/src/nmp/core/entities/app/repository/sqlalchemy/entity.py b/services/core/entities/src/nmp/core/entities/app/repository/sqlalchemy/entity.py index 136d0c182a..a280fce6c3 100644 --- a/services/core/entities/src/nmp/core/entities/app/repository/sqlalchemy/entity.py +++ b/services/core/entities/src/nmp/core/entities/app/repository/sqlalchemy/entity.py @@ -222,15 +222,18 @@ async def count_entities_by( filter_repo = SQLAlchemyFilterRepository( DBEntity, relationship_child_workspaces=relationship_child_workspaces ) - group_column, is_json = filter_repo.get_text_column(group_by) - if is_json and self._is_sqlite(sess): - json_path = "$." + ".".join(group_by.split(".")[1:]) - json_type = func.json_type(DBEntity.data, json_path) - group_column = case( - (json_type == "true", "true"), - (json_type == "false", "false"), - else_=group_column, - ) + group_column, raw_group_column, is_json = filter_repo.get_text_column(group_by) + if is_json: + if self._is_sqlite(sess): + json_path = "$." + ".".join(group_by.split(".")[1:]) + json_type = func.json_type(DBEntity.data, json_path) + group_column = case( + (json_type == "true", "true"), + (json_type == "false", "false"), + else_=group_column, + ) + else: + json_type = func.json_typeof(raw_group_column) query = select(group_column, func.count()).select_from(DBEntity).where(DBEntity.entity_type == entity_type) if workspace != ALL_WORKSPACES: @@ -239,9 +242,10 @@ async def count_entities_by( if filter_op is not None: query = query.where(filter_op.apply(filter_repo)) - query = query.where(group_column.is_not(None)) if is_json: - query = query.where(group_column != "null") + query = query.where(json_type.is_not(None), json_type != "null") + else: + query = query.where(group_column.is_not(None)) rows = (await sess.execute(query.group_by(group_column).limit(MAX_GROUP_COUNT_ROWS + 1))).all() if len(rows) > MAX_GROUP_COUNT_ROWS: diff --git a/services/core/entities/src/nmp/core/entities/app/repository/sqlalchemy/filter.py b/services/core/entities/src/nmp/core/entities/app/repository/sqlalchemy/filter.py index 7bc60f69e9..a6add997fa 100644 --- a/services/core/entities/src/nmp/core/entities/app/repository/sqlalchemy/filter.py +++ b/services/core/entities/src/nmp/core/entities/app/repository/sqlalchemy/filter.py @@ -96,11 +96,11 @@ def _cast_json_to_text(self, column: Any) -> Any: # This handles both SQLite and PostgreSQL JSON string extraction return func.trim(cast(column, String), '"') - def get_text_column(self, field: str) -> tuple[ColumnElement, bool]: - """Return a text-normalized field expression and whether it is JSON.""" + def get_text_column(self, field: str) -> tuple[ColumnElement, ColumnElement, bool]: + """Return text-normalized and raw field expressions and whether the field is JSON.""" column, is_json = self._get_column(field) text_column = self._cast_json_to_text(column) if is_json else cast(column, String) - return text_column, is_json + return text_column, column, is_json def _cast_json_to_numeric(self, column: Any) -> Any: """Cast a JSON column element to a float for numeric comparisons. diff --git a/services/core/entities/tests/integration/test_generic_entities.py b/services/core/entities/tests/integration/test_generic_entities.py index 76278d4cec..f00b141523 100644 --- a/services/core/entities/tests/integration/test_generic_entities.py +++ b/services/core/entities/tests/integration/test_generic_entities.py @@ -189,7 +189,7 @@ async def test_list_entities_returns_group_counts_for_filtered_field(self, clien assert result["group_counts"] == {"insight-a": 2, "insight-b": 1} assert len(result["data"]) == 1 - @pytest.mark.parametrize("count_by", ["not_a_field", "metadata", "data."]) + @pytest.mark.parametrize("count_by", ["not_a_field", "metadata"]) async def test_list_entities_rejects_invalid_count_field(self, client: AsyncClient, ctx, count_by: str): response = await client.get( "/apis/entities/v2/workspaces/default/entities/experiment_group", diff --git a/services/core/entities/tests/repository/test_entity_group_counts.py b/services/core/entities/tests/repository/test_entity_group_counts.py index edaf0a2d34..910b165498 100644 --- a/services/core/entities/tests/repository/test_entity_group_counts.py +++ b/services/core/entities/tests/repository/test_entity_group_counts.py @@ -72,6 +72,7 @@ async def test_omits_missing_and_null_json_group_keys(entity_repo: SQLAlchemyEnt ("linked", {"insight_id": "insight-a"}), ("missing", {}), ("null", {"insight_id": None}), + ("string-null", {"insight_id": "null"}), ): await entity_repo.create_entity( workspace="workspace-1", @@ -86,7 +87,7 @@ async def test_omits_missing_and_null_json_group_keys(entity_repo: SQLAlchemyEnt group_by="data.insight_id", ) - assert counts == {"insight-a": 1} + assert counts == {"insight-a": 1, "null": 1} async def test_normalizes_sqlite_json_booleans_without_changing_numbers( @@ -126,10 +127,24 @@ async def test_rejects_more_than_1000_group_values( name=f"group-{index}", data={"value": f"group-{index}"}, ) - for index in range(1001) + for index in range(1000) ) await session.commit() + counts = await entity_repo.count_entities_by( + workspace="workspace-1", + entity_type="experiment_group", + group_by="data.value", + ) + assert len(counts) == 1000 + + await entity_repo.create_entity( + workspace="workspace-1", + entity_type="experiment_group", + name="group-1000", + data={"value": "group-1000"}, + ) + with pytest.raises(ValueError, match="more than 1000 distinct values"): await entity_repo.count_entities_by( workspace="workspace-1", From 91b9d126adb2dd237380b7899e7ada761354d5ba Mon Sep 17 00:00:00 2001 From: Aditya Pandey Date: Wed, 22 Jul 2026 16:30:05 -0600 Subject: [PATCH 09/15] docs: narrow entity grouped count design Signed-off-by: Aditya Pandey --- ...07-22-narrow-entity-group-counts-design.md | 59 +++++++++++++++++++ 1 file changed, 59 insertions(+) create mode 100644 docs/superpowers/specs/2026-07-22-narrow-entity-group-counts-design.md diff --git a/docs/superpowers/specs/2026-07-22-narrow-entity-group-counts-design.md b/docs/superpowers/specs/2026-07-22-narrow-entity-group-counts-design.md new file mode 100644 index 0000000000..e5183f99ed --- /dev/null +++ b/docs/superpowers/specs/2026-07-22-narrow-entity-group-counts-design.md @@ -0,0 +1,59 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# Narrow Entity Group Counts + +## Goal + +Count child entities for several parents in one request without fetching every +matching entity. Current uses are Experiment Groups by `insight_id` and +Evaluations by `experiment_group_id`. + +## API + +Keep grouped counts as optional metadata on the existing entity-list endpoint: + +```text +GET .../entities/{entity_type}?count_by=data.{field}&filter=... +``` + +When `count_by` is absent, list behavior is unchanged. When present, +`group_counts` maps each observed string value to its matching entity count. +Missing values are not returned. + +Expose the narrow client method: + +```python +await entity_client.count_by( + entity_type, + field, + workspace=workspace, + filter_obj=filter_obj, +) +``` + +The client accepts one direct entity data field and adds the `data.` prefix. +It does not accept base fields, nested paths, filter strings, or filter +operation objects. + +## Repository + +Reuse existing workspace authorization and filter application. Accept only a +`data.` path with exactly one segment below `data`, and include only +string values. Ignore missing, JSON-null, boolean, numeric, array, and object +values. Keep the 1,000-group bound. The endpoint returns HTTP 400 for an +unsupported field or bound overflow. + +The list-endpoint integration intentionally retains its normal page and total +queries. A dedicated count endpoint is excluded because it increases API and +SDK surface beyond the current need. + +## Scope Control + +Remove support and tests for base-column grouping, nested JSON paths, boolean +and numeric normalization, and expanded client base-field mappings. Avoid +refactoring existing list/filter behavior solely to share code with +`count_by`. + +Test only the client request mapping, filtered direct-string grouping, +unsupported group fields, cardinality overflow, and endpoint response wiring. From 22da54b2da01a7235f1c2baa671704396338f497 Mon Sep 17 00:00:00 2001 From: Aditya Pandey Date: Wed, 22 Jul 2026 16:45:19 -0600 Subject: [PATCH 10/15] refactor(entities): narrow grouped counts Signed-off-by: Aditya Pandey --- .../entities/api/v2/entities/endpoints.py | 2 +- .../app/repository/sqlalchemy/entity.py | 39 +++--- .../app/repository/sqlalchemy/filter.py | 6 - .../integration/test_generic_entities.py | 23 +--- .../repository/test_entity_group_counts.py | 122 ++++-------------- 5 files changed, 44 insertions(+), 148 deletions(-) diff --git a/services/core/entities/src/nmp/core/entities/api/v2/entities/endpoints.py b/services/core/entities/src/nmp/core/entities/api/v2/entities/endpoints.py index 7d95deeb9d..cb3c4b6c13 100644 --- a/services/core/entities/src/nmp/core/entities/api/v2/entities/endpoints.py +++ b/services/core/entities/src/nmp/core/entities/api/v2/entities/endpoints.py @@ -339,7 +339,7 @@ async def list_entities( ), count_by: str | None = Query( default=None, - description="Optional scalar field whose matching values should be counted.", + description="Optional direct string data field whose matching values should be counted.", ), ) -> EntitiesPage: """List entities with filtering, supporting cross-workspace queries.""" diff --git a/services/core/entities/src/nmp/core/entities/app/repository/sqlalchemy/entity.py b/services/core/entities/src/nmp/core/entities/app/repository/sqlalchemy/entity.py index a280fce6c3..f0b948c247 100644 --- a/services/core/entities/src/nmp/core/entities/app/repository/sqlalchemy/entity.py +++ b/services/core/entities/src/nmp/core/entities/app/repository/sqlalchemy/entity.py @@ -15,7 +15,7 @@ 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 JSON, case, func, select +from sqlalchemy import String, cast, func, select from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker from sqlalchemy.orm.exc import StaleDataError @@ -209,31 +209,25 @@ async def count_entities_by( relationship_child_workspaces: set[str] | None = None, session: AsyncSession | None = None, ) -> dict[str, int]: - """Count filtered entities grouped by a scalar field.""" + """Count filtered entities grouped by a direct string data field.""" async with self._get_session(session) as sess: - if group_by.startswith("data."): - if not all(group_by.split(".")[1:]): - raise ValueError(f"Field '{group_by}' does not exist on model DBEntity") + 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: - column = DBEntity.__table__.columns.get(group_by) - if column is None or isinstance(column.type, JSON): - raise ValueError(f"Field '{group_by}' does not exist on model DBEntity") + json_type = func.json_typeof(raw_group_column) + string_type = "string" filter_repo = SQLAlchemyFilterRepository( DBEntity, relationship_child_workspaces=relationship_child_workspaces ) - group_column, raw_group_column, is_json = filter_repo.get_text_column(group_by) - if is_json: - if self._is_sqlite(sess): - json_path = "$." + ".".join(group_by.split(".")[1:]) - json_type = func.json_type(DBEntity.data, json_path) - group_column = case( - (json_type == "true", "true"), - (json_type == "false", "false"), - else_=group_column, - ) - else: - json_type = func.json_typeof(raw_group_column) query = select(group_column, func.count()).select_from(DBEntity).where(DBEntity.entity_type == entity_type) if workspace != ALL_WORKSPACES: @@ -242,10 +236,7 @@ async def count_entities_by( if filter_op is not None: query = query.where(filter_op.apply(filter_repo)) - if is_json: - query = query.where(json_type.is_not(None), json_type != "null") - else: - query = query.where(group_column.is_not(None)) + 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: diff --git a/services/core/entities/src/nmp/core/entities/app/repository/sqlalchemy/filter.py b/services/core/entities/src/nmp/core/entities/app/repository/sqlalchemy/filter.py index a6add997fa..40eb859d82 100644 --- a/services/core/entities/src/nmp/core/entities/app/repository/sqlalchemy/filter.py +++ b/services/core/entities/src/nmp/core/entities/app/repository/sqlalchemy/filter.py @@ -96,12 +96,6 @@ def _cast_json_to_text(self, column: Any) -> Any: # This handles both SQLite and PostgreSQL JSON string extraction return func.trim(cast(column, String), '"') - def get_text_column(self, field: str) -> tuple[ColumnElement, ColumnElement, bool]: - """Return text-normalized and raw field expressions and whether the field is JSON.""" - column, is_json = self._get_column(field) - text_column = self._cast_json_to_text(column) if is_json else cast(column, String) - return text_column, column, is_json - def _cast_json_to_numeric(self, column: Any) -> Any: """Cast a JSON column element to a float for numeric comparisons. diff --git a/services/core/entities/tests/integration/test_generic_entities.py b/services/core/entities/tests/integration/test_generic_entities.py index f00b141523..480e34cd69 100644 --- a/services/core/entities/tests/integration/test_generic_entities.py +++ b/services/core/entities/tests/integration/test_generic_entities.py @@ -4,7 +4,6 @@ """Integration tests for generic entity API v2 endpoints.""" import json -from unittest.mock import AsyncMock import pytest from httpx import AsyncClient @@ -189,30 +188,14 @@ async def test_list_entities_returns_group_counts_for_filtered_field(self, clien assert result["group_counts"] == {"insight-a": 2, "insight-b": 1} assert len(result["data"]) == 1 - @pytest.mark.parametrize("count_by", ["not_a_field", "metadata"]) - async def test_list_entities_rejects_invalid_count_field(self, client: AsyncClient, ctx, count_by: str): + async def test_list_entities_rejects_unsupported_count_field(self, client: AsyncClient, ctx): response = await client.get( "/apis/entities/v2/workspaces/default/entities/experiment_group", - params={"count_by": count_by}, + params={"count_by": "name"}, ) assert response.status_code == 400 - assert "does not exist" in response.json()["detail"] - - async def test_list_entities_translates_group_count_overflow(self, client: AsyncClient, ctx, repos, monkeypatch): - monkeypatch.setattr( - repos["entity"], - "count_entities_by", - AsyncMock(side_effect=ValueError("Grouped count has more than 1000 distinct values")), - ) - - response = await client.get( - "/apis/entities/v2/workspaces/default/entities/experiment_group", - params={"count_by": "data.value"}, - ) - - assert response.status_code == 400 - assert response.json()["detail"] == "Grouped count has more than 1000 distinct values" + assert "direct string data field" in response.json()["detail"] async def test_update_entity_by_name(self, client: AsyncClient, ctx): """Test updating an entity by name.""" diff --git a/services/core/entities/tests/repository/test_entity_group_counts.py b/services/core/entities/tests/repository/test_entity_group_counts.py index 910b165498..6a690a53d4 100644 --- a/services/core/entities/tests/repository/test_entity_group_counts.py +++ b/services/core/entities/tests/repository/test_entity_group_counts.py @@ -6,28 +6,29 @@ import pytest from nmp.common.api.filter import ComparisonOperation, FilterOperator, LogicalOperation from nmp.core.entities.app.repository import SQLAlchemyEntityRepository -from nmp.core.entities.app.repository.sqlalchemy.models import DBEntity +from nmp.core.entities.app.repository.sqlalchemy import entity as entity_repository pytestmark = pytest.mark.asyncio -async def test_counts_filtered_entities_grouped_by_json_field( +async def test_counts_filtered_entities_grouped_by_direct_string_data_field( entity_repo: SQLAlchemyEntityRepository, setup_workspaces ): """Count only live experiment groups for the requested insights.""" entities = ( - ("workspace-1", "experiment_group", "a-1", {"insight_id": "insight-a", "is_deleted": False}), - ("workspace-1", "experiment_group", "a-2", {"insight_id": "insight-a", "is_deleted": False}), - ("workspace-1", "experiment_group", "b-1", {"insight_id": "insight-b", "is_deleted": False}), - ("workspace-1", "experiment_group", "deleted", {"insight_id": "insight-a", "is_deleted": True}), - ("workspace-1", "experiment_group", "unlinked", {"is_deleted": False}), - ("workspace-2", "experiment_group", "other-workspace", {"insight_id": "insight-a", "is_deleted": False}), - ("workspace-1", "other_type", "other-type", {"insight_id": "insight-a", "is_deleted": False}), + ("live-a-1", {"insight_id": "insight-a", "is_deleted": False}), + ("live-a-2", {"insight_id": "insight-a", "is_deleted": False}), + ("live-b", {"insight_id": "insight-b", "is_deleted": False}), + ("deleted", {"insight_id": "insight-a", "is_deleted": True}), + ("missing", {"is_deleted": False}), + ("null", {"insight_id": None, "is_deleted": False}), + ("boolean", {"insight_id": True, "is_deleted": False}), + ("numeric", {"insight_id": 1, "is_deleted": False}), ) - for workspace, entity_type, name, data in entities: + for name, data in entities: await entity_repo.create_entity( - workspace=workspace, - entity_type=entity_type, + workspace="workspace-1", + entity_type="experiment_group", name=name, data=data, ) @@ -50,102 +51,29 @@ async def test_counts_filtered_entities_grouped_by_json_field( assert counts == {"insight-a": 2, "insight-b": 1} -async def test_counts_base_field_literal_null(entity_repo: SQLAlchemyEntityRepository, setup_workspaces): - await entity_repo.create_entity( - workspace="workspace-1", - entity_type="experiment_group", - name="null", - data={}, - ) - - counts = await entity_repo.count_entities_by( - workspace="workspace-1", - entity_type="experiment_group", - group_by="name", - ) - - assert counts == {"null": 1} - - -async def test_omits_missing_and_null_json_group_keys(entity_repo: SQLAlchemyEntityRepository, setup_workspaces): - for name, data in ( - ("linked", {"insight_id": "insight-a"}), - ("missing", {}), - ("null", {"insight_id": None}), - ("string-null", {"insight_id": "null"}), - ): - await entity_repo.create_entity( +@pytest.mark.parametrize("field", ["name", "data.nested.value", "data.", "data.not-valid"]) +async def test_rejects_unsupported_group_fields(entity_repo: SQLAlchemyEntityRepository, setup_workspaces, field: str): + with pytest.raises(ValueError, match="direct string data field"): + await entity_repo.count_entities_by( workspace="workspace-1", entity_type="experiment_group", - name=name, - data=data, + group_by=field, ) - counts = await entity_repo.count_entities_by( - workspace="workspace-1", - entity_type="experiment_group", - group_by="data.insight_id", - ) - - assert counts == {"insight-a": 1, "null": 1} - -async def test_normalizes_sqlite_json_booleans_without_changing_numbers( - entity_repo: SQLAlchemyEntityRepository, setup_workspaces +async def test_rejects_group_counts_over_limit( + entity_repo: SQLAlchemyEntityRepository, setup_workspaces, monkeypatch: pytest.MonkeyPatch ): - for name, value in ( - ("true", True), - ("false", False), - ("one", 1), - ("zero", 0), - ): + monkeypatch.setattr(entity_repository, "MAX_GROUP_COUNT_ROWS", 2) + for index in range(3): await entity_repo.create_entity( workspace="workspace-1", entity_type="experiment_group", - name=name, - data={"value": value}, - ) - - counts = await entity_repo.count_entities_by( - workspace="workspace-1", - entity_type="experiment_group", - group_by="data.value", - ) - - assert counts == {"true": 1, "false": 1, "1": 1, "0": 1} - - -async def test_rejects_more_than_1000_group_values( - entity_repo: SQLAlchemyEntityRepository, session_maker, setup_workspaces -): - async with session_maker() as session: - session.add_all( - DBEntity( - id=f"experiment-group-{index}", - workspace="workspace-1", - entity_type="experiment_group", - name=f"group-{index}", - data={"value": f"group-{index}"}, - ) - for index in range(1000) + name=f"group-{index}", + data={"value": f"group-{index}"}, ) - await session.commit() - - counts = await entity_repo.count_entities_by( - workspace="workspace-1", - entity_type="experiment_group", - group_by="data.value", - ) - assert len(counts) == 1000 - - await entity_repo.create_entity( - workspace="workspace-1", - entity_type="experiment_group", - name="group-1000", - data={"value": "group-1000"}, - ) - with pytest.raises(ValueError, match="more than 1000 distinct values"): + with pytest.raises(ValueError, match="more than 2 distinct values"): await entity_repo.count_entities_by( workspace="workspace-1", entity_type="experiment_group", From 58ea3a057516c05af20487a44e4d8617a83b1793 Mon Sep 17 00:00:00 2001 From: Aditya Pandey Date: Wed, 22 Jul 2026 16:48:22 -0600 Subject: [PATCH 11/15] fix(entities): validate empty count field Signed-off-by: Aditya Pandey --- .../src/nmp/core/entities/api/v2/entities/endpoints.py | 2 +- .../core/entities/tests/integration/test_generic_entities.py | 5 +++-- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/services/core/entities/src/nmp/core/entities/api/v2/entities/endpoints.py b/services/core/entities/src/nmp/core/entities/api/v2/entities/endpoints.py index cb3c4b6c13..d0fe28e5c3 100644 --- a/services/core/entities/src/nmp/core/entities/api/v2/entities/endpoints.py +++ b/services/core/entities/src/nmp/core/entities/api/v2/entities/endpoints.py @@ -371,7 +371,7 @@ async def list_entities( relationship_child_workspaces=accessible_workspaces, ) group_counts = None - if count_by: + if count_by is not None: try: group_counts = await repository.count_entities_by( workspace=query_workspace, diff --git a/services/core/entities/tests/integration/test_generic_entities.py b/services/core/entities/tests/integration/test_generic_entities.py index 480e34cd69..d71a199b9a 100644 --- a/services/core/entities/tests/integration/test_generic_entities.py +++ b/services/core/entities/tests/integration/test_generic_entities.py @@ -188,10 +188,11 @@ async def test_list_entities_returns_group_counts_for_filtered_field(self, clien assert result["group_counts"] == {"insight-a": 2, "insight-b": 1} assert len(result["data"]) == 1 - async def test_list_entities_rejects_unsupported_count_field(self, client: AsyncClient, ctx): + @pytest.mark.parametrize("count_by", ["name", ""]) + async def test_list_entities_rejects_unsupported_count_field(self, client: AsyncClient, ctx, count_by: str): response = await client.get( "/apis/entities/v2/workspaces/default/entities/experiment_group", - params={"count_by": "name"}, + params={"count_by": count_by}, ) assert response.status_code == 400 From 963263acb96fd77ad35e9221c3a1e46831f275f1 Mon Sep 17 00:00:00 2001 From: Aditya Pandey Date: Wed, 22 Jul 2026 16:50:06 -0600 Subject: [PATCH 12/15] refactor(entities): narrow grouped count client Signed-off-by: Aditya Pandey --- .../src/nemo_platform_plugin/entities.py | 97 +++++++++---------- .../tests/test_entity_client.py | 40 +++----- 2 files changed, 57 insertions(+), 80 deletions(-) diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/entities.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/entities.py index e89dfbf31e..0351d6e3e0 100644 --- a/packages/nemo_platform_plugin/src/nemo_platform_plugin/entities.py +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/entities.py @@ -17,19 +17,7 @@ # Regex pattern for valid workspace names ID_PATTERN = r"^[\w\-\+.@:]+$" -BASE_FIELDS = { - "id", - "name", - "workspace", - "created_at", - "created_by", - "updated_at", - "updated_by", - "entity_type", - "parent", - "project", - "db_version", -} +BASE_FIELDS = {"id", "name", "workspace", "created_at", "updated_at", "entity_type", "project"} # Default workspace when none is specified DEFAULT_WORKSPACE = "default" @@ -251,8 +239,6 @@ async def count_by( field: str, *, workspace: str = DEFAULT_WORKSPACE, - filter_operation: FilterOperation | None = None, - filter_str: str | None = None, filter_obj: dict[str, Any] | None = None, ) -> dict[str, int]: ... async def get(self, entity_type: EntityTypeLike, name: str, *, workspace: Optional[str] = None) -> EntityT: ... @@ -314,20 +300,26 @@ def _get_entity_type(entity_class: EntityTypeLike) -> str: return "".join(["_" + c.lower() if c.isupper() else c for c in name]).lstrip("_") -def _convert_field_to_api_field(field: str) -> str: - """Convert an entity field name to its API field name.""" - if field in BASE_FIELDS or field.startswith("data."): - return field - return f"data.{field}" - - def _convert_filter_obj_to_filter_str(filter_obj: Dict[str, Any]) -> Dict[str, Any]: """Convert a filter dict to API filter format. For EntityBase entities, fields are stored in the data JSON column, so we prefix them with 'data.' unless they're base fields. """ - return {_convert_field_to_api_field(field): value for field, value in filter_obj.items()} + filter_dict: Dict[str, Any] = {} + + for field, value in filter_obj.items(): + # Base fields don't need the data. prefix + if field in BASE_FIELDS: + api_field = field + # Already has data. prefix - don't double-prefix + elif field.startswith("data."): + api_field = field + # All other fields are stored in the data JSON column + else: + api_field = f"data.{field}" + filter_dict[api_field] = value + return filter_dict def _convert_sort_to_api_sort(sort: str) -> str: @@ -336,30 +328,11 @@ def _convert_sort_to_api_sort(sort: str) -> str: For EntityBase entities, fields are stored in the data JSON column, so we prefix them with 'data.' unless they're base fields. """ - direction = "-" if sort.startswith("-") else "" - return f"{direction}{_convert_field_to_api_field(sort.lstrip('-'))}" - - -def _get_effective_filter( - filter_operation: FilterOperation | None, - filter_str: str | None, - filter_obj: dict[str, Any] | None, -) -> str | None: - """Resolve structured, string, and shorthand filters into an API filter string.""" - if filter_operation is not None and filter_str is not None: - raise ValueError( - "EntityClient.list: pass either filter_operation or filter_str, not both. " - "Combining them previously silently dropped one — merge into a single filter_operation " - "via ParsedFilter.and_with." - ) + field = sort.lstrip("-") + if field not in BASE_FIELDS: + return f"{'-' if sort.startswith('-') else ''}data.{field}" - if filter_operation is not None: - return json.dumps(filter_operation.to_dict()) - if filter_str: - return filter_str - if filter_obj: - return json.dumps(_convert_filter_obj_to_filter_str(filter_obj)) - return None + return sort class EntityClient: @@ -485,7 +458,24 @@ async def list( ListResponse containing data and pagination info """ - effective_filter_str = _get_effective_filter(filter_operation, filter_str, filter_obj) + if filter_operation is not None and filter_str is not None: + raise ValueError( + "EntityClient.list: pass either filter_operation or filter_str, not both. " + "Combining them previously silently dropped one — merge into a single filter_operation " + "via ParsedFilter.and_with." + ) + + if filter_operation is not None: + effective_filter_str = json.dumps(filter_operation.to_dict()) + else: + effective_filter_str = filter_str + + # Build filter string from filter_obj if provided + if filter_obj and not effective_filter_str: + # Convert filter_obj to filter JSON format + filter_dict = _convert_filter_obj_to_filter_str(filter_obj) + if filter_dict: + effective_filter_str = json.dumps(filter_dict) response = await self.entities_api.list( _get_entity_type(entity_type), @@ -517,19 +507,22 @@ async def count_by( field: str, *, workspace: str = DEFAULT_WORKSPACE, - filter_operation: FilterOperation | None = None, - filter_str: str | None = None, filter_obj: dict[str, Any] | None = None, ) -> dict[str, int]: """Return the number of matching entities grouped by ``field``.""" - effective_filter_str = _get_effective_filter(filter_operation, filter_str, filter_obj) + 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_str if effective_filter_str else omit, + filter=effective_filter, page=1, page_size=1, - extra_query={"count_by": _convert_field_to_api_field(field)}, + extra_query={"count_by": f"data.{field}"}, ) group_counts = getattr(response, "group_counts", None) if group_counts is None: diff --git a/packages/nemo_platform_plugin/tests/test_entity_client.py b/packages/nemo_platform_plugin/tests/test_entity_client.py index a4d6ed366a..5eb6080c93 100644 --- a/packages/nemo_platform_plugin/tests/test_entity_client.py +++ b/packages/nemo_platform_plugin/tests/test_entity_client.py @@ -21,17 +21,6 @@ def _entities_page(group_counts: dict[str, int] | None = None) -> EntitiesPage: ) -@pytest.mark.asyncio -async def test_list_uses_shorthand_filter_when_filter_string_is_empty() -> None: - mock_api = Mock() - mock_api.list = AsyncMock(return_value=_entities_page()) - client = EntityClient(mock_api) - - await client.list(ExperimentGroup, filter_str="", filter_obj={"insight_id": "insight-a"}) - - assert mock_api.list.await_args.kwargs["filter"] == '{"data.insight_id": "insight-a"}' - - @pytest.mark.asyncio async def test_count_by_returns_grouped_counts_for_shorthand_filter() -> None: mock_api = Mock() @@ -41,36 +30,31 @@ async def test_count_by_returns_grouped_counts_for_shorthand_filter() -> None: counts = await client.count_by( ExperimentGroup, "insight_id", - filter_obj={"insight_id": {"$in": ["insight-a"]}}, + filter_obj={ + "insight_id": {"$in": ["insight-a"]}, + "is_deleted": False, + }, ) assert counts == {"insight-a": 2} - mock_api.list.assert_awaited_once_with( - "experiment_group", - workspace="default", - filter='{"data.insight_id": {"$in": ["insight-a"]}}', - page=1, - page_size=1, - extra_query={"count_by": "data.insight_id"}, - ) + assert mock_api.list.await_args.kwargs["extra_query"] == {"count_by": "data.insight_id"} @pytest.mark.asyncio -async def test_count_by_preserves_top_level_parent_field() -> None: +async def test_count_by_rejects_response_without_grouped_counts() -> None: mock_api = Mock() - mock_api.list = AsyncMock(return_value=_entities_page(group_counts={"parent-id": 1})) + mock_api.list = AsyncMock(return_value=_entities_page()) client = EntityClient(mock_api) - await client.count_by(ExperimentGroup, "parent") - - assert mock_api.list.await_args.kwargs["extra_query"] == {"count_by": "parent"} + 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_response_without_grouped_counts() -> None: +async def test_count_by_rejects_non_direct_field() -> 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") + with pytest.raises(ValueError, match="direct entity data field"): + await client.count_by(ExperimentGroup, "data.insight_id") From 0ac25e46df78c0c3ec4c5b093a1a45850ccedfc8 Mon Sep 17 00:00:00 2001 From: Aditya Pandey Date: Wed, 22 Jul 2026 16:55:21 -0600 Subject: [PATCH 13/15] chore(entities): finalize narrow grouped counts Signed-off-by: Aditya Pandey --- ...07-22-narrow-entity-group-counts-design.md | 59 ------------------- openapi/ga/individual/platform.openapi.yaml | 6 +- openapi/ga/openapi.yaml | 6 +- openapi/openapi.yaml | 6 +- 4 files changed, 12 insertions(+), 65 deletions(-) delete mode 100644 docs/superpowers/specs/2026-07-22-narrow-entity-group-counts-design.md diff --git a/docs/superpowers/specs/2026-07-22-narrow-entity-group-counts-design.md b/docs/superpowers/specs/2026-07-22-narrow-entity-group-counts-design.md deleted file mode 100644 index e5183f99ed..0000000000 --- a/docs/superpowers/specs/2026-07-22-narrow-entity-group-counts-design.md +++ /dev/null @@ -1,59 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -# Narrow Entity Group Counts - -## Goal - -Count child entities for several parents in one request without fetching every -matching entity. Current uses are Experiment Groups by `insight_id` and -Evaluations by `experiment_group_id`. - -## API - -Keep grouped counts as optional metadata on the existing entity-list endpoint: - -```text -GET .../entities/{entity_type}?count_by=data.{field}&filter=... -``` - -When `count_by` is absent, list behavior is unchanged. When present, -`group_counts` maps each observed string value to its matching entity count. -Missing values are not returned. - -Expose the narrow client method: - -```python -await entity_client.count_by( - entity_type, - field, - workspace=workspace, - filter_obj=filter_obj, -) -``` - -The client accepts one direct entity data field and adds the `data.` prefix. -It does not accept base fields, nested paths, filter strings, or filter -operation objects. - -## Repository - -Reuse existing workspace authorization and filter application. Accept only a -`data.` path with exactly one segment below `data`, and include only -string values. Ignore missing, JSON-null, boolean, numeric, array, and object -values. Keep the 1,000-group bound. The endpoint returns HTTP 400 for an -unsupported field or bound overflow. - -The list-endpoint integration intentionally retains its normal page and total -queries. A dedicated count endpoint is excluded because it increases API and -SDK surface beyond the current need. - -## Scope Control - -Remove support and tests for base-column grouping, nested JSON paths, boolean -and numeric normalization, and expanded client base-field mappings. Avoid -refactoring existing list/filter behavior solely to share code with -`count_by`. - -Test only the client request mapping, filtered direct-string grouping, -unsupported group fields, cardinality overflow, and endpoint response wiring. diff --git a/openapi/ga/individual/platform.openapi.yaml b/openapi/ga/individual/platform.openapi.yaml index 1c7ad894bf..a405df280e 100644 --- a/openapi/ga/individual/platform.openapi.yaml +++ b/openapi/ga/individual/platform.openapi.yaml @@ -703,10 +703,12 @@ paths: in: query required: false schema: - description: Optional scalar field whose matching values should be counted. + description: Optional direct string data field whose matching values should + be counted. title: Count By type: string - description: Optional scalar field whose matching values should be counted. + description: Optional direct string data field whose matching values should + be counted. - name: filter in: query required: false diff --git a/openapi/ga/openapi.yaml b/openapi/ga/openapi.yaml index 1c7ad894bf..a405df280e 100644 --- a/openapi/ga/openapi.yaml +++ b/openapi/ga/openapi.yaml @@ -703,10 +703,12 @@ paths: in: query required: false schema: - description: Optional scalar field whose matching values should be counted. + description: Optional direct string data field whose matching values should + be counted. title: Count By type: string - description: Optional scalar field whose matching values should be counted. + description: Optional direct string data field whose matching values should + be counted. - name: filter in: query required: false diff --git a/openapi/openapi.yaml b/openapi/openapi.yaml index 1c7ad894bf..a405df280e 100644 --- a/openapi/openapi.yaml +++ b/openapi/openapi.yaml @@ -703,10 +703,12 @@ paths: in: query required: false schema: - description: Optional scalar field whose matching values should be counted. + description: Optional direct string data field whose matching values should + be counted. title: Count By type: string - description: Optional scalar field whose matching values should be counted. + description: Optional direct string data field whose matching values should + be counted. - name: filter in: query required: false From c18525912063d263673f72e2ce2920c025b413a5 Mon Sep 17 00:00:00 2001 From: Aditya Pandey Date: Wed, 22 Jul 2026 17:02:35 -0600 Subject: [PATCH 14/15] chore(entities): apply final-review cleanup to grouped counts Signed-off-by: Aditya Pandey --- .../nemo_platform_plugin/tests/test_entity_client.py | 4 +++- .../src/nmp/core/entities/app/repository/entity.py | 2 +- .../tests/repository/test_entity_group_counts.py | 12 +++--------- 3 files changed, 7 insertions(+), 11 deletions(-) diff --git a/packages/nemo_platform_plugin/tests/test_entity_client.py b/packages/nemo_platform_plugin/tests/test_entity_client.py index 5eb6080c93..3847de42ab 100644 --- a/packages/nemo_platform_plugin/tests/test_entity_client.py +++ b/packages/nemo_platform_plugin/tests/test_entity_client.py @@ -37,6 +37,9 @@ async def test_count_by_returns_grouped_counts_for_shorthand_filter() -> None: ) 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"} @@ -53,7 +56,6 @@ async def test_count_by_rejects_response_without_grouped_counts() -> None: @pytest.mark.asyncio async def test_count_by_rejects_non_direct_field() -> None: mock_api = Mock() - mock_api.list = AsyncMock(return_value=_entities_page()) client = EntityClient(mock_api) with pytest.raises(ValueError, match="direct entity data field"): diff --git a/services/core/entities/src/nmp/core/entities/app/repository/entity.py b/services/core/entities/src/nmp/core/entities/app/repository/entity.py index 17a7171698..ef8e168a2d 100644 --- a/services/core/entities/src/nmp/core/entities/app/repository/entity.py +++ b/services/core/entities/src/nmp/core/entities/app/repository/entity.py @@ -122,7 +122,7 @@ async def count_entities_by( relationship_child_workspaces: set[str] | None = None, session: AsyncSession | None = None, ) -> dict[str, int]: - """Count filtered entities grouped by a scalar field.""" + """Count filtered entities grouped by a direct string data field.""" pass @abstractmethod diff --git a/services/core/entities/tests/repository/test_entity_group_counts.py b/services/core/entities/tests/repository/test_entity_group_counts.py index 6a690a53d4..b63abd8519 100644 --- a/services/core/entities/tests/repository/test_entity_group_counts.py +++ b/services/core/entities/tests/repository/test_entity_group_counts.py @@ -4,7 +4,7 @@ """Tests for grouped entity counts.""" import pytest -from nmp.common.api.filter import ComparisonOperation, FilterOperator, LogicalOperation +from nmp.common.api.filter import ComparisonOperation, FilterOperator from nmp.core.entities.app.repository import SQLAlchemyEntityRepository from nmp.core.entities.app.repository.sqlalchemy import entity as entity_repository @@ -14,7 +14,7 @@ async def test_counts_filtered_entities_grouped_by_direct_string_data_field( entity_repo: SQLAlchemyEntityRepository, setup_workspaces ): - """Count only live experiment groups for the requested insights.""" + """Count live experiment groups grouped by string insight_id values.""" entities = ( ("live-a-1", {"insight_id": "insight-a", "is_deleted": False}), ("live-a-2", {"insight_id": "insight-a", "is_deleted": False}), @@ -33,13 +33,7 @@ async def test_counts_filtered_entities_grouped_by_direct_string_data_field( data=data, ) - filter_op = LogicalOperation( - operator=FilterOperator.AND, - operations=[ - ComparisonOperation(field="data.insight_id", operator=FilterOperator.IN, value=["insight-a", "insight-b"]), - ComparisonOperation(field="data.is_deleted", operator=FilterOperator.EQ, value=False), - ], - ) + filter_op = ComparisonOperation(field="data.is_deleted", operator=FilterOperator.EQ, value=False) counts = await entity_repo.count_entities_by( workspace="workspace-1", From 91f68e75888ff9bb040931adcfad0023bf354527 Mon Sep 17 00:00:00 2001 From: Aditya Pandey Date: Thu, 23 Jul 2026 14:11:15 -0600 Subject: [PATCH 15/15] fix(entities): decode grouped JSON string values Signed-off-by: Aditya Pandey --- sdk/python/nemo-platform/.nmpcontext/openapi.yaml | 15 +++++++++++++++ .../nemo_platform/resources/entities/entities.py | 8 ++++++++ .../nemo_platform/types/entities/entities_page.py | 2 ++ .../types/entities/entity_list_params.py | 3 +++ .../tests/api_resources/test_entities.py | 2 ++ .../entities/app/repository/sqlalchemy/entity.py | 4 ++-- .../tests/repository/test_entity_group_counts.py | 3 ++- 7 files changed, 34 insertions(+), 3 deletions(-) diff --git a/sdk/python/nemo-platform/.nmpcontext/openapi.yaml b/sdk/python/nemo-platform/.nmpcontext/openapi.yaml index f91c0ade04..a405df280e 100644 --- a/sdk/python/nemo-platform/.nmpcontext/openapi.yaml +++ b/sdk/python/nemo-platform/.nmpcontext/openapi.yaml @@ -699,6 +699,16 @@ paths: default: -created_at title: Sort description: Sort field + - name: count_by + in: query + required: false + schema: + description: Optional direct string data field whose matching values should + be counted. + title: Count By + type: string + description: Optional direct string data field whose matching values should + be counted. - name: filter in: query required: false @@ -10276,6 +10286,11 @@ components: description: Filtering information. additionalProperties: true type: object + group_counts: + title: Group Counts + additionalProperties: + type: integer + type: object type: object required: - data diff --git a/sdk/python/nemo-platform/src/nemo_platform/resources/entities/entities.py b/sdk/python/nemo-platform/src/nemo_platform/resources/entities/entities.py index 3375494483..881cc756b8 100644 --- a/sdk/python/nemo-platform/src/nemo_platform/resources/entities/entities.py +++ b/sdk/python/nemo-platform/src/nemo_platform/resources/entities/entities.py @@ -145,6 +145,7 @@ def list( entity_type: str, *, workspace: str | None = None, + count_by: str | Omit = omit, filter: str | Omit = omit, page: int | Omit = omit, page_size: int | Omit = omit, @@ -176,6 +177,8 @@ def list( ``` Args: + count_by: Optional direct string data field whose matching values should be counted. + filter: Query filter expression. Supports text and JSON syntaxes: @@ -221,6 +224,7 @@ def list( timeout=timeout, query=maybe_transform( { + "count_by": count_by, "filter": filter, "page": page, "page_size": page_size, @@ -580,6 +584,7 @@ def list( entity_type: str, *, workspace: str | None = None, + count_by: str | Omit = omit, filter: str | Omit = omit, page: int | Omit = omit, page_size: int | Omit = omit, @@ -611,6 +616,8 @@ def list( ``` Args: + count_by: Optional direct string data field whose matching values should be counted. + filter: Query filter expression. Supports text and JSON syntaxes: @@ -656,6 +663,7 @@ def list( timeout=timeout, query=maybe_transform( { + "count_by": count_by, "filter": filter, "page": page, "page_size": page_size, diff --git a/sdk/python/nemo-platform/src/nemo_platform/types/entities/entities_page.py b/sdk/python/nemo-platform/src/nemo_platform/types/entities/entities_page.py index 1482dce772..8fe357de60 100644 --- a/sdk/python/nemo-platform/src/nemo_platform/types/entities/entities_page.py +++ b/sdk/python/nemo-platform/src/nemo_platform/types/entities/entities_page.py @@ -30,6 +30,8 @@ class EntitiesPage(BaseModel): filter: Optional[Dict[str, object]] = None """Filtering information.""" + group_counts: Optional[Dict[str, int]] = None + pagination: Optional[PaginationData] = None """Pagination information.""" diff --git a/sdk/python/nemo-platform/src/nemo_platform/types/entities/entity_list_params.py b/sdk/python/nemo-platform/src/nemo_platform/types/entities/entity_list_params.py index 6ef35bfff0..79be468fab 100644 --- a/sdk/python/nemo-platform/src/nemo_platform/types/entities/entity_list_params.py +++ b/sdk/python/nemo-platform/src/nemo_platform/types/entities/entity_list_params.py @@ -25,6 +25,9 @@ class EntityListParams(TypedDict, total=False): workspace: str + count_by: str + """Optional direct string data field whose matching values should be counted.""" + filter: str """Query filter expression. Supports text and JSON syntaxes: diff --git a/sdk/python/nemo-platform/tests/api_resources/test_entities.py b/sdk/python/nemo-platform/tests/api_resources/test_entities.py index a76aaf406b..57781ebf66 100644 --- a/sdk/python/nemo-platform/tests/api_resources/test_entities.py +++ b/sdk/python/nemo-platform/tests/api_resources/test_entities.py @@ -121,6 +121,7 @@ def test_method_list_with_all_params(self, client: NeMoPlatform) -> None: entity = client.entities.list( entity_type="entity_type", workspace="workspace", + count_by="count_by", filter="filter", page=1, page_size=1, @@ -539,6 +540,7 @@ async def test_method_list_with_all_params(self, async_client: AsyncNeMoPlatform entity = await async_client.entities.list( entity_type="entity_type", workspace="workspace", + count_by="count_by", filter="filter", page=1, page_size=1, diff --git a/services/core/entities/src/nmp/core/entities/app/repository/sqlalchemy/entity.py b/services/core/entities/src/nmp/core/entities/app/repository/sqlalchemy/entity.py index f0b948c247..202b29612f 100644 --- a/services/core/entities/src/nmp/core/entities/app/repository/sqlalchemy/entity.py +++ b/services/core/entities/src/nmp/core/entities/app/repository/sqlalchemy/entity.py @@ -15,7 +15,7 @@ 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 String, cast, func, select +from sqlalchemy import func, select from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker from sqlalchemy.orm.exc import StaleDataError @@ -217,7 +217,7 @@ async def count_entities_by( field = parts[1] raw_group_column = DBEntity.data[field] - group_column = func.trim(cast(raw_group_column, String), '"') + group_column = raw_group_column.as_string() if self._is_sqlite(sess): json_type = func.json_type(DBEntity.data, f"$.{field}") string_type = "text" diff --git a/services/core/entities/tests/repository/test_entity_group_counts.py b/services/core/entities/tests/repository/test_entity_group_counts.py index b63abd8519..2c0b6e7769 100644 --- a/services/core/entities/tests/repository/test_entity_group_counts.py +++ b/services/core/entities/tests/repository/test_entity_group_counts.py @@ -19,6 +19,7 @@ async def test_counts_filtered_entities_grouped_by_direct_string_data_field( ("live-a-1", {"insight_id": "insight-a", "is_deleted": False}), ("live-a-2", {"insight_id": "insight-a", "is_deleted": False}), ("live-b", {"insight_id": "insight-b", "is_deleted": False}), + ("escaped", {"insight_id": 'insight-"quoted"\\path', "is_deleted": False}), ("deleted", {"insight_id": "insight-a", "is_deleted": True}), ("missing", {"is_deleted": False}), ("null", {"insight_id": None, "is_deleted": False}), @@ -42,7 +43,7 @@ async def test_counts_filtered_entities_grouped_by_direct_string_data_field( filter_op=filter_op, ) - assert counts == {"insight-a": 2, "insight-b": 1} + assert counts == {'insight-"quoted"\\path': 1, "insight-a": 2, "insight-b": 1} @pytest.mark.parametrize("field", ["name", "data.nested.value", "data.", "data.not-valid"])