[OPIK-555] Add retries for get dataset items operation

.github/workflows/documentation_cookbook_tests.yml

@@ -1,6 +1,15 @@
 name: Documentation - Test cookbooks
 on:
   workflow_dispatch:
+    inputs:
+      install_opik:
+        description: 'Enable opik installation from source files'
+        required: false
+        default: 'false'
+        type: choice
+        options:
+          - 'false'
+          - 'true'
   schedule:
     - cron: '0 0 * * *'  # Run once a day at midnight UTC
 jobs:
@@ -43,8 +52,9 @@ jobs:
         run: |
           python -m pip install --upgrade pip
           python -m pip install -U ipython nbconvert
-      # - name: Install opik # Uncomment if you want to run cookbooks with local version of opik
-      #   run: pip install sdks/python
+      - name: Install opik from source files (optional)
+        if: ${{ github.event_name == 'workflow_dispatch' && inputs.install_opik == 'true' }}
+        run: pip install sdks/python
       - name: Prepare env variables
         run: |
           directory=$(dirname -- "${NOTEBOOK_TO_TEST}")

sdks/python/src/opik/api_objects/dataset/dataset.py

@@ -6,7 +6,7 @@
 from opik.rest_api.types import dataset_item_write as rest_dataset_item
 from opik.message_processing.batching import sequence_splitter
 from opik import exceptions, config
-
+from opik.rest_client_configurator import retry_decorators
 from .. import constants
 from . import dataset_item, converters
 
@@ -206,6 +206,7 @@ def get_items(self, nb_samples: Optional[int] = None) -> List[Dict[str, Any]]:
 
         return dataset_items_as_dicts
 
+    @retry_decorators.connection_retry
     def __internal_api__get_items_as_dataclasses__(
         self, nb_samples: Optional[int] = None
     ) -> List[dataset_item.DatasetItem]:

sdks/python/src/opik/rest_client_configurator/retry_decorators.py

@@ -3,7 +3,7 @@
 
 connection_retry = tenacity.retry(
     stop=tenacity.stop_after_attempt(3),
-    wait=tenacity.wait_exponential(multiplier=1, min=1, max=10),
+    wait=tenacity.wait_exponential(multiplier=2, min=3, max=10),
     retry=tenacity.retry_if_exception_type(
         (
             httpx.RemoteProtocolError,  # handle retries for expired connections