diff --git a/openapi/ga/individual/platform.openapi.yaml b/openapi/ga/individual/platform.openapi.yaml index 9be005085e..f16a77fcae 100644 --- a/openapi/ga/individual/platform.openapi.yaml +++ b/openapi/ga/individual/platform.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 @@ -10325,6 +10335,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 9be005085e..f16a77fcae 100644 --- a/openapi/ga/openapi.yaml +++ b/openapi/ga/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 @@ -10325,6 +10335,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 9be005085e..f16a77fcae 100644 --- a/openapi/openapi.yaml +++ b/openapi/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 @@ -10325,6 +10335,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/packages/nemo_platform_plugin/src/nemo_platform_plugin/entities.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/entities.py index 56dde5bd12..0351d6e3e0 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,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. 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..3847de42ab --- /dev/null +++ b/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") diff --git a/sdk/python/nemo-platform/.nmpcontext/openapi.yaml b/sdk/python/nemo-platform/.nmpcontext/openapi.yaml index 9be005085e..f16a77fcae 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 @@ -10325,6 +10335,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/api/v2/entities/endpoints.py b/services/core/entities/src/nmp/core/entities/api/v2/entities/endpoints.py index 80d8842e3a..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 @@ -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, ) 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..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 @@ -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, 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..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 @@ -19,6 +19,8 @@ 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 = raw_group_column.as_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, *, diff --git a/services/core/entities/tests/integration/test_generic_entities.py b/services/core/entities/tests/integration/test_generic_entities.py index 6175090916..d71a199b9a 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,50 @@ 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 + + @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": count_by}, + ) + + assert response.status_code == 400 + 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.""" 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 new file mode 100644 index 0000000000..2c0b6e7769 --- /dev/null +++ b/services/core/entities/tests/repository/test_entity_group_counts.py @@ -0,0 +1,76 @@ +# 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 +from nmp.core.entities.app.repository import SQLAlchemyEntityRepository +from nmp.core.entities.app.repository.sqlalchemy import entity as entity_repository + +pytestmark = pytest.mark.asyncio + + +async def test_counts_filtered_entities_grouped_by_direct_string_data_field( + entity_repo: SQLAlchemyEntityRepository, setup_workspaces +): + """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}), + ("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}), + ("boolean", {"insight_id": True, "is_deleted": False}), + ("numeric", {"insight_id": 1, "is_deleted": False}), + ) + for name, data in entities: + await entity_repo.create_entity( + workspace="workspace-1", + entity_type="experiment_group", + name=name, + data=data, + ) + + filter_op = 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-"quoted"\\path': 1, "insight-a": 2, "insight-b": 1} + + +@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", + group_by=field, + ) + + +async def test_rejects_group_counts_over_limit( + entity_repo: SQLAlchemyEntityRepository, setup_workspaces, monkeypatch: pytest.MonkeyPatch +): + 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=f"group-{index}", + data={"value": f"group-{index}"}, + ) + + with pytest.raises(ValueError, match="more than 2 distinct values"): + await entity_repo.count_entities_by( + workspace="workspace-1", + entity_type="experiment_group", + group_by="data.value", + )