[ENHANCEMENT] argilla: Add limit argument when getting records (#5525)

# 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 adds a new `limit` argument when getting records from the
Argilla server.

**Type of change**
<!-- Please delete options that are not relevant. Remember to title the
PR according to the type of change -->

- Improvement (change adding some improvement to an existing
functionality)
- Documentation update

**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: burtenshaw <[email protected]>

Author: frascuchon

argilla/CHANGELOG.md

@@ -16,6 +16,10 @@ These are the section headers that we use:
 
 ## [Unreleased]()
 
+### Added
+
+- Added `limit` argument when fetching records. ([#5525](https://github.com/argilla-io/argilla/pull/5525)
+
 ### Fixed
 
 - Fixed the deployment yaml used to create a new Argilla server in K8s. Added `USERNAME` and `PASSWORD` to the environment variables of pod template. ([#5434](https://github.com/argilla-io/argilla/issues/5434))

argilla/src/argilla/records/_dataset_records.py

@@ -52,6 +52,7 @@ def __init__(
         with_suggestions: bool = False,
         with_responses: bool = False,
         with_vectors: Optional[Union[str, List[str], bool]] = None,
+        limit: Optional[int] = None,
     ):
         self.__dataset = dataset
         self.__client = client
@@ -62,19 +63,41 @@ def __init__(
         self.__with_responses = with_responses
         self.__with_vectors = with_vectors
         self.__records_batch = []
+        self.__limit = limit
+
+        if self.__limit is not None and self.__limit <= 0:
+            warnings.warn(f"Limit {self.__limit} is invalid: must be greater than 0. Setting limit to 1.")
+            self.__limit = 1
+
+        if self.__limit is not None and self.__limit < self.__batch_size:
+            self.__batch_size = self.__limit
 
     def __iter__(self):
         return self
 
     def __next__(self) -> Record:
+        if self._limit_reached():
+            raise StopIteration()
+
         if not self._has_local_records():
             self._fetch_next_batch()
             if not self._has_local_records():
                 raise StopIteration()
+
         return self._next_record()
 
+    def _limit_reached(self) -> bool:
+        if self.__limit is None:
+            return False
+        return self.__limit <= 0
+
     def _next_record(self) -> Record:
-        return self.__records_batch.pop(0)
+        record = self.__records_batch.pop(0)
+
+        if self.__limit is not None:
+            self.__limit -= 1
+
+        return record
 
     def _has_local_records(self) -> bool:
         return len(self.__records_batch) > 0
@@ -170,6 +193,7 @@ def __call__(
         with_suggestions: bool = True,
         with_responses: bool = True,
         with_vectors: Optional[Union[List, bool, str]] = None,
+        limit: Optional[int] = None,
     ) -> DatasetRecordsIterator:
         """Returns an iterator over the records in the dataset on the server.
 
@@ -182,6 +206,7 @@ def __call__(
             with_vectors: A list of vector names to include in the records. The default is None.
                 If a list is provided, only the specified vectors will be included.
                 If True is provided, all vectors will be included.
+            limit: The maximum number of records to fetch. The default is None.
 
         Returns:
             An iterator over the records in the dataset on the server.
@@ -202,6 +227,7 @@ def __call__(
             with_suggestions=with_suggestions,
             with_responses=with_responses,
             with_vectors=with_vectors,
+            limit=limit,
         )
 
     def __repr__(self) -> str:

argilla/tests/integration/test_list_records.py

@@ -66,6 +66,93 @@ def test_list_records_with_start_offset(client: Argilla, dataset: Dataset):
     ]
 
 
+def test_list_records_with_limit(client: Argilla, dataset: Dataset):
+    dataset.records.log(
+        [
+            {"text": "The record text field", "id": 1},
+            {"text": "The record text field", "id": 2},
+            {"text": "The record text field", "id": 3},
+            {"text": "The record text field", "id": 4},
+            {"text": "The record text field", "id": 5},
+        ]
+    )
+
+    records = list(dataset.records(limit=2))
+    assert len(records) == 2
+
+    assert [record.to_dict() for record in records] == [
+        {
+            "_server_id": str(records[0]._server_id),
+            "fields": {"text": "The record text field"},
+            "id": "1",
+            "status": "pending",
+            "metadata": {},
+            "responses": {},
+            "suggestions": {},
+            "vectors": {},
+        },
+        {
+            "_server_id": str(records[1]._server_id),
+            "fields": {"text": "The record text field"},
+            "id": "2",
+            "status": "pending",
+            "metadata": {},
+            "responses": {},
+            "suggestions": {},
+            "vectors": {},
+        },
+    ]
+
+
+def test_list_records_with_limit_greater_than_batch_size(client: Argilla, dataset: Dataset):
+    dataset.records.log(
+        [
+            {"text": "The record text field", "id": 1},
+            {"text": "The record text field", "id": 2},
+            {"text": "The record text field", "id": 3},
+            {"text": "The record text field", "id": 4},
+            {"text": "The record text field", "id": 5},
+        ]
+    )
+
+    records = list(dataset.records(limit=2, batch_size=1))
+
+    assert len(records) == 2
+    assert records[0].id == "1"
+    assert records[1].id == "2"
+
+
+@pytest.mark.parametrize("limit", [0, -1, -10])
+def test_list_records_with_invalid_limit(client: Argilla, dataset: Dataset, limit: int):
+    dataset.records.log(
+        [
+            {"text": "The record text field", "id": 1},
+            {"text": "The record text field", "id": 2},
+            {"text": "The record text field", "id": 3},
+            {"text": "The record text field", "id": 4},
+            {"text": "The record text field", "id": 5},
+        ]
+    )
+    with pytest.warns(UserWarning, match=f"Limit {limit} is invalid: must be greater than 0. Setting limit to 1."):
+        records = list(dataset.records(limit=limit))
+        assert len(records) == 1
+
+
+def test_list_records_with_limit_greater_than_total(client: Argilla, dataset: Dataset):
+    dataset.records.log(
+        [
+            {"text": "The record text field", "id": 1},
+            {"text": "The record text field", "id": 2},
+            {"text": "The record text field", "id": 3},
+            {"text": "The record text field", "id": 4},
+            {"text": "The record text field", "id": 5},
+        ]
+    )
+
+    records = list(dataset.records(limit=10))
+    assert len(records) == 5
+
+
 def test_get_record_by_id(client: Argilla, dataset: Dataset):
     dataset.records.log(
         [