Feature/persistent-record-mappings-in-settings (#5466)

# 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 brings changes to persist the record mapping dict used in the
log function within dataset settings.

**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: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: burtenshaw

argilla/src/argilla/datasets/_resource.py

@@ -73,7 +73,7 @@ def __init__(
         self._model = DatasetModel(name=name)
         self._settings = settings._copy() if settings else Settings(_dataset=self)
         self._settings.dataset = self
-        self.__records = DatasetRecords(client=self._client, dataset=self)
+        self.__records = DatasetRecords(client=self._client, dataset=self, mapping=self._settings.mapping)
 
     #####################
     #  Properties       #

argilla/src/argilla/records/_dataset_records.py

@@ -139,14 +139,17 @@ class DatasetRecords(Iterable[Record], LoggingMixin):
     DEFAULT_BATCH_SIZE = 256
     DEFAULT_DELETE_BATCH_SIZE = 64
 
-    def __init__(self, client: "Argilla", dataset: "Dataset"):
+    def __init__(
+        self, client: "Argilla", dataset: "Dataset", mapping: Optional[Dict[str, Union[str, Sequence[str]]]] = None
+    ):
         """Initializes a DatasetRecords object with a client and a dataset.
         Args:
             client: An Argilla client object.
             dataset: A Dataset object.
         """
         self.__client = client
         self.__dataset = dataset
+        self._mapping = mapping or {}
         self._api = self.__client.api.records
 
     def __iter__(self):
@@ -380,6 +383,7 @@ def _ingest_records(
     ) -> List[RecordModel]:
         """Ingests records from a list of dictionaries, a Hugging Face Dataset, or a list of Record objects."""
 
+        mapping = mapping or self._mapping
         if len(records) == 0:
             raise ValueError("No records provided to ingest.")
 

argilla/src/argilla/settings/_resource.py

@@ -50,6 +50,7 @@ def __init__(
         guidelines: Optional[str] = None,
         allow_extra_metadata: bool = False,
         distribution: Optional[TaskDistribution] = None,
+        mapping: Optional[Dict[str, Union[str, Sequence[str]]]] = None,
         _dataset: Optional["Dataset"] = None,
     ) -> None:
         """
@@ -64,11 +65,13 @@ def __init__(
                 Dataset. Defaults to False.
             distribution (TaskDistribution): The annotation task distribution configuration.
                 Default to DEFAULT_TASK_DISTRIBUTION
+            mapping (Dict[str, Union[str, Sequence[str]]]): A dictionary that maps incoming data names to Argilla dataset attributes in DatasetRecords.
         """
         super().__init__(client=_dataset._client if _dataset else None)
 
         self._dataset = _dataset
         self._distribution = distribution
+        self._mapping = mapping
         self.__guidelines = self.__process_guidelines(guidelines)
         self.__allow_extra_metadata = allow_extra_metadata
 
@@ -137,6 +140,14 @@ def distribution(self) -> TaskDistribution:
     def distribution(self, value: TaskDistribution) -> None:
         self._distribution = value
 
+    @property
+    def mapping(self) -> Dict[str, Union[str, Sequence[str]]]:
+        return self._mapping
+
+    @mapping.setter
+    def mapping(self, value: Dict[str, Union[str, Sequence[str]]]):
+        self._mapping = value
+
     @property
     def dataset(self) -> "Dataset":
         return self._dataset
@@ -220,6 +231,7 @@ def serialize(self):
                 "metadata": self.metadata.serialize(),
                 "allow_extra_metadata": self.allow_extra_metadata,
                 "distribution": self.distribution.to_dict(),
+                "mapping": self.mapping,
             }
         except Exception as e:
             raise ArgillaSerializeError(f"Failed to serialize the settings. {e.__class__.__name__}") from e
@@ -271,6 +283,7 @@ def _from_dict(cls, settings_dict: dict) -> "Settings":
         guidelines = settings_dict.get("guidelines")
         distribution = settings_dict.get("distribution")
         allow_extra_metadata = settings_dict.get("allow_extra_metadata")
+        mapping = settings_dict.get("mapping")
 
         questions = [question_from_dict(question) for question in settings_dict.get("questions", [])]
         fields = [_field_from_dict(field) for field in fields]
@@ -280,6 +293,9 @@ def _from_dict(cls, settings_dict: dict) -> "Settings":
         if distribution:
             distribution = TaskDistribution.from_dict(distribution)
 
+        if mapping:
+            mapping = cls._validate_mapping(mapping)
+
         return cls(
             questions=questions,
             fields=fields,
@@ -288,6 +304,7 @@ def _from_dict(cls, settings_dict: dict) -> "Settings":
             guidelines=guidelines,
             allow_extra_metadata=allow_extra_metadata,
             distribution=distribution,
+            mapping=mapping,
         )
 
     def _copy(self) -> "Settings":
@@ -362,6 +379,18 @@ def _validate_duplicate_names(self) -> None:
                     )
                 dataset_properties_by_name[property.name] = property
 
+    @classmethod
+    def _validate_mapping(cls, mapping: Dict[str, Union[str, Sequence[str]]]) -> None:
+        validate_mapping = {}
+        for key, value in mapping.items():
+            if isinstance(value, str):
+                validate_mapping[key] = value
+            elif isinstance(value, list) or isinstance(value, tuple):
+                validate_mapping[key] = tuple(value)
+            else:
+                raise SettingsError(f"Invalid mapping value for key {key!r}: {value}")
+        return validate_mapping
+
     def __process_guidelines(self, guidelines):
         if guidelines is None:
             return guidelines

argilla/tests/unit/test_settings/test_settings_mapping_record_ingestion.py

@@ -0,0 +1,88 @@
+# Copyright 2024-present, Argilla, Inc.
+#
+# 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 tempfile import TemporaryDirectory
+from uuid import uuid4
+
+import pytest
+
+import argilla as rg
+
+
+@pytest.fixture
+def dataset():
+    mock_mapping = {
+        "true_label": "label.response",
+        "my_label": "label.suggestion.value",
+        "score": "label.suggestion.score",
+        "model": "label.suggestion.agent",
+        "my_prompt": ("prompt_field", "prompt_question"),
+    }
+    settings = rg.Settings(
+        fields=[rg.TextField(name="prompt_field")],
+        questions=[
+            rg.LabelQuestion(name="label", labels=["negative", "positive"]),
+            rg.TextQuestion(name="prompt_question"),
+        ],
+        metadata=[rg.FloatMetadataProperty(name="score")],
+        vectors=[rg.VectorField(name="vector", dimensions=3)],
+        mapping=mock_mapping,
+    )
+    workspace = rg.Workspace(name="workspace", id=uuid4())
+    dataset = rg.Dataset(
+        name="test_dataset",
+        settings=settings,
+        workspace=workspace,
+    )
+    return dataset
+
+
+def test_settings_with_record_mapping(dataset):
+    mock_user_id = uuid4()
+    record_api_models = dataset.records._ingest_records(
+        records=[
+            {
+                "my_prompt": "What is the capital of France?",
+                "my_label": "positive",
+                "true_label": "positive",
+                "score": 0.9,
+                "model": "model_name",
+            }
+        ],
+        user_id=mock_user_id,
+    )
+    record = record_api_models[0]
+    assert record.fields["prompt_field"] == "What is the capital of France?"
+    assert record.suggestions[0].value == "positive"
+    assert record.suggestions[0].question_name == "label"
+    assert record.suggestions[0].score == 0.9
+    assert record.suggestions[0].agent == "model_name"
+    assert record.responses[0].values["label"]["value"] == "positive"
+    assert record.responses[0].user_id == mock_user_id
+
+    record = record_api_models[0]
+    suggestions = [s.value for s in record.suggestions]
+    assert record.fields["prompt_field"] == "What is the capital of France?"
+    assert "positive" in suggestions
+    assert "What is the capital of France?" in suggestions
+
+
+def test_settings_with_record_mapping_export(dataset):
+    with TemporaryDirectory() as temp_dir:
+        path = f"{temp_dir}/test_dataset.json"
+        dataset.settings.to_json(path)
+        loaded_settings = rg.Settings.from_json(path)
+
+    assert dataset.settings.mapping == loaded_settings.mapping
+    assert dataset.settings == loaded_settings