[FEATURE] argilla server: Add annotators progress endpoint (#5367)

# Description
<!-- Please include a summary of the changes and the related issue.
Please also include relevant motivation and context. List any
dependencies that are required for this change. -->

This PR adds a new endpoint to compute the dataset annotators' progress.

The endpoint is  defined as
```json
GET /api/v1/datasets/:dataset_id/progress/annotators

{
  "annotators": [
      { 
         "username": "userA",  
         "completed": { "submitted": 20, "draft": 5, "discarded": 1 },
         "pending": { "submitted": 2, "draft": 50, "discarded": 10 }
      },
      ...
 ]
}
```

Progress distribution is split by record status, so values represent
user responses for each kind of record status.


**Type of change**
<!-- Please delete options that are not relevant. Remember to title the
PR according to the type of change -->

- New feature (non-breaking change which adds functionality)

**How Has This Been Tested**
<!-- Please add some reference about how your feature has been tested.
-->

**Checklist**
<!-- Please go over the list and make sure you've taken everything into
account -->

- I added relevant documentation
- I followed the style guidelines of this project
- I did a self-review of my code
- I made corresponding changes to the documentation
- I confirm My changes generate no new warnings
- I have added tests that prove my fix is effective or that my feature
works
- I have added relevant notes to the CHANGELOG.md file (See
https://keepachangelog.com/)

---------

Co-authored-by: José Francisco Calvo <[email protected]>

Author: frascuchon

argilla-server/CHANGELOG.md

@@ -16,6 +16,10 @@ These are the section headers that we use:
 
 ## [Unreleased]()
 
+### Added
+
+- Added new endpoint `GET /api/v1/datsets/:dataset_id/users/progress` to compute the users progress. ([#5367](https://github.com/argilla-io/argilla/pull/5367))
+
 ## [2.0.0](https://github.com/argilla-io/argilla/compare/v2.0.0rc1...v2.0.0)
 
 > [!IMPORTANT]

argilla-server/src/argilla_server/api/handlers/v1/datasets/datasets.py

@@ -22,6 +22,7 @@
 from argilla_server.api.policies.v1 import DatasetPolicy, MetadataPropertyPolicy, authorize, is_authorized
 from argilla_server.api.schemas.v1.datasets import (
     Dataset as DatasetSchema,
+    UsersProgress,
 )
 from argilla_server.api.schemas.v1.datasets import (
     DatasetCreate,
@@ -39,7 +40,6 @@
 from argilla_server.api.schemas.v1.vector_settings import VectorSettings, VectorSettingsCreate, VectorsSettings
 from argilla_server.contexts import datasets
 from argilla_server.database import get_async_db
-from argilla_server.enums import ResponseStatus
 from argilla_server.models import Dataset, User
 from argilla_server.search_engine import (
     SearchEngine,
@@ -161,7 +161,23 @@ async def get_dataset_progress(
 
     await authorize(current_user, DatasetPolicy.get(dataset))
 
-    return await datasets.get_dataset_progress(db, dataset_id)
+    return await datasets.get_dataset_progress(db, dataset.id)
+
+
+@router.get("/datasets/{dataset_id}/users/progress", response_model=UsersProgress, response_model_exclude_unset=True)
+async def get_dataset_users_progress(
+    *,
+    current_user: User = Security(auth.get_current_user),
+    dataset_id: UUID,
+    db: AsyncSession = Depends(get_async_db),
+):
+    dataset = await Dataset.get_or_raise(db, dataset_id)
+
+    await authorize(current_user, DatasetPolicy.get(dataset))
+
+    progress = await datasets.get_dataset_users_progress(dataset.id)
+
+    return UsersProgress(users=progress)
 
 
 @router.post("/datasets", status_code=status.HTTP_201_CREATED, response_model=DatasetSchema)

argilla-server/src/argilla_server/api/schemas/v1/datasets.py

@@ -31,13 +31,11 @@
 DATASET_GUIDELINES_MIN_LENGTH = 1
 DATASET_GUIDELINES_MAX_LENGTH = 10000
 
-
 DatasetName = Annotated[
     constr(regex=DATASET_NAME_REGEX, min_length=DATASET_NAME_MIN_LENGTH, max_length=DATASET_NAME_MAX_LENGTH),
     Field(..., description="Dataset name"),
 ]
 
-
 DatasetGuidelines = Annotated[
     constr(min_length=DATASET_GUIDELINES_MIN_LENGTH, max_length=DATASET_GUIDELINES_MAX_LENGTH),
     Field(..., description="Dataset guidelines"),
@@ -88,6 +86,22 @@ class DatasetProgress(BaseModel):
     pending: int
 
 
+class RecordResponseDistribution(BaseModel):
+    submitted: Optional[int]
+    discarded: Optional[int]
+    draft: Optional[int]
+
+
+class UserProgress(BaseModel):
+    username: str
+    completed: RecordResponseDistribution
+    pending: RecordResponseDistribution
+
+
+class UsersProgress(BaseModel):
+    users: List[UserProgress]
+
+
 class Dataset(BaseModel):
     id: UUID
     name: str

argilla-server/src/argilla_server/contexts/datasets.py

@@ -14,6 +14,8 @@
 
 import asyncio
 import copy
+from collections import defaultdict
+
 import sqlalchemy
 
 from datetime import datetime
@@ -59,6 +61,7 @@
 )
 from argilla_server.api.schemas.v1.vectors import Vector as VectorSchema
 from argilla_server.contexts import accounts, distribution
+from argilla_server.database import get_async_db
 from argilla_server.enums import DatasetStatus, UserRole, RecordStatus
 from argilla_server.errors.future import NotUniqueError, UnprocessableEntityError
 from argilla_server.models import (
@@ -425,15 +428,19 @@ async def get_user_dataset_metrics(db: AsyncSession, user_id: UUID, dataset_id:
 async def get_dataset_progress(db: AsyncSession, dataset_id: UUID) -> dict:
     records_completed, records_pending = await asyncio.gather(
         db.execute(
-            select(func.count(Record.id)).filter(
-                Record.dataset_id == dataset_id,
-                Record.status == RecordStatus.completed,
+            select(func.count(Record.id)).where(
+                and_(
+                    Record.dataset_id == dataset_id,
+                    Record.status == RecordStatus.completed,
+                )
             ),
         ),
         db.execute(
-            select(func.count(Record.id)).filter(
-                Record.dataset_id == dataset_id,
-                Record.status == RecordStatus.pending,
+            select(func.count(Record.id)).where(
+                and_(
+                    Record.dataset_id == dataset_id,
+                    Record.status == RecordStatus.pending,
+                )
             ),
         ),
     )
@@ -449,6 +456,25 @@ async def get_dataset_progress(db: AsyncSession, dataset_id: UUID) -> dict:
     }
 
 
+async def get_dataset_users_progress(dataset_id: UUID) -> List[dict]:
+    query = (
+        select(User.username, Record.status, Response.status, func.count(Response.id))
+        .join(Record)
+        .join(User)
+        .where(Record.dataset_id == dataset_id)
+        .group_by(User.username, Record.status, Response.status)
+    )
+
+    async for session in get_async_db():
+        annotators_progress = defaultdict(lambda: defaultdict(dict))
+        results = (await session.execute(query)).all()
+
+        for username, record_status, response_status, count in results:
+            annotators_progress[username][record_status][response_status] = count
+
+        return [{"username": username, **progress} for username, progress in annotators_progress.items()]
+
+
 _EXTRA_METADATA_FLAG = "extra"
 
 

argilla-server/tests/unit/api/handlers/v1/datasets/test_get_dataset_users_progress.py

@@ -0,0 +1,125 @@
+#  Copyright 2021-present, the Recognai S.L. team.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+from uuid import UUID, uuid4
+
+import pytest
+from httpx import AsyncClient
+
+from argilla_server.constants import API_KEY_HEADER_NAME
+from argilla_server.enums import RecordStatus, UserRole, ResponseStatus
+from tests.factories import DatasetFactory, RecordFactory, AnnotatorFactory, ResponseFactory, UserFactory
+
+
+@pytest.mark.asyncio
+class TestGetDatasetUsersProgress:
+    def url(self, dataset_id: UUID) -> str:
+        return f"/api/v1/datasets/{dataset_id}/users/progress"
+
+    async def test_get_dataset_users_progress(self, async_client: AsyncClient, owner_auth_header: dict):
+        dataset = await DatasetFactory.create()
+
+        user_with_submitted = await AnnotatorFactory.create()
+        user_with_draft = await AnnotatorFactory.create()
+        user_with_discarded = await AnnotatorFactory.create()
+
+        records_completed = await RecordFactory.create_batch(3, status=RecordStatus.completed, dataset=dataset)
+        records_pending = await RecordFactory.create_batch(2, status=RecordStatus.pending, dataset=dataset)
+
+        for record in records_completed + records_pending:
+            await ResponseFactory.create(record=record, user=user_with_submitted, status=ResponseStatus.submitted)
+            await ResponseFactory.create(record=record, user=user_with_draft, status=ResponseStatus.draft)
+            await ResponseFactory.create(record=record, user=user_with_discarded, status=ResponseStatus.discarded)
+
+        response = await async_client.get(self.url(dataset.id), headers=owner_auth_header)
+
+        assert response.status_code == 200, response.json()
+        assert response.json() == {
+            "users": [
+                {
+                    "username": user_with_submitted.username,
+                    "completed": {"submitted": 3},
+                    "pending": {"submitted": 2},
+                },
+                {
+                    "username": user_with_draft.username,
+                    "completed": {"draft": 3},
+                    "pending": {"draft": 2},
+                },
+                {
+                    "username": user_with_discarded.username,
+                    "completed": {"discarded": 3},
+                    "pending": {"discarded": 2},
+                },
+            ]
+        }
+
+    async def test_get_dataset_users_progress_with_empty_dataset(
+        self, async_client: AsyncClient, owner_auth_header: dict
+    ):
+        dataset = await DatasetFactory.create()
+
+        response = await async_client.get(self.url(dataset.id), headers=owner_auth_header)
+
+        assert response.status_code == 200
+        assert response.json() == {"users": []}
+
+    @pytest.mark.parametrize("user_role", [UserRole.admin, UserRole.annotator])
+    async def test_get_dataset_users_progress_as_restricted_user(self, async_client: AsyncClient, user_role: UserRole):
+        dataset = await DatasetFactory.create()
+        user = await UserFactory.create(workspaces=[dataset.workspace], role=user_role)
+
+        response = await async_client.get(self.url(dataset.id), headers={API_KEY_HEADER_NAME: user.api_key})
+
+        assert response.status_code == 200
+
+    @pytest.mark.parametrize("user_role", [UserRole.admin, UserRole.annotator])
+    async def test_get_dataset_users_progress_as_restricted_user_from_different_workspace(
+        self, async_client: AsyncClient, user_role: UserRole
+    ):
+        dataset = await DatasetFactory.create()
+
+        other_dataset = await DatasetFactory.create()
+        user = await UserFactory.create(workspaces=[other_dataset.workspace], role=user_role)
+
+        response = await async_client.get(self.url(dataset.id), headers={API_KEY_HEADER_NAME: user.api_key})
+
+        assert response.status_code == 403
+        assert response.json() == {
+            "detail": {
+                "code": "argilla.api.errors::ForbiddenOperationError",
+                "params": {"detail": "Operation not allowed"},
+            },
+        }
+
+    async def test_get_dataset_users_progress_without_authentication(self, async_client: AsyncClient):
+        response = await async_client.get(self.url(uuid4()))
+
+        assert response.status_code == 401
+        assert response.json() == {
+            "detail": {
+                "code": "argilla.api.errors::UnauthorizedError",
+                "params": {"detail": "Could not validate credentials"},
+            },
+        }
+
+    async def test_get_dataset_users_progress_with_nonexistent_dataset_id(
+        self, async_client: AsyncClient, owner_auth_header: dict
+    ):
+        dataset_id = uuid4()
+
+        response = await async_client.get(self.url(dataset_id), headers=owner_auth_header)
+
+        assert response.status_code == 404
+        assert response.json() == {"detail": f"Dataset with id `{dataset_id}` not found"}

argilla-server/tests/unit/conftest.py

@@ -22,7 +22,7 @@
 from opensearchpy import OpenSearch
 
 from argilla_server import telemetry
-from argilla_server.contexts import distribution
+from argilla_server.contexts import distribution, datasets
 from argilla_server.api.routes import api_v1
 from argilla_server.constants import API_KEY_HEADER_NAME, DEFAULT_API_KEY
 from argilla_server.database import get_async_db
@@ -92,6 +92,7 @@ async def override_get_search_engine():
         yield mock_search_engine
 
     mocker.patch.object(distribution, "_get_async_db", override_get_async_db)
+    mocker.patch.object(datasets, "get_async_db", override_get_async_db)
 
     api_v1.dependency_overrides.update(
         {