feat: Support for update records from SDK

src/argilla/client/feedback/dataset/base.py

@@ -14,7 +14,7 @@
 
 import logging
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union
+from typing import TYPE_CHECKING, Any, Dict, Generic, Iterable, List, Literal, Optional, TypeVar, Union
 
 from pydantic import ValidationError
 
@@ -55,8 +55,10 @@
 
 _LOGGER = logging.getLogger(__name__)
 
+R = TypeVar("R", bound=FeedbackRecord)
 
-class FeedbackDatasetBase(ABC, HuggingFaceDatasetMixin):
+
+class FeedbackDatasetBase(ABC, HuggingFaceDatasetMixin, Generic[R]):
     """Base class with shared functionality for `FeedbackDataset` and `RemoteFeedbackDataset`."""
 
     def __init__(
@@ -166,10 +168,22 @@ def __init__(
 
     @property
     @abstractmethod
-    def records(self) -> Any:
+    def records(self) -> Iterable[R]:
         """Returns the records of the dataset."""
         pass
 
+    @abstractmethod
+    def update_records(self, records: Union[R, List[R]]) -> None:
+        """Updates the records of the dataset.
+
+        Args:
+            records: the records to update the dataset with.
+
+        Raises:
+            ValueError: if the provided `records` are invalid.
+        """
+        pass
+
     @property
     def guidelines(self) -> str:
         """Returns the guidelines for annotating the dataset."""
@@ -364,11 +378,12 @@ def _validate_records(self, records: List[FeedbackRecord]) -> None:
 
     def _parse_and_validate_records(
         self,
-        records: Union[FeedbackRecord, Dict[str, Any], List[Union[FeedbackRecord, Dict[str, Any]]]],
-    ) -> List[FeedbackRecord]:
+        records: Union[R, Dict[str, Any], List[Union[R, Dict[str, Any]]]],
+    ) -> List[R]:
         """Convenient method for calling `_parse_records` and `_validate_records` in sequence."""
         records = self._parse_records(records)
         self._validate_records(records)
+
         return records
 
     @requires_dependencies("datasets")

src/argilla/client/feedback/dataset/local.py

@@ -30,7 +30,7 @@
     )
 
 
-class FeedbackDataset(FeedbackDatasetBase, ArgillaMixin, UnificationMixin):
+class FeedbackDataset(FeedbackDatasetBase["FeedbackRecord"], ArgillaMixin, UnificationMixin):
     def __init__(
         self,
         *,
@@ -127,6 +127,12 @@ def records(self) -> List["FeedbackRecord"]:
         """Returns the records in the dataset."""
         return self._records
 
+    def update_records(self, records: Union["FeedbackRecord", List["FeedbackRecord"]]) -> None:
+        warnings.warn(
+            "`update_records` method only works for `FeedbackDataset` pushed to Argilla. "
+            "If your are working with local data, you can just iterate over the records and update them."
+        )
+
     def __repr__(self) -> str:
         """Returns a string representation of the dataset."""
         return f"<FeedbackDataset fields={self.fields} questions={self.questions} guidelines={self.guidelines}>"

src/argilla/client/feedback/dataset/remote/dataset.py

@@ -66,10 +66,6 @@ def __init__(
                 and/or attributes.
         """
         self._dataset = dataset
-        # TODO: review why this is here !
-        self._question_id_to_name = {question.id: question.name for question in self._dataset.questions}
-        self._question_name_to_id = {value: key for key, value in self._question_id_to_name.items()}
-        # TODO END
 
         if response_status and not isinstance(response_status, list):
             response_status = [response_status]
@@ -106,6 +102,14 @@ def _client(self) -> "httpx.Client":
         """Returns the `httpx.Client` instance that will be used to send requests to Argilla."""
         return self.dataset._client
 
+    @property
+    def _question_id_to_name(self) -> Dict["UUID", str]:
+        return self.dataset._question_id_to_name_id
+
+    @property
+    def _question_name_to_id(self) -> Dict[str, "UUID"]:
+        return self.dataset._question_name_to_id
+
     @allowed_for_roles(roles=[UserRole.owner, UserRole.admin])
     def __len__(self) -> int:
         """Returns the number of records in the current `FeedbackDataset` in Argilla."""
@@ -233,7 +237,7 @@ def _create_from_dataset(
         )
 
 
-class RemoteFeedbackDataset(FeedbackDatasetBase):
+class RemoteFeedbackDataset(FeedbackDatasetBase[RemoteFeedbackRecord]):
     # TODO: Call super method once the base init contains only commons init attributes
     def __init__(
         self,
@@ -304,6 +308,14 @@ def records(self) -> RemoteFeedbackRecords:
         """
         return self._records
 
+    def update_records(self, records: Union[RemoteFeedbackRecord, List[RemoteFeedbackRecord]]) -> None:
+        if not isinstance(records, list):
+            records = [records]
+
+        # TODO: Use the batch version of endpoint once is implemented
+        for record in records:
+            record.update()
+
     @property
     def id(self) -> "UUID":
         """Returns the ID of the dataset in Argilla."""
@@ -334,6 +346,14 @@ def updated_at(self) -> datetime:
         """Returns the datetime when the dataset was last updated in Argilla."""
         return self._updated_at
 
+    @property
+    def _question_id_to_name_id(self) -> Dict["UUID", str]:
+        return {question.id: question.name for question in self._questions}
+
+    @property
+    def _question_name_to_id(self) -> Dict[str, "UUID"]:
+        return {question.name: question.id for question in self._questions}
+
     def __repr__(self) -> str:
         """Returns a string representation of the dataset."""
         return (

src/argilla/client/feedback/dataset/remote/mixins.py

@@ -21,13 +21,13 @@
 from argilla.client.utils import allowed_for_roles
 
 if TYPE_CHECKING:
-    from argilla.client.feedback.dataset.remote.base import RemoteFeedbackRecordsBase
+    from argilla.client.feedback.dataset.remote.dataset import RemoteFeedbackRecords
 
 
 class ArgillaRecordsMixin:
     @allowed_for_roles(roles=[UserRole.owner, UserRole.admin])
     def __getitem__(
-        self: "RemoteFeedbackRecordsBase", key: Union[slice, int]
+        self: "RemoteFeedbackRecords", key: Union[slice, int]
     ) -> Union["RemoteFeedbackRecord", List["RemoteFeedbackRecord"]]:
         """Returns the record(s) at the given index(es) from Argilla.
 
@@ -103,7 +103,7 @@ def __getitem__(
 
     @allowed_for_roles(roles=[UserRole.owner, UserRole.admin])
     def __iter__(
-        self: "RemoteFeedbackRecordsBase",
+        self: "RemoteFeedbackRecords",
     ) -> Iterator["RemoteFeedbackRecord"]:
         """Iterates over the `FeedbackRecord`s of the current `FeedbackDataset` in Argilla."""
         current_batch = 0

src/argilla/client/feedback/schemas/records.py

@@ -88,6 +88,7 @@ def to_server_payload(self) -> Dict[str, Any]:
         """Method that will be used to create the payload that will be sent to Argilla
         to create a `ResponseSchema` for a `FeedbackRecord`."""
         return {
+            # UUID is not json serializable!!!
             "user_id": self.user_id,
             "values": {question_name: value.dict() for question_name, value in self.values.items()}
             if self.values is not None
@@ -194,9 +195,7 @@ class FeedbackRecord(BaseModel):
     fields: Dict[str, Union[str, None]]
     metadata: Dict[str, Any] = Field(default_factory=dict)
     responses: List[ResponseSchema] = Field(default_factory=list)
-    suggestions: Union[Tuple[SuggestionSchema], List[SuggestionSchema]] = Field(
-        default_factory=tuple, allow_mutation=False
-    )
+    suggestions: Union[Tuple[SuggestionSchema], List[SuggestionSchema]] = Field(default_factory=tuple)
     external_id: Optional[str] = None
 
     _unified_responses: Optional[Dict[str, List["UnifiedValueSchema"]]] = PrivateAttr(default_factory=dict)

src/argilla/client/feedback/schemas/remote/records.py

@@ -22,7 +22,6 @@
 from argilla.client.feedback.schemas.records import FeedbackRecord, ResponseSchema, SuggestionSchema
 from argilla.client.feedback.schemas.remote.shared import RemoteSchema
 from argilla.client.sdk.users.models import UserRole
-from argilla.client.sdk.v1.datasets import api as datasets_api_v1
 from argilla.client.sdk.v1.records import api as records_api_v1
 from argilla.client.sdk.v1.suggestions import api as suggestions_api_v1
 from argilla.client.utils import allowed_for_roles
@@ -31,7 +30,7 @@
     import httpx
 
     from argilla.client.sdk.v1.datasets.models import FeedbackResponseModel, FeedbackSuggestionModel
-    from argilla.client.sdk.v1.records.models import FeedbackItemModel
+    from argilla.client.sdk.v1.records.models import FeedbackRecordModel
 
 
 class RemoteSuggestionSchema(SuggestionSchema, RemoteSchema):
@@ -95,12 +94,16 @@ def from_api(cls, payload: "FeedbackResponseModel") -> "RemoteResponseSchema":
         return RemoteResponseSchema(
             user_id=payload.user_id,
             values=payload.values,
+            # TODO: Review type mismatch between API and SDK
             status=payload.status,
             inserted_at=payload.inserted_at,
             updated_at=payload.updated_at,
         )
 
 
+AllowedSuggestionSchema = Union[RemoteSuggestionSchema, SuggestionSchema]
+
+
 class RemoteFeedbackRecord(FeedbackRecord, RemoteSchema):
     """Schema for the records of a `RemoteFeedbackDataset`.
 
@@ -117,37 +120,29 @@ class RemoteFeedbackRecord(FeedbackRecord, RemoteSchema):
             question. Defaults to an empty list.
     """
 
+    # TODO: remote record should receive a dataset instead of this
     question_name_to_id: Optional[Dict[str, UUID]] = Field(..., exclude=True, repr=False)
 
     responses: List[RemoteResponseSchema] = Field(default_factory=list)
-    suggestions: Union[Tuple[RemoteSuggestionSchema], List[RemoteSuggestionSchema]] = Field(
-        default_factory=tuple, allow_mutation=False
-    )
+    suggestions: Union[Tuple[AllowedSuggestionSchema], List[AllowedSuggestionSchema]] = Field(default_factory=tuple)
 
     class Config:
+        allow_mutation = True
         validate_assignment = True
 
-    def __update_suggestions(
+    def __normalize_suggestions_to_update(
         self,
         suggestions: Union[
-            RemoteSuggestionSchema,
-            List[RemoteSuggestionSchema],
-            SuggestionSchema,
-            List[SuggestionSchema],
-            Dict[str, Any],
-            List[Dict[str, Any]],
+            Dict[str, Any], List[Dict[str, Any]], AllowedSuggestionSchema, List[AllowedSuggestionSchema]
         ],
-    ) -> None:
-        """Updates the suggestions for the record in Argilla. Note that the suggestions
-        must exist in Argilla to be updated.
-
-        Note that this method will update the record in Argilla directly.
+    ) -> List[AllowedSuggestionSchema]:
+        """
+        Normalizes the suggestions to update.
 
         Args:
-            suggestions: can be a single `RemoteSuggestionSchema` or `SuggestionSchema`,
-                a list of `RemoteSuggestionSchema` or `SuggestionSchema`, a single
-                dictionary, or a list of dictionaries. If a dictionary is provided,
-                it will be converted to a `RemoteSuggestionSchema` internally.
+            suggestions: can be a single `RemoteSuggestionSchema` or `SuggestionSchema`, a dictionary, a list of
+            `RemoteSuggestionSchema` or `SuggestionSchema`, or a list of dictionaries. If a dictionary is provided,
+            it will be converted to a `SuggestionSchema` internally.
         """
         if isinstance(suggestions, (dict, SuggestionSchema)):
             suggestions = [suggestions]
@@ -203,34 +198,40 @@ def __update_suggestions(
             else:
                 new_suggestions[suggestion.question_name] = suggestion
 
-        for suggestion in new_suggestions.values():
+        return list(new_suggestions.values())
+
+    def __update_suggestions(self, suggestions: List[AllowedSuggestionSchema]) -> None:
+        """Updates the suggestions for the record in Argilla.
+
+        Note that this method will update the record in Argilla directly.
+
+        Args:
+            suggestions: can be a list of `RemoteSuggestionSchema` or `SuggestionSchema`.
+        """
+
+        pushed_suggestions = []
+
+        for suggestion in suggestions:
             if isinstance(suggestion, RemoteSuggestionSchema):
                 suggestion = suggestion.to_local()
-            pushed_suggestion = datasets_api_v1.set_suggestion(
+            # TODO: review the existence of bulk endpoint for record suggestions
+            pushed_suggestion = records_api_v1.set_suggestion(
                 client=self.client,
                 record_id=self.id,
                 **suggestion.to_server_payload(question_name_to_id=self.question_name_to_id),
             )
-            existing_suggestions[suggestion.question_name] = RemoteSuggestionSchema.from_api(
-                payload=pushed_suggestion.parsed,
-                question_id_to_name={value: key for key, value in self.question_name_to_id.items()},
-                client=self.client,
+            pushed_suggestions.append(
+                RemoteSuggestionSchema.from_api(
+                    payload=pushed_suggestion.parsed,
+                    question_id_to_name={value: key for key, value in self.question_name_to_id.items()},
+                    client=self.client,
+                )
             )
 
-        self.__dict__["suggestions"] = tuple(existing_suggestions.values())
+        self.__dict__["suggestions"] = tuple(pushed_suggestions)
 
     @allowed_for_roles(roles=[UserRole.owner, UserRole.admin])
-    def update(
-        self,
-        suggestions: Union[
-            RemoteSuggestionSchema,
-            List[RemoteSuggestionSchema],
-            SuggestionSchema,
-            List[SuggestionSchema],
-            Dict[str, Any],
-            List[Dict[str, Any]],
-        ],
-    ) -> None:
+    def update(self, suggestions: Optional[AllowedSuggestionSchema] = None) -> None:
         """Update a `RemoteFeedbackRecord`. Currently just `suggestions` are supported.
 
         Note that this method will update the record in Argilla directly.
@@ -244,7 +245,27 @@ def update(
         Raises:
             PermissionError: if the user does not have either `owner` or `admin` role.
         """
-        self.__update_suggestions(suggestions=suggestions)
+        if suggestions:
+            suggestions = self.__normalize_suggestions_to_update(suggestions)
+        else:
+            suggestions = suggestions or [s for s in self.suggestions]
+
+        self.__updated_record_data()
+        if suggestions:
+            self.__update_suggestions(suggestions=suggestions)
+
+    def __updated_record_data(self) -> None:
+        response = records_api_v1.update_record(self.client, self.id, self.to_server_payload())
+
+        updated_record = self.from_api(
+            payload=response.parsed,
+            question_id_to_name={value: key for key, value in self.question_name_to_id.items()}
+            if self.question_name_to_id
+            else None,
+            client=self.client,
+        )
+
+        self.__dict__.update(updated_record.__dict__)
 
     @allowed_for_roles(roles=[UserRole.owner, UserRole.admin])
     def delete_suggestions(self, suggestions: Union[RemoteSuggestionSchema, List[RemoteSuggestionSchema]]) -> None:
@@ -282,7 +303,8 @@ def delete_suggestions(self, suggestions: Union[RemoteSuggestionSchema, List[Rem
             self.__dict__["suggestions"] = tuple(existing_suggestions.values())
         except Exception as e:
             raise RuntimeError(
-                f"Failed to delete suggestions with IDs `{[suggestion.id for suggestion in delete_suggestions]}` from record with ID `{self.id}` from Argilla."
+                f"Failed to delete suggestions with IDs `{[suggestion.id for suggestion in delete_suggestions]}` from "
+                f"record with ID `{self.id}` from Argilla."
             ) from e
 
     @allowed_for_roles(roles=[UserRole.owner, UserRole.admin])
@@ -316,7 +338,7 @@ def to_local(self) -> "FeedbackRecord":
     @classmethod
     def from_api(
         cls,
-        payload: "FeedbackItemModel",
+        payload: "FeedbackRecordModel",
         question_id_to_name: Optional[Dict[UUID, str]] = None,
         client: Optional["httpx.Client"] = None,
     ) -> "RemoteFeedbackRecord":

src/argilla/client/feedback/schemas/remote/shared.py

@@ -21,6 +21,7 @@
 
 
 class RemoteSchema(BaseModel, ABC):
+    # TODO(@alvarobartt): Review optional id configuration for remote schemas
     id: Optional[UUID] = None
     client: Optional[httpx.Client] = None
 
@@ -30,6 +31,7 @@ def _client(self) -> Optional[httpx.Client]:
         return self.client
 
     class Config:
+        # TODO(@alvarobart) Not sure if we need this at this level
         allow_mutation = False
         arbitrary_types_allowed = True