[SPIKE] feat: refresh record count_submitted_responses column using context functions

argilla-server/CHANGELOG.md

@@ -16,12 +16,17 @@ These are the section headers that we use:
 
 ## [Unreleased]()
 
-## [2.0.0rc1](https://github.com/argilla-io/argilla/compare/v1.29.0...v2.0.0rc1)
+### Added
+
+- Added support to specify `distribution` attribute when creating a dataset. ([#5013](https://github.com/argilla-io/argilla/pull/5013))
+- Added support to change `distribution` attribute when updating a dataset. ([#5028](https://github.com/argilla-io/argilla/pull/5028))
 
 ### Changed
 
 - Change `responses` table to delete rows on cascade when a user is deleted. ([#5126](https://github.com/argilla-io/argilla/pull/5126))
 
+## [2.0.0rc1](https://github.com/argilla-io/argilla/compare/v1.29.0...v2.0.0rc1)
+
 ### Removed
 
 - Removed all API v0 endpoints. ([#4852](https://github.com/argilla-io/argilla/pull/4852))

argilla-server/src/argilla_server/alembic/versions/3ff6484f8b37_add_record_metadata_column.py → argilla-server/src/argilla_server/alembic/versions/3ff6484f8b37_add_metadata_column_to_records_table.py

@@ -12,7 +12,7 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
-"""add record metadata column
+"""add metadata column to records table
 
 Revision ID: 3ff6484f8b37
 Revises: ae5522b4c674
@@ -31,12 +31,8 @@
 
 
 def upgrade() -> None:
-    # ### commands auto generated by Alembic - please adjust! ###
     op.add_column("records", sa.Column("metadata", sa.JSON(), nullable=True))
-    # ### end Alembic commands ###
 
 
 def downgrade() -> None:
-    # ### commands auto generated by Alembic - please adjust! ###
     op.drop_column("records", "metadata")
-    # ### end Alembic commands ###

argilla-server/src/argilla_server/alembic/versions/45a12f74448b_add_distribution_column_to_datasets_table.py

@@ -0,0 +1,45 @@
+#  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.
+
+"""add distribution column to datasets table
+
+Revision ID: 45a12f74448b
+Revises: d00f819ccc67
+Create Date: 2024-06-13 11:23:43.395093
+
+"""
+
+import json
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "45a12f74448b"
+down_revision = "d00f819ccc67"
+branch_labels = None
+depends_on = None
+
+DISTRIBUTION_VALUE = json.dumps({"strategy": "overlap", "min_submitted": 1})
+
+
+def upgrade() -> None:
+    op.add_column("datasets", sa.Column("distribution", sa.JSON(), nullable=True))
+    op.execute(f"UPDATE datasets SET distribution = '{DISTRIBUTION_VALUE}'")
+    with op.batch_alter_table("datasets") as batch_op:
+        batch_op.alter_column("distribution", nullable=False)
+
+
+def downgrade() -> None:
+    op.drop_column("datasets", "distribution")

argilla-server/src/argilla_server/alembic/versions/b4e101b124d2_add_count_submitted_responses_to_.py

@@ -0,0 +1,47 @@
+#  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.
+
+"""add count_submitted_responses to records table
+
+Revision ID: b4e101b124d2
+Revises: 45a12f74448b
+Create Date: 2024-06-24 17:07:18.614728
+
+"""
+
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision = "b4e101b124d2"
+down_revision = "45a12f74448b"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    op.add_column("records", sa.Column("count_submitted_responses", sa.Integer(), server_default="0", nullable=False))
+    op.execute("""
+        UPDATE records
+        SET count_submitted_responses = (
+            SELECT COUNT(*)
+            FROM responses
+            WHERE responses.record_id = records.id AND responses.status = 'submitted'
+        )
+    """)
+
+
+def downgrade() -> None:
+    op.drop_column("records", "count_submitted_responses")

argilla-server/src/argilla_server/alembic/versions/b8458008b60e_add_allow_extra_metadata_column_to_.py → argilla-server/src/argilla_server/alembic/versions/b8458008b60e_add_allow_extra_metadata_column_to_datasets_table.py

@@ -12,7 +12,7 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
-"""add allow_extra_metadata column to dataset table
+"""add allow_extra_metadata column to datasets table
 
 Revision ID: b8458008b60e
 Revises: 7cbcccf8b57a
@@ -31,14 +31,10 @@
 
 
 def upgrade() -> None:
-    # ### commands auto generated by Alembic - please adjust! ###
     op.add_column(
         "datasets", sa.Column("allow_extra_metadata", sa.Boolean(), server_default=sa.text("true"), nullable=False)
     )
-    # ### end Alembic commands ###
 
 
 def downgrade() -> None:
-    # ### commands auto generated by Alembic - please adjust! ###
     op.drop_column("datasets", "allow_extra_metadata")
-    # ### end Alembic commands ###

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

@@ -189,7 +189,7 @@ async def create_dataset(
 ):
     await authorize(current_user, DatasetPolicy.create(dataset_create.workspace_id))
 
-    return await datasets.create_dataset(db, dataset_create)
+    return await datasets.create_dataset(db, dataset_create.dict())
 
 
 @router.post("/datasets/{dataset_id}/fields", status_code=status.HTTP_201_CREATED, response_model=Field)
@@ -302,4 +302,4 @@ async def update_dataset(
 
     await authorize(current_user, DatasetPolicy.update(dataset))
 
-    return await datasets.update_dataset(db, dataset, dataset_update)
+    return await datasets.update_dataset(db, dataset, dataset_update.dict(exclude_unset=True))

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

@@ -13,11 +13,11 @@
 #  limitations under the License.
 
 from datetime import datetime
-from typing import List, Optional
+from typing import List, Literal, Optional, Union
 from uuid import UUID
 
 from argilla_server.api.schemas.v1.commons import UpdateSchema
-from argilla_server.enums import DatasetStatus
+from argilla_server.enums import DatasetDistributionStrategy, DatasetStatus
 from argilla_server.pydantic_v1 import BaseModel, Field, constr
 
 try:
@@ -44,6 +44,32 @@
 ]
 
 
+class DatasetOverlapDistribution(BaseModel):
+    strategy: Literal[DatasetDistributionStrategy.overlap]
+    min_submitted: int
+
+
+DatasetDistribution = DatasetOverlapDistribution
+
+
+class DatasetOverlapDistributionCreate(BaseModel):
+    strategy: Literal[DatasetDistributionStrategy.overlap]
+    min_submitted: int = Field(
+        ge=1,
+        description="Minimum number of submitted responses to consider a record as completed",
+    )
+
+
+DatasetDistributionCreate = DatasetOverlapDistributionCreate
+
+
+class DatasetOverlapDistributionUpdate(DatasetDistributionCreate):
+    pass
+
+
+DatasetDistributionUpdate = DatasetOverlapDistributionUpdate
+
+
 class RecordMetrics(BaseModel):
     count: int
 
@@ -74,6 +100,7 @@ class Dataset(BaseModel):
     guidelines: Optional[str]
     allow_extra_metadata: bool
     status: DatasetStatus
+    distribution: DatasetDistribution
     workspace_id: UUID
     last_activity_at: datetime
     inserted_at: datetime
@@ -91,12 +118,17 @@ class DatasetCreate(BaseModel):
     name: DatasetName
     guidelines: Optional[DatasetGuidelines]
     allow_extra_metadata: bool = True
+    distribution: DatasetDistributionCreate = DatasetOverlapDistributionCreate(
+        strategy=DatasetDistributionStrategy.overlap,
+        min_submitted=1,
+    )
     workspace_id: UUID
 
 
 class DatasetUpdate(UpdateSchema):
     name: Optional[DatasetName]
     guidelines: Optional[DatasetGuidelines]
     allow_extra_metadata: Optional[bool]
+    distribution: Optional[DatasetDistributionUpdate]
 
-    __non_nullable_fields__ = {"name", "allow_extra_metadata"}
+    __non_nullable_fields__ = {"name", "allow_extra_metadata", "distribution"}

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

@@ -37,10 +37,7 @@
 from sqlalchemy.ext.asyncio import AsyncSession
 from sqlalchemy.orm import contains_eager, joinedload, selectinload
 
-from argilla_server.api.schemas.v1.datasets import (
-    DatasetCreate,
-    DatasetProgress,
-)
+from argilla_server.api.schemas.v1.datasets import DatasetProgress
 from argilla_server.api.schemas.v1.fields import FieldCreate
 from argilla_server.api.schemas.v1.metadata_properties import MetadataPropertyCreate, MetadataPropertyUpdate
 from argilla_server.api.schemas.v1.records import (
@@ -82,6 +79,7 @@
 )
 from argilla_server.models.suggestions import SuggestionCreateWithRecordId
 from argilla_server.search_engine import SearchEngine
+from argilla_server.validators.datasets import DatasetCreateValidator
 from argilla_server.validators.responses import (
     ResponseCreateValidator,
     ResponseUpdateValidator,
@@ -122,22 +120,18 @@ async def list_datasets_by_workspace_id(db: AsyncSession, workspace_id: UUID) ->
     return result.scalars().all()
 
 
-async def create_dataset(db: AsyncSession, dataset_create: DatasetCreate):
-    if await Workspace.get(db, dataset_create.workspace_id) is None:
-        raise UnprocessableEntityError(f"Workspace with id `{dataset_create.workspace_id}` not found")
+async def create_dataset(db: AsyncSession, dataset_attrs: dict):
+    dataset = Dataset(
+        name=dataset_attrs["name"],
+        guidelines=dataset_attrs["guidelines"],
+        allow_extra_metadata=dataset_attrs["allow_extra_metadata"],
+        distribution=dataset_attrs["distribution"],
+        workspace_id=dataset_attrs["workspace_id"],
+    )
 
-    if await Dataset.get_by(db, name=dataset_create.name, workspace_id=dataset_create.workspace_id):
-        raise NotUniqueError(
-            f"Dataset with name `{dataset_create.name}` already exists for workspace with id `{dataset_create.workspace_id}`"
-        )
+    await DatasetCreateValidator.validate(db, dataset)
 
-    return await Dataset.create(
-        db,
-        name=dataset_create.name,
-        guidelines=dataset_create.guidelines,
-        allow_extra_metadata=dataset_create.allow_extra_metadata,
-        workspace_id=dataset_create.workspace_id,
-    )
+    return await dataset.save(db)
 
 
 async def _count_required_fields_by_dataset_id(db: AsyncSession, dataset_id: UUID) -> int:
@@ -176,6 +170,10 @@ async def publish_dataset(db: AsyncSession, search_engine: SearchEngine, dataset
     return dataset
 
 
+async def update_dataset(db: AsyncSession, dataset: Dataset, dataset_attrs: dict) -> Dataset:
+    return await dataset.update(db, **dataset_attrs)
+
+
 async def delete_dataset(db: AsyncSession, search_engine: SearchEngine, dataset: Dataset) -> Dataset:
     async with db.begin_nested():
         dataset = await dataset.delete(db, autocommit=False)
@@ -186,11 +184,6 @@ async def delete_dataset(db: AsyncSession, search_engine: SearchEngine, dataset:
     return dataset
 
 
-async def update_dataset(db: AsyncSession, dataset: Dataset, dataset_update: "DatasetUpdate") -> Dataset:
-    params = dataset_update.dict(exclude_unset=True)
-    return await dataset.update(db, **params)
-
-
 async def create_field(db: AsyncSession, dataset: Dataset, field_create: FieldCreate) -> Field:
     if dataset.is_ready:
         raise UnprocessableEntityError("Field cannot be created for a published dataset")
@@ -943,6 +936,15 @@ async def create_response(
         )
 
         await db.flush([response])
+
+        if response_create.status == ResponseStatus.submitted:
+            await db.execute(
+                sqlalchemy.update(Record)
+                .where(Record.id == record.id)
+                .values(count_submitted_responses=Record.count_submitted_responses + 1)
+            )
+        # TODO: refresh record at search engine
+
         await _touch_dataset_last_activity_at(db, record.dataset)
         await search_engine.update_record_response(response)
 
@@ -957,6 +959,8 @@ async def update_response(
     ResponseUpdateValidator(response_update).validate_for(response.record)
 
     async with db.begin_nested():
+        previous_response_status = response.status
+
         response = await response.update(
             db,
             values=jsonable_encoder(response_update.values),
@@ -966,6 +970,21 @@ async def update_response(
         )
 
         await _load_users_from_responses(response)
+
+        if response_update.status == ResponseStatus.submitted and previous_response_status != ResponseStatus.submitted:
+            await db.execute(
+                sqlalchemy.update(Record)
+                .where(Record.id == response.record_id)
+                .values(count_submitted_responses=Record.count_submitted_responses + 1)
+            )
+        if response_update.status != ResponseStatus.submitted and previous_response_status == ResponseStatus.submitted:
+            await db.execute(
+                sqlalchemy.update(Record)
+                .where(Record.id == response.record_id)
+                .values(count_submitted_responses=Record.count_submitted_responses - 1)
+            )
+        # TODO: refresh record at search engine
+
         await _touch_dataset_last_activity_at(db, response.record.dataset)
         await search_engine.update_record_response(response)
 
@@ -1007,6 +1026,15 @@ async def delete_response(db: AsyncSession, search_engine: SearchEngine, respons
     async with db.begin_nested():
         response = await response.delete(db, autocommit=False)
         await _load_users_from_responses(response)
+
+        if response.status == ResponseStatus.submitted:
+            await db.execute(
+                sqlalchemy.update(Record)
+                .where(Record.id == response.record_id)
+                .values(count_submitted_responses=Record.count_submitted_responses - 1)
+            )
+        # TODO: refresh record at search engine
+
         await _touch_dataset_last_activity_at(db, response.record.dataset)
         await search_engine.delete_record_response(response)
 

argilla-server/src/argilla_server/enums.py

@@ -43,6 +43,10 @@ class DatasetStatus(str, Enum):
     ready = "ready"
 
 
+class DatasetDistributionStrategy(str, Enum):
+    overlap = "overlap"
+
+
 class UserRole(str, Enum):
     owner = "owner"
     admin = "admin"