refactor(#1044): include last updated field for sort (#1093)

* refactor(API): sort by last_updated field * test: add tests * test: fix tests (cherry picked from commit 60dca5e)
argilla-io · Feb 7, 2022 · 1da7d5a · 1da7d5a
1 parent 780d54a
commit 1da7d5a
Show file tree

Hide file tree

Showing 6 changed files with 39 additions and 8 deletions.
diff --git a/src/rubrix/server/tasks/commons/api/model.py b/src/rubrix/server/tasks/commons/api/model.py
@@ -45,6 +45,7 @@ class EsRecordDataFieldNames(str, Enum):
     score = "score"
     words = "words"
     event_timestamp = "event_timestamp"
+    last_updated = "last_updated"
 
     def __str__(self):
         return self.value

diff --git a/src/rubrix/server/tasks/text2text/service/service.py b/src/rubrix/server/tasks/text2text/service/service.py
@@ -104,13 +104,12 @@ def search(
                     sort_by,
                     valid_fields=[
                         "metadata",
-                        EsRecordDataFieldNames.score,
-                        EsRecordDataFieldNames.predicted,
                         EsRecordDataFieldNames.predicted_as,
                         EsRecordDataFieldNames.predicted_by,
                         EsRecordDataFieldNames.annotated_as,
                         EsRecordDataFieldNames.annotated_by,
                         EsRecordDataFieldNames.status,
+                        EsRecordDataFieldNames.last_updated,
                         EsRecordDataFieldNames.event_timestamp,
                     ],
                 ),

diff --git a/src/rubrix/server/tasks/text_classification/service/service.py b/src/rubrix/server/tasks/text_classification/service/service.py
@@ -146,6 +146,7 @@ def search(
                     sort_by,
                     valid_fields=[
                         "metadata",
+                        EsRecordDataFieldNames.last_updated,
                         EsRecordDataFieldNames.score,
                         EsRecordDataFieldNames.predicted,
                         EsRecordDataFieldNames.predicted_as,

diff --git a/src/rubrix/server/tasks/token_classification/service/service.py b/src/rubrix/server/tasks/token_classification/service/service.py
@@ -140,6 +140,7 @@ def search(
                     sort_by,
                     valid_fields=[
                         "metadata",
+                        EsRecordDataFieldNames.last_updated,
                         EsRecordDataFieldNames.score,
                         EsRecordDataFieldNames.predicted,
                         EsRecordDataFieldNames.predicted_as,

diff --git a/tests/server/text_classification/test_api.py b/tests/server/text_classification/test_api.py
@@ -264,6 +264,33 @@ def test_partial_record_update():
     )
 
 
+def test_sort_by_last_updated():
+    dataset = "test_sort_by_last_updated"
+    assert client.delete(f"/api/datasets/{dataset}").status_code == 200
+    for i in range(0, 10):
+        client.post(
+            f"/api/datasets/{dataset}/TextClassification:bulk",
+            json=TextClassificationBulkData(
+                records=[
+                    TextClassificationRecord(
+                        **{
+                            "id": i,
+                            "inputs": {"data": "my data"},
+                            "metadata": {"s": "value"},
+                        }
+                    )
+                ],
+            ).dict(by_alias=True),
+        )
+
+    response = client.post(
+        f"/api/datasets/{dataset}/TextClassification:search?from=0&limit=10",
+        json={"sort": [{"id": "last_updated", "order": "asc"}]},
+    )
+
+    assert [r["id"] for r in response.json()["records"]] == list(range(0, 10))
+
+
 def test_sort_by_id_as_default():
     dataset = "test_sort_by_id_as_default"
     assert client.delete(f"/api/datasets/{dataset}").status_code == 200
@@ -341,9 +368,10 @@ def test_some_sort_by():
             "code": "rubrix.api.errors::BadRequestError",
             "params": {
                 "message": "Wrong sort id wrong_field. Valid values "
-                "are: ['metadata', 'score', 'predicted', "
-                "'predicted_as', 'predicted_by', "
-                "'annotated_as', 'annotated_by', 'status', "
+                "are: ['metadata', 'last_updated', 'score', "
+                "'predicted', 'predicted_as', "
+                "'predicted_by', 'annotated_as', "
+                "'annotated_by', 'status', "
                 "'event_timestamp']"
             },
         }

diff --git a/tests/server/token_classification/test_api.py b/tests/server/token_classification/test_api.py
@@ -137,9 +137,10 @@ def test_some_sort():
             "code": "rubrix.api.errors::BadRequestError",
             "params": {
                 "message": "Wrong sort id babba. Valid values are: "
-                "['metadata', 'score', 'predicted', "
-                "'predicted_as', 'predicted_by', "
-                "'annotated_as', 'annotated_by', 'status', "
+                "['metadata', 'last_updated', 'score', "
+                "'predicted', 'predicted_as', "
+                "'predicted_by', 'annotated_as', "
+                "'annotated_by', 'status', "
                 "'event_timestamp']"
             },
         }