fix(#1658): align agents for predictions annotations (#1742)

* fix: align agents for predictions annotations

* fix: empty agent for annotations dict but check it on single annotation

src/rubrix/server/daos/models/records.py

@@ -22,6 +22,7 @@
 from rubrix._constants import MAX_KEYWORD_LENGTH
 from rubrix.server.commons.models import PredictionStatus, TaskStatus, TaskType
 from rubrix.server.daos.backend.search.model import BackendRecordsQuery, SortConfig
+from rubrix.server.errors import ValidationError
 from rubrix.server.helpers import flatten_dict
 from rubrix.utils import limit_value_length
 
@@ -76,7 +77,14 @@ def update_annotation(values, annotation_field: str):
         annotation = values.get(annotation_field)
         annotations = values.get(field_to_update) or {}
 
+        if annotations:
+            for key, value in annotations.items():
+                value.agent = None  # Maybe we want key and agents with different values
+
         if annotation:
+            if not annotation.agent:
+                raise AssertionError("Agent must be defined!")
+
             annotations.update(
                 {
                     annotation.agent: annotation.__class__.parse_obj(
@@ -87,12 +95,11 @@ def update_annotation(values, annotation_field: str):
             values[field_to_update] = annotations
 
         if annotations and not annotation:
-            for key, value in annotations.items():
-                new_annotation = value.__class__.parse_obj(
-                    {**value.dict(), "agent": key}
-                )
-                values[annotation_field] = new_annotation
-                break
+            # set first annotation
+            key, value = list(annotations.items())[0]
+            values[annotation_field] = value.__class__(
+                agent=key, **value.dict(exclude={"agent"})
+            )
 
         return values
 

tests/server/text_classification/test_model.py

@@ -398,3 +398,38 @@ def test_annotated_without_labels_for_multilabel():
     )
 
     assert record.predicted == PredictionStatus.OK
+
+
+def test_using_predictions_dict():
+    record = ServiceTextClassificationRecord(
+        inputs={"text": "this is a text"},
+        predictions={
+            "carl": TextClassificationAnnotation(
+                agent="wat at", labels=[ClassPrediction(class_label="YES")]
+            ),
+            "BOB": TextClassificationAnnotation(
+                agent="wot wot", labels=[ClassPrediction(class_label="NO")]
+            ),
+        },
+    )
+
+    assert record.prediction.dict() == {
+        "agent": "carl",
+        "labels": [{"class_label": "YES", "score": 1.0}],
+    }
+    assert record.predictions == {
+        "BOB": TextClassificationAnnotation(labels=[ClassPrediction(class_label="NO")]),
+        "carl": TextClassificationAnnotation(
+            labels=[ClassPrediction(class_label="YES")]
+        ),
+    }
+
+
+def test_with_no_agent_at_all():
+    with pytest.raises(ValidationError):
+        ServiceTextClassificationRecord(
+            inputs={"text": "this is a text"},
+            prediction=TextClassificationAnnotation(
+                labels=[ClassPrediction(class_label="YES")]
+            ),
+        )