Update argilla integration to use argilla_sdk v2

src/distilabel/steps/argilla/base.py

@@ -20,23 +20,24 @@
 from pydantic import Field, PrivateAttr, SecretStr
 
 try:
-    import argilla as rg
+    import argilla_sdk as rg
 except ImportError:
     pass
 
 from distilabel.mixins.runtime_parameters import RuntimeParameter
 from distilabel.steps.base import Step, StepInput
 
 if TYPE_CHECKING:
-    from argilla.client.feedback.dataset.remote.dataset import RemoteFeedbackDataset
+    from argilla_sdk import Argilla, Dataset
 
     from distilabel.steps.typing import StepOutput
 
 
+_ARGILLA_API_URL_ENV_VAR_NAME = "ARGILLA_API_URL"
 _ARGILLA_API_KEY_ENV_VAR_NAME = "ARGILLA_API_KEY"
 
 
-class Argilla(Step, ABC):
+class _Argilla(Step, ABC):
     """Abstract step that provides a class to subclass from, that contains the boilerplate code
     required to interact with Argilla, as well as some extra validations on top of it. It also defines
     the abstract methods that need to be implemented in order to add a new dataset type as a step.
@@ -75,50 +76,67 @@ class Argilla(Step, ABC):
     )
 
     api_url: Optional[RuntimeParameter[str]] = Field(
-        default_factory=lambda: os.getenv("ARGILLA_API_URL"),
+        default_factory=lambda: os.getenv(_ARGILLA_API_URL_ENV_VAR_NAME),
         description="The base URL to use for the Argilla API requests.",
     )
     api_key: Optional[RuntimeParameter[SecretStr]] = Field(
         default_factory=lambda: os.getenv(_ARGILLA_API_KEY_ENV_VAR_NAME),
         description="The API key to authenticate the requests to the Argilla API.",
     )
 
-    _rg_dataset: Optional["RemoteFeedbackDataset"] = PrivateAttr(...)
+    _client: Optional["Argilla"] = PrivateAttr(...)
+    _dataset: Optional["Dataset"] = PrivateAttr(...)
 
     def model_post_init(self, __context: Any) -> None:
         """Checks that the Argilla Python SDK is installed, and then filters the Argilla warnings."""
         super().model_post_init(__context)
 
         try:
-            import argilla as rg  # noqa
+            import argilla_sdk as rg  # noqa
         except ImportError as ie:
             raise ImportError(
-                "Argilla is not installed. Please install it using `pip install argilla`."
+                "Argilla is not installed. Please install it using `pip install argilla_sdk --upgrade`."
             ) from ie
 
         warnings.filterwarnings("ignore")
 
-    def _rg_init(self) -> None:
+    def _client_init(self) -> None:
         """Initializes the Argilla API client with the provided `api_url` and `api_key`."""
         try:
-            if "hf.space" in self.api_url and "HF_TOKEN" in os.environ:
-                headers = {"Authorization": f"Bearer {os.environ['HF_TOKEN']}"}
-            else:
-                headers = None
-            rg.init(
+            self._client = rg.Argilla(  # type: ignore
                 api_url=self.api_url,
-                api_key=self.api_key.get_secret_value(),
-                extra_headers=headers,
-            )  # type: ignore
+                api_key=self.api_key.get_secret_value(),  # type: ignore
+                headers={"Authorization": f"Bearer {os.environ['HF_TOKEN']}"}
+                if isinstance(self.api_url, str)
+                and "hf.space" in self.api_url
+                and "HF_TOKEN" in os.environ
+                else {},
+            )
         except Exception as e:
             raise ValueError(f"Failed to initialize the Argilla API: {e}") from e
 
-    def _rg_dataset_exists(self) -> bool:
-        """Checks if the dataset already exists in Argilla."""
-        return self.dataset_name in [
-            dataset.name
-            for dataset in rg.FeedbackDataset.list(workspace=self.dataset_workspace)  # type: ignore
-        ]
+    @property
+    def _dataset_exists_in_workspace(self) -> bool:
+        """Checks if the dataset already exists in Argilla in the provided workspace if any."""
+        return (
+            True
+            if self._client.datasets(  # type: ignore
+                name=self.dataset_name,  # type: ignore
+                workspace=self._client.workspaces(name=self.dataset_workspace)  # type: ignore
+                if self.dataset_workspace is not None
+                else None,
+            ).exists()
+            is not None
+            else False
+        )
+        # Alternative way but the above should work
+        # return self.dataset_name in (
+        #     [
+        #         dataset.name for dataset in self._client.workspaces(self.dataset_workspace)  # type: ignore
+        #     ]
+        #     if self.dataset_workspace is not None
+        #     else [dataset.name for dataset in self._client.datasets.list()]  # type: ignore
+        # )
 
     @property
     def outputs(self) -> List[str]:
@@ -133,7 +151,7 @@ def load(self) -> None:
         """
         super().load()
 
-        self._rg_init()
+        self._client_init()
 
     @property
     @abstractmethod

src/distilabel/steps/argilla/text_generation.py

@@ -19,18 +19,18 @@
 from typing_extensions import override
 
 try:
-    import argilla as rg
+    import argilla_sdk as rg
 except ImportError:
     pass
 
-from distilabel.steps.argilla.base import Argilla
+from distilabel.steps.argilla.base import _Argilla
 from distilabel.steps.base import StepInput
 
 if TYPE_CHECKING:
     from distilabel.steps.typing import StepOutput
 
 
-class TextGenerationToArgilla(Argilla):
+class TextGenerationToArgilla(_Argilla):
     """Creates a text generation dataset in Argilla.
 
     `Step` that creates a dataset in Argilla during the load phase, and then pushes the input
@@ -74,26 +74,30 @@ def load(self) -> None:
         self._instruction = self.input_mappings.get("instruction", "instruction")
         self._generation = self.input_mappings.get("generation", "generation")
 
-        if self._rg_dataset_exists():
-            _rg_dataset = rg.FeedbackDataset.from_argilla(  # type: ignore
-                name=self.dataset_name,
-                workspace=self.dataset_workspace,
+        if self._dataset_exists_in_workspace:
+            _dataset = self._client.datasets(  # type: ignore
+                name=self.dataset_name,  # type: ignore
+                workspace=self._client.workspaces(name=self.dataset_workspace)  # type: ignore
+                if self.dataset_workspace is not None
+                else None,
             )
 
-            for field in _rg_dataset.fields:
+            for field in _dataset.fields:
+                if not isinstance(field, rg.TextField):  # type: ignore
+                    continue
                 if (
                     field.name not in [self._id, self._instruction, self._generation]
                     and field.required
                 ):
                     raise ValueError(
                         f"The dataset {self.dataset_name} in the workspace {self.dataset_workspace} already exists,"
                         f" but contains at least a required field that is neither `{self._id}`, `{self._instruction}`"
-                        f", nor `{self._generation}`."
+                        f", nor `{self._generation}`, so it cannot be reused for this dataset."
                     )
 
-            self._rg_dataset = _rg_dataset
+            self._dataset = _dataset
         else:
-            _rg_dataset = rg.FeedbackDataset(  # type: ignore
+            _settings = rg.Settings(  # type: ignore
                 fields=[
                     rg.TextField(name=self._id, title=self._id),  # type: ignore
                     rg.TextField(name=self._instruction, title=self._instruction),  # type: ignore
@@ -103,14 +107,18 @@ def load(self) -> None:
                     rg.LabelQuestion(  # type: ignore
                         name="quality",
                         title=f"What's the quality of the {self._generation} for the given {self._instruction}?",
-                        labels={"bad": "👎", "good": "👍"},
+                        labels={"bad": "👎", "good": "👍"},  # type: ignore
                     )
                 ],
             )
-            self._rg_dataset = _rg_dataset.push_to_argilla(
-                name=self.dataset_name,  # type: ignore
+            _dataset = rg.Dataset(  # type: ignore
+                name=self.dataset_name,
                 workspace=self.dataset_workspace,
+                settings=_settings,
+                client=self._client,
             )
+            self._logger.info(f"Creating the dataset {self.dataset_name} in Argilla.")
+            self._dataset = _dataset.create()
 
     @property
     def inputs(self) -> List[str]:
@@ -151,13 +159,13 @@ def process(self, inputs: StepInput) -> "StepOutput":  # type: ignore
                 generations_set.add(generation)
 
                 records.append(
-                    rg.FeedbackRecord(  # type: ignore
+                    rg.Record(  # type: ignore
                         fields={
                             self._id: instruction_id,
                             self._instruction: input["instruction"],
                             self._generation: generation,
                         },
-                    )
+                    ),
                 )
-        self._rg_dataset.add_records(records)  # type: ignore
+        self._dataset.log(records)  # type: ignore
         yield inputs