rapidata 2.17.0__py3-none-any.whl → 2.18.0__py3-none-any.whl

rapidata/api_client/rest.py

@@ -11,114 +11,86 @@
     Do not edit the class manually.
 """  # noqa: E501
 
-
 import io
 import json
 import re
-import ssl
+from typing import Dict, Optional
 
-import urllib3
+import requests
+from authlib.integrations.requests_client import OAuth2Session
+from requests.adapters import HTTPAdapter
+from urllib3 import Retry
 
 from rapidata.api_client.exceptions import ApiException, ApiValueError
 
-SUPPORTED_SOCKS_PROXIES = {"socks5", "socks5h", "socks4", "socks4a"}
-RESTResponseType = urllib3.HTTPResponse
-
-
-def is_socks_proxy_url(url):
-    if url is None:
-        return False
-    split_section = url.split("://")
-    if len(split_section) < 2:
-        return False
-    else:
-        return split_section[0].lower() in SUPPORTED_SOCKS_PROXIES
-
 
 class RESTResponse(io.IOBase):
 
-    def __init__(self, resp) -> None:
+    def __init__(self, resp: requests.Response) -> None:
         self.response = resp
-        self.status = resp.status
+        self.status = resp.status_code
         self.reason = resp.reason
         self.data = None
 
     def read(self):
         if self.data is None:
-            self.data = self.response.data
+            self.data = self.response.content
         return self.data
 
-    def getheaders(self):
+    def getheaders(self) -> Dict[str, str]:
         """Returns a dictionary of the response headers."""
-        return self.response.headers
+        return dict(self.response.headers)
 
-    def getheader(self, name, default=None):
+    def getheader(self, name, default=None) -> Optional[str]:
         """Returns a given response header."""
         return self.response.headers.get(name, default)
 
 
-class RESTClientObject:
-
-    def __init__(self, configuration) -> None:
-        # urllib3.PoolManager will pass all kw parameters to connectionpool
-        # https://github.com/shazow/urllib3/blob/f9409436f83aeb79fbaf090181cd81b784f1b8ce/urllib3/poolmanager.py#L75  # noqa: E501
-        # https://github.com/shazow/urllib3/blob/f9409436f83aeb79fbaf090181cd81b784f1b8ce/urllib3/connectionpool.py#L680  # noqa: E501
-        # Custom SSL certificates and client certificates: http://urllib3.readthedocs.io/en/latest/advanced-usage.html  # noqa: E501
-
-        # cert_reqs
-        if configuration.verify_ssl:
-            cert_reqs = ssl.CERT_REQUIRED
-        else:
-            cert_reqs = ssl.CERT_NONE
-
-        pool_args = {
-            "cert_reqs": cert_reqs,
-            "ca_certs": configuration.ssl_ca_cert,
-            "cert_file": configuration.cert_file,
-            "key_file": configuration.key_file,
-        }
-        if configuration.assert_hostname is not None:
-            pool_args['assert_hostname'] = (
-                configuration.assert_hostname
-            )
+RESTResponseType = RESTResponse
 
-        if configuration.retries is not None:
-            pool_args['retries'] = configuration.retries
 
-        if configuration.tls_server_name:
-            pool_args['server_hostname'] = configuration.tls_server_name
-
-
-        if configuration.socket_options is not None:
-            pool_args['socket_options'] = configuration.socket_options
+class RESTClientObject:
 
-        if configuration.connection_pool_maxsize is not None:
-            pool_args['maxsize'] = configuration.connection_pool_maxsize
+    def __init__(self, configuration) -> None:
+        self.configuration = configuration
 
-        # https pool manager
-        self.pool_manager: urllib3.PoolManager
+        self.session: Optional[OAuth2Session] = None
 
-        if configuration.proxy:
-            if is_socks_proxy_url(configuration.proxy):
-                from urllib3.contrib.socks import SOCKSProxyManager
-                pool_args["proxy_url"] = configuration.proxy
-                pool_args["headers"] = configuration.proxy_headers
-                self.pool_manager = SOCKSProxyManager(**pool_args)
-            else:
-                pool_args["proxy_url"] = configuration.proxy
-                pool_args["proxy_headers"] = configuration.proxy_headers
-                self.pool_manager = urllib3.ProxyManager(**pool_args)
-        else:
-            self.pool_manager = urllib3.PoolManager(**pool_args)
+    def setup_oauth_client_credentials(
+            self, client_id: str, client_secret: str, token_endpoint: str, scope: str
+    ):
+        self.session = OAuth2Session(
+            client_id=client_id,
+            client_secret=client_secret,
+            token_endpoint=token_endpoint,
+            scope=scope,
+        )
+        self._configure_session_defaults()
+        self.session.fetch_token()
+
+    def setup_oauth_with_token(self,
+                               client_id: str | None,
+                               client_secret: str | None,
+                               token: dict,
+                               token_endpoint: str,
+                               leeway: int = 60):
+        self.session = OAuth2Session(
+            token=token,
+            token_endpoint=token_endpoint,
+            client_id=client_id,
+            client_secret=client_secret,
+            leeway=leeway,
+        )
+        self._configure_session_defaults()
 
     def request(
-        self,
-        method,
-        url,
-        headers=None,
-        body=None,
-        post_params=None,
-        _request_timeout=None
+            self,
+            method,
+            url,
+            headers=None,
+            body=None,
+            post_params=None,
+            _request_timeout=None,
     ):
         """Perform requests.
 
@@ -135,15 +107,7 @@ class RESTClientObject:
                                  (connection, read) timeouts.
         """
         method = method.upper()
-        assert method in [
-            'GET',
-            'HEAD',
-            'DELETE',
-            'POST',
-            'PUT',
-            'PATCH',
-            'OPTIONS'
-        ]
+        assert method in ["GET", "HEAD", "DELETE", "POST", "PUT", "PATCH", "OPTIONS"]
 
         if post_params and body:
             raise ApiValueError(
@@ -153,105 +117,115 @@ class RESTClientObject:
         post_params = post_params or {}
         headers = headers or {}
 
+        if not self.session:
+            raise ApiValueError(
+                "OAuth2 session is not initialized. Please initialize it before making requests."
+            )
+
+        session = self.session
+
         timeout = None
         if _request_timeout:
             if isinstance(_request_timeout, (int, float)):
-                timeout = urllib3.Timeout(total=_request_timeout)
-            elif (
-                    isinstance(_request_timeout, tuple)
-                    and len(_request_timeout) == 2
-                ):
-                timeout = urllib3.Timeout(
-                    connect=_request_timeout[0],
-                    read=_request_timeout[1]
-                )
+                timeout = _request_timeout
+            elif isinstance(_request_timeout, tuple) and len(_request_timeout) == 2:
+                timeout = _request_timeout
 
         try:
-            # For `POST`, `PUT`, `PATCH`, `OPTIONS`, `DELETE`
-            if method in ['POST', 'PUT', 'PATCH', 'OPTIONS', 'DELETE']:
-
-                # no content type provided or payload is json
-                content_type = headers.get('Content-Type')
-                if (
-                    not content_type
-                    or re.search('json', content_type, re.IGNORECASE)
-                ):
+            if method in ["POST", "PUT", "PATCH", "OPTIONS", "DELETE"]:
+                content_type = headers.get("Content-Type")
+
+                if not content_type or re.search("json", content_type, re.IGNORECASE):
                     request_body = None
                     if body is not None:
                         request_body = json.dumps(body)
-                    r = self.pool_manager.request(
-                        method,
-                        url,
-                        body=request_body,
-                        timeout=timeout,
-                        headers=headers,
-                        preload_content=False
-                    )
+                    r = session.request(method, url, data=request_body, timeout=timeout, headers=headers)
+
                 elif content_type == 'application/x-www-form-urlencoded':
-                    r = self.pool_manager.request(
-                        method,
-                        url,
-                        fields=post_params,
-                        encode_multipart=False,
-                        timeout=timeout,
-                        headers=headers,
-                        preload_content=False
-                    )
+                    r = session.request(method, url, data=post_params, timeout=timeout, headers=headers)
+
                 elif content_type == 'multipart/form-data':
-                    # must del headers['Content-Type'], or the correct
-                    # Content-Type which generated by urllib3 will be
-                    # overwritten.
                     del headers['Content-Type']
-                    # Ensures that dict objects are serialized
-                    post_params = [(a, json.dumps(b)) if isinstance(b, dict) else (a,b) for a, b in post_params]
-                    r = self.pool_manager.request(
-                        method,
-                        url,
-                        fields=post_params,
-                        encode_multipart=True,
-                        timeout=timeout,
-                        headers=headers,
-                        preload_content=False
-                    )
-                # Pass a `string` parameter directly in the body to support
-                # other content types than JSON when `body` argument is
-                # provided in serialized form.
+                    files = []
+                    data = {}
+
+                    for key, value in post_params:
+                        if isinstance(value, tuple) and len(value) >= 2:
+                            # This is a file tuple (filename, file_data, [content_type])
+                            filename, file_data = value[0], value[1]
+                            content_type = value[2] if len(value) > 2 else None
+                            files.append((key, (filename, file_data, content_type)))
+                        elif isinstance(value, dict):
+                            # JSON-serialize dictionary values
+                            if key in data:
+                                # If we already have this key, handle as needed
+                                # (convert to list or append to existing list)
+                                if not isinstance(data[key], list):
+                                    data[key] = [data[key]]
+                                data[key].append(json.dumps(value))
+                            else:
+                                data[key] = json.dumps(value)
+                        else:
+                            # Regular form data
+                            if key in data:
+                                if not isinstance(data[key], list):
+                                    data[key] = [data[key]]
+                                data[key].append(value)
+                            else:
+                                data[key] = value
+                    r = session.request(method, url, files=files, data=data, timeout=timeout, headers=headers)
+
                 elif isinstance(body, str) or isinstance(body, bytes):
-                    r = self.pool_manager.request(
-                        method,
-                        url,
-                        body=body,
-                        timeout=timeout,
-                        headers=headers,
-                        preload_content=False
-                    )
+                    r = session.request(method, url, data=body, timeout=timeout, headers=headers)
+
                 elif headers['Content-Type'].startswith('text/') and isinstance(body, bool):
-                    request_body = "true" if body else "false"
-                    r = self.pool_manager.request(
-                        method,
-                        url,
-                        body=request_body,
-                        preload_content=False,
-                        timeout=timeout,
-                        headers=headers)
+                    request_body = 'true' if body else 'false'
+                    r = session.request(method, url, data=request_body, timeout=timeout, headers=headers)
+
                 else:
-                    # Cannot generate the request from given parameters
-                    msg = """Cannot prepare a request message for provided
-                             arguments. Please check that your arguments match
-                             declared content type."""
+                    msg = '''Cannot prepare a request message for provided arguments.
+                             Please check that your arguments match declared content type.'''
                     raise ApiException(status=0, reason=msg)
-            # For `GET`, `HEAD`
+
             else:
-                r = self.pool_manager.request(
-                    method,
-                    url,
-                    fields={},
-                    timeout=timeout,
-                    headers=headers,
-                    preload_content=False
-                )
-        except urllib3.exceptions.SSLError as e:
-            msg = "\n".join([type(e).__name__, str(e)])
+                r = session.request(method, url, params={}, timeout=timeout, headers=headers)
+
+        except requests.exceptions.SSLError as e:
+            msg = '\n'.join([type(e).__name__, str(e)])
             raise ApiException(status=0, reason=msg)
 
         return RESTResponse(r)
+
+    def _configure_session_defaults(self):
+        self.session.verify = (
+            self.configuration.ssl_ca_cert
+            if self.configuration.ssl_ca_cert
+            else self.configuration.verify_ssl
+        )
+
+        if self.configuration.cert_file and self.configuration.key_file:
+            self.session.cert = (
+                self.configuration.cert_file,
+                self.configuration.key_file,
+            )
+
+        if self.configuration.retries is not None:
+            retry = Retry(
+                total=self.configuration.retries,
+                backoff_factor=0.3,
+                status_forcelist=[429, 500, 502, 503, 504],
+            )
+            adapter = HTTPAdapter(max_retries=retry)
+            self.session.mount("https://", adapter)
+            # noinspection HttpUrlsUsage
+            self.session.mount("http://", adapter)
+
+        if self.configuration.proxy:
+            self.session.proxies = {
+                "http": self.configuration.proxy,
+                "https": self.configuration.proxy,
+            }
+
+            if self.configuration.proxy_headers:
+                for key, value in self.configuration.proxy_headers.items():
+                    self.session.headers[key] = value

rapidata/rapidata_client/order/rapidata_order.py

@@ -51,7 +51,7 @@ class RapidataOrder:
         """Runs the order to start collecting responses."""
         self.__openapi_service.order_api.order_submit_post(self.order_id)
         if print_link:
-            print(f"Order '{self.name}' is now viewable under: https://app.{self.__openapi_service.enviroment}/order/detail/{self.order_id}")
+            print(f"Order '{self.name}' is now viewable under: https://app.{self.__openapi_service.environment}/order/detail/{self.order_id}")
         return self
 
     def pause(self) -> None:
@@ -167,7 +167,7 @@ class RapidataOrder:
             Exception: If the order is not in processing state.
         """        
         campaign_id = self.__get_campaign_id()
-        auth_url = f"https://app.{self.__openapi_service.enviroment}/order/detail/{self.order_id}/preview?campaignId={campaign_id}"
+        auth_url = f"https://app.{self.__openapi_service.environment}/order/detail/{self.order_id}/preview?campaignId={campaign_id}"
         could_open_browser = webbrowser.open(auth_url)
         if not could_open_browser:
             encoded_url = urllib.parse.quote(auth_url, safe="%/:=&?~#+!$,;'@()*[]")

rapidata/rapidata_client/rapidata_client.py

@@ -2,28 +2,37 @@ from rapidata.service.openapi_service import OpenAPIService
 
 from rapidata.rapidata_client.order.rapidata_order_manager import RapidataOrderManager
 
-from rapidata.rapidata_client.validation.validation_set_manager import ValidationSetManager
+from rapidata.rapidata_client.validation.validation_set_manager import (
+    ValidationSetManager,
+)
 
 from rapidata.rapidata_client.demographic.demographic_manager import DemographicManager
 
+
 class RapidataClient:
     """The Rapidata client is the main entry point for interacting with the Rapidata API. It allows you to create orders and validation sets."""
-    
+
     def __init__(
-        self,
-        client_id: str | None = None,
-        client_secret: str | None = None,
-        enviroment: str = "rapidata.ai",
-        oauth_scope: str = "openid",
-        cert_path: str | None = None,
+            self,
+            client_id: str | None = None,
+            client_secret: str | None = None,
+            environment: str = "rapidata.ai",
+            oauth_scope: str = "openid",
+            cert_path: str | None = None,
+            token: dict | None = None,
+            leeway: int = 60
     ):
-        """Initialize the RapidataClient. If both the client_id and client_secret are None, it will try using your credentials under "~/.config/rapidata/credentials.json". 
-        If this is not successful, it will open a browser windown and ask you to log in, then save your new credentials in said json file.
+        """Initialize the RapidataClient. If both the client_id and client_secret are None, it will try using your credentials under "~/.config/rapidata/credentials.json".
+        If this is not successful, it will open a browser window and ask you to log in, then save your new credentials in said json file.
 
         Args:
             client_id (str): The client ID for authentication.
             client_secret (str): The client secret for authentication.
-            enviroment (str, optional): The API endpoint.
+            environment (str, optional): The API endpoint.
+            oauth_scope (str, optional): The scopes to use for authentication. In general this does not need to be changed.
+            cert_path (str, optional): An optional path to a certificate file useful for development.
+            token (dict, optional): If you already have a token that the client should use for authentication. Important, if set, this needs to be the complete token object containing the access token, token type and expiration time.
+            leeway (int, optional): An optional leeway to use to determine if a token is expired. Defaults to 60 seconds.
 
         Attributes:
             order (RapidataOrderManager): The RapidataOrderManager instance.
@@ -32,17 +41,19 @@ class RapidataClient:
         self._openapi_service = OpenAPIService(
             client_id=client_id,
             client_secret=client_secret,
-            enviroment=enviroment,
+            environment=environment,
             oauth_scope=oauth_scope,
-            cert_path=cert_path
+            cert_path=cert_path,
+            token=token,
+            leeway=leeway,
         )
-        
+
         self.order = RapidataOrderManager(openapi_service=self._openapi_service)
-        
+
         self.validation = ValidationSetManager(openapi_service=self._openapi_service)
 
         self._demographic = DemographicManager(openapi_service=self._openapi_service)
 
     def reset_credentials(self):
-        """Reset the credentials saved in the configuration file for the current enviroment."""
+        """Reset the credentials saved in the configuration file for the current environment."""
         self._openapi_service.reset_credentials()

rapidata/rapidata_client/selection/__init__.py

@@ -5,3 +5,5 @@ from .validation_selection import ValidationSelection
 from .conditional_validation_selection import ConditionalValidationSelection
 from .capped_selection import CappedSelection
 from .shuffling_selection import ShufflingSelection
+from .ab_test_selection import AbTestSelection
+from .static_selection import StaticSelection

rapidata/rapidata_client/selection/ab_test_selection.py

@@ -20,7 +20,7 @@ class AbTestSelection(RapidataSelection):
         b_selections (Sequence[RapidataSelection]): List of selections for group B.
     """
 
-    def __init__(self, a_selections: Sequence[RapidataSelection], b_selections: Sequence[RapidataSelection], max_rapids: int):
+    def __init__(self, a_selections: Sequence[RapidataSelection], b_selections: Sequence[RapidataSelection]):
         self.a_selections = a_selections
         self.b_selections = b_selections
 

rapidata/rapidata_client/selection/static_selection.py

@@ -0,0 +1,22 @@
+
+from rapidata.api_client.models.static_selection import StaticSelection as StaticSelectionModel
+from rapidata.rapidata_client.selection._base_selection import RapidataSelection
+
+class StaticSelection(RapidataSelection):
+    """StaticSelection Class
+
+    Given a list of RapidIds, theses specific rapids will be shown in order for every session.
+    
+    Args:
+        selections (list[str]): List of rapid ids to show.
+    """
+
+    def __init__(self, rapid_ids: list[str]):
+        self.rapid_ids = rapid_ids
+
+    def _to_model(self) -> StaticSelectionModel:
+        return StaticSelectionModel(
+            _t="StaticSelection",
+            rapidIds=self.rapid_ids
+        )
+    

rapidata/rapidata_client/validation/rapidata_validation_set.py

@@ -31,13 +31,13 @@ class RapidataValidationSet:
         rapid._add_to_validation_set(self.id, self.__openapi_service, self.__session)
         return self
 
-    def update_dimensions(self, dimensions: list[str] | None):
+    def update_dimensions(self, dimensions: list[str]):
         """Update the dimensions of the validation set.
 
         Args:
             dimensions (list[str]): The new dimensions of the validation set.
         """
-        self.__openapi_service.validation_api.validation_validation_set_id_dimensions_patch(self.id, UpdateDimensionsModel(dimensions=dimensions) if dimensions else None)
+        self.__openapi_service.validation_api.validation_validation_set_id_dimensions_patch(self.id, UpdateDimensionsModel(dimensions=dimensions))
         return self
 
     def _get_session(self, max_retries: int = 5, max_workers: int = 10) -> requests.Session:   

rapidata/rapidata_client/validation/validation_set_manager.py

@@ -39,7 +39,7 @@ class ValidationSetManager:
         data_type: str = RapidataDataTypes.MEDIA,
         contexts: list[str] | None = None,
         explanations: list[str | None] | None = None,
-        dimensions: list[str] | None = None,
+        dimensions: list[str] = [],
         print_confirmation: bool = True,
     ) -> RapidataValidationSet:
         """Create a classification validation set.
@@ -59,7 +59,7 @@ class ValidationSetManager:
                 If provided has to be the same length as datapoints and will be shown in addition to the instruction and answer options. (Therefore will be different for each datapoint)
                 Will be match up with the datapoints using the list index.
             explanations (list[str | None], optional): The explanations for each datapoint. Will be given to the annotators in case the answer is wrong. Defaults to None.
-            dimensions (list[str] | None, optional): The dimensions of the validation set. If not provided will be set to the default dimensions.
+            dimensions (list[str], optional): The dimensions to add to the validation set accross which users will be tracked. Defaults to [] which is the default dimension.
             print_confirmation (bool, optional): Whether to print a confirmation message that validation set has been created. Defaults to True.
 
         Example:
@@ -107,7 +107,7 @@ class ValidationSetManager:
         data_type: str = RapidataDataTypes.MEDIA,
         contexts: list[str] | None = None,
         explanation: list[str | None] | None = None,
-        dimensions: list[str] | None = None,
+        dimensions: list[str] = [],
         print_confirmation: bool = True,
     ) -> RapidataValidationSet:
         """Create a comparison validation set.
@@ -127,7 +127,7 @@ class ValidationSetManager:
                 If provided has to be the same length as datapoints and will be shown in addition to the instruction and truth. (Therefore will be different for each datapoint)
                 Will be match up with the datapoints using the list index.
             explanation (list[str | None], optional): The explanations for each datapoint. Will be given to the annotators in case the answer is wrong. Defaults to None.
-            dimensions (list[str] | None, optional): The dimensions of the validation set. If not provided will be set to the default dimensions.
+            dimensions (list[str], optional): The dimensions to add to the validation set accross which users will be tracked. Defaults to [] which is the default dimension.
             print_confirmation (bool, optional): Whether to print a confirmation message that validation set has been created. Defaults to True.
 
         Example:
@@ -175,7 +175,7 @@ class ValidationSetManager:
         required_precision: float = 1.0,
         required_completeness: float = 1.0,
         explanation: list[str | None] | None = None,
-        dimensions: list[str] | None = None,
+        dimensions: list[str] = [],
         print_confirmation: bool = True,
     ) -> RapidataValidationSet:
         """Create a select words validation set.
@@ -194,7 +194,7 @@ class ValidationSetManager:
             required_precision (float, optional): The required precision for the labeler to get the rapid correct (minimum ratio of the words selected that need to be correct). Defaults to 1.0 (no wrong word can be selected).
             required_completeness (float, optional): The required completeness for the labeler to get the rapid correct (miminum ratio of total correct words selected). Defaults to 1.0 (all correct words need to be selected).
             explanation (list[str | None], optional): The explanations for each datapoint. Will be given to the annotators in case the answer is wrong. Defaults to None.
-            dimensions (list[str] | None, optional): The dimensions of the validation set. If not provided will be set to the default dimensions.
+            dimensions (list[str], optional): The dimensions to add to the validation set accross which users will be tracked. Defaults to [] which is the default dimension.
             print_confirmation (bool, optional): Whether to print a confirmation message that validation set has been created. Defaults to True.
 
         Example:
@@ -238,7 +238,7 @@ class ValidationSetManager:
         datapoints: list[str],
         contexts: list[str] | None = None,
         explanation: list[str | None] | None = None,
-        dimensions: list[str] | None = None,
+        dimensions: list[str] = [],
         print_confirmation: bool = True,
     ) -> RapidataValidationSet:
         """Create a locate validation set.
@@ -253,7 +253,7 @@ class ValidationSetManager:
             datapoints (list[str]): The datapoints that will be used for validation.
             contexts (list[str], optional): The contexts for each datapoint. Defaults to None.
             explanation (list[str | None], optional): The explanations for each datapoint. Will be given to the annotators in case the answer is wrong. Defaults to None.
-            dimensions (list[str] | None, optional): The dimensions of the validation set. If not provided will be set to the default dimensions.
+            dimensions (list[str], optional): The dimensions to add to the validation set accross which users will be tracked. Defaults to [] which is the default dimension.
             print_confirmation (bool, optional): Whether to print a confirmation message that validation set has been created. Defaults to True.
 
         Example:
@@ -299,7 +299,7 @@ class ValidationSetManager:
         datapoints: list[str],
         contexts: list[str] | None = None,
         explanation: list[str | None] | None = None,
-        dimensions: list[str] | None = None,
+        dimensions: list[str] = [],
         print_confirmation: bool = True,
     ) -> RapidataValidationSet:
         """Create a draw validation set.
@@ -314,7 +314,7 @@ class ValidationSetManager:
             datapoints (list[str]): The datapoints that will be used for validation.
             contexts (list[str], optional): The contexts for each datapoint. Defaults to None.
             explanation (list[str | None], optional): The explanations for each datapoint. Will be given to the annotators in case the answer is wrong. Defaults to None.
-            dimensions (list[str] | None, optional): The dimensions of the validation set. If not provided will be set to the default dimensions.
+            dimensions (list[str], optional): The dimensions to add to the validation set accross which users will be tracked. Defaults to [] which is the default dimension.
             print_confirmation (bool, optional): Whether to print a confirmation message that validation set has been created. Defaults to True.
 
         Example:
@@ -359,7 +359,7 @@ class ValidationSetManager:
         datapoints: list[str],
         contexts: list[str] | None = None,
         explanation: list[str | None] | None = None,
-        dimensions: list[str] | None = None,
+        dimensions: list[str] = [],
         print_confirmation: bool = True,
     ) -> RapidataValidationSet:
         """Create a timestamp validation set.
@@ -375,7 +375,7 @@ class ValidationSetManager:
             datapoints (list[str]): The datapoints that will be used for validation.
             contexts (list[str], optional): The contexts for each datapoint. Defaults to None.
             explanation (list[str | None], optional): The explanations for each datapoint. Will be given to the annotators in case the answer is wrong. Defaults to None.
-            dimensions (list[str] | None, optional): The dimensions of the validation set. If not provided will be set to the default dimensions.
+            dimensions (list[str], optional): The dimensions to add to the validation set accross which users will be tracked. Defaults to [] which is the default dimension.
             print_confirmation (bool, optional): Whether to print a confirmation message that validation set has been created. Defaults to True.
 
         Example:
@@ -416,7 +416,7 @@ class ValidationSetManager:
     def create_mixed_set(self,
         name: str,
         rapids: list[Rapid],
-        dimensions: list[str] | None = None,
+        dimensions: list[str] = [],
         print_confirmation: bool = True
     ) -> RapidataValidationSet:
         """Create a validation set with a list of rapids.
@@ -424,7 +424,7 @@ class ValidationSetManager:
         Args:
             name (str): The name of the validation set. (will not be shown to the labeler)
             rapids (list[Rapid]): The list of rapids to add to the validation set.
-            dimensions (list[str] | None, optional): The dimensions of the validation set. If not provided will be set to the default dimensions.
+            dimensions (list[str], optional): The dimensions to add to the validation set accross which users will be tracked. Defaults to [] which is the default dimension.
             print_confirmation (bool, optional): Whether to print a confirmation message that validation set has been created. Defaults to True.
         """
 
@@ -468,10 +468,12 @@ class ValidationSetManager:
         if print_confirmation:
             print()
             print(f"Validation set '{name}' created with ID {validation_set_id}\n",
-                  f"Now viewable under: https://app.{self.__openapi_service.enviroment}/validation-set/detail/{validation_set_id}",
+                  f"Now viewable under: https://app.{self.__openapi_service.environment}/validation-set/detail/{validation_set_id}",
                   sep="")
         
-        validation_set.update_dimensions(dimensions)
+        if dimensions:
+            validation_set.update_dimensions(dimensions)
+        
         return validation_set
 
 

rapidata/service/credential_manager.py

@@ -133,9 +133,9 @@ class CredentialManager:
                 return credential
 
         return self._create_new_credentials()
-    
+
     def reset_credentials(self) -> None:
-        """Reset the stored credentials for current enviroment."""
+        """Reset the stored credentials for current environment."""
         credentials = self._read_credentials()
         if self.endpoint in credentials:
             del credentials[self.endpoint]

rapidata/service/openapi_service.py

@@ -1,3 +1,6 @@
+import subprocess
+from importlib.metadata import version, PackageNotFoundError
+
 from rapidata.api_client.api.campaign_api import CampaignApi
 from rapidata.api_client.api.dataset_api import DatasetApi
 from rapidata.api_client.api.order_api import OrderApi
@@ -7,31 +10,31 @@ from rapidata.api_client.api.validation_api import ValidationApi
 from rapidata.api_client.api.workflow_api import WorkflowApi
 from rapidata.api_client.api_client import ApiClient
 from rapidata.api_client.configuration import Configuration
-from rapidata.service.token_manager import TokenManager, TokenInfo
 from rapidata.service.credential_manager import CredentialManager
 
-from importlib.metadata import version, PackageNotFoundError
-
 
 class OpenAPIService:
     def __init__(
-        self,
-        client_id: str | None,
-        client_secret: str | None,
-        enviroment: str,
-        oauth_scope: str,
-        cert_path: str | None = None,
+            self,
+            client_id: str | None,
+            client_secret: str | None,
+            environment: str,
+            oauth_scope: str,
+            cert_path: str | None = None,
+            token: dict | None = None,
+            leeway: int = 60,
     ):
-        self.enviroment = enviroment
-        endpoint = f"https://api.{enviroment}"
-        self._token_url = f"https://auth.{enviroment}"
-        token_manager = TokenManager(
-            client_id=client_id,
-            client_secret=client_secret,
-            endpoint=self._token_url,
-            oauth_scope=oauth_scope,
-            cert_path=cert_path,
+        self.environment = environment
+        endpoint = f"https://api.{environment}"
+        auth_endpoint = f"https://auth.{environment}"
+
+        if environment == "rapidata.dev" and not cert_path:
+            cert_path = _get_local_certificate()
+
+        self.credential_manager = CredentialManager(
+            endpoint=auth_endpoint, cert_path=cert_path
         )
+
         client_configuration = Configuration(host=endpoint, ssl_ca_cert=cert_path)
         self.api_client = ApiClient(
             configuration=client_configuration,
@@ -39,16 +42,32 @@ class OpenAPIService:
             header_value=f"RapidataPythonSDK/{self._get_rapidata_package_version()}",
         )
 
-        self.api_client.configuration.api_key["bearer"] = (
-            f"Bearer {token_manager.fetch_token().access_token}"
+        if token:
+            self.api_client.rest_client.setup_oauth_with_token(
+                token=token,
+                token_endpoint=f"{auth_endpoint}/connect/token",
+                client_id=client_id,
+                client_secret=client_secret,
+                leeway=leeway,
+            )
+            return
+
+        if not client_id or not client_secret:
+            credentials = self.credential_manager.get_client_credentials()
+            if not credentials:
+                raise ValueError("Failed to fetch client credentials")
+            client_id = credentials.client_id
+            client_secret = credentials.client_secret
+
+        self.api_client.rest_client.setup_oauth_client_credentials(
+            client_id=client_id,
+            client_secret=client_secret,
+            token_endpoint=f"{auth_endpoint}/connect/token",
+            scope=oauth_scope,
         )
 
-        self._cert_path = cert_path
-
-        token_manager.start_token_refresh(token_callback=self._set_token)
-
     def reset_credentials(self):
-        CredentialManager(endpoint=self._token_url, cert_path=self._cert_path).reset_credentials()
+        self.credential_manager.reset_credentials()
 
     @property
     def order_api(self) -> OrderApi:
@@ -78,9 +97,6 @@ class OpenAPIService:
     def workflow_api(self) -> WorkflowApi:
         return WorkflowApi(self.api_client)
 
-    def _set_token(self, token: TokenInfo):
-        self.api_client.configuration.api_key["bearer"] = f"Bearer {token.access_token}"
-
     def _get_rapidata_package_version(self):
         """
         Returns the version of the currently installed rapidata package.
@@ -93,3 +109,15 @@ class OpenAPIService:
             return version("rapidata")
         except PackageNotFoundError:
             return None
+
+
+def _get_local_certificate() -> str | None:
+    result = subprocess.run(["mkcert", "-CAROOT"], capture_output=True)
+    if result.returncode != 0:
+        return None
+
+    output = result.stdout.decode("utf-8").strip()
+    if not output:
+        return None
+
+    return f"{output}/rootCA.pem"

{rapidata-2.17.0.dist-info → rapidata-2.18.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: rapidata
-Version: 2.17.0
+Version: 2.18.0
 Summary: Rapidata package containing the Rapidata Python Client to interact with the Rapidata Web API in an easy way.
 License: Apache-2.0
 Author: Rapidata AG
@@ -12,6 +12,7 @@ Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
+Requires-Dist: authlib (>=1.5.1,<2.0.0)
 Requires-Dist: colorama (==0.4.6)
 Requires-Dist: deprecated (>=1.2.14,<2.0.0)
 Requires-Dist: pandas (>=2.2.3,<3.0.0)

{rapidata-2.17.0.dist-info → rapidata-2.18.0.dist-info}/RECORD

@@ -403,7 +403,7 @@ rapidata/api_client/models/workflow_labeling_step_model.py,sha256=iXeIb78bdMhGFj
 rapidata/api_client/models/workflow_split_model.py,sha256=zthOSaUl8dbLhLymLK_lrPTBpeV1a4cODLxnHmNCAZw,4474
 rapidata/api_client/models/workflow_split_model_filter_configs_inner.py,sha256=1Fx9uZtztiiAdMXkj7YeCqt7o6VkG9lKf7D7UP_h088,7447
 rapidata/api_client/models/workflow_state.py,sha256=5LAK1se76RCoozeVB6oxMPb8p_5bhLZJqn7q5fFQWis,850
-rapidata/api_client/rest.py,sha256=zmCIFQC2l1t-KZcq-TgEm3vco3y_LK6vRm3Q07K-xRI,9423
+rapidata/api_client/rest.py,sha256=Nnn1XE9sVUprPm_6AsUmetb_bd9dMjynDOob6y8NJNE,8775
 rapidata/api_client_README.md,sha256=97mR2UeWNIhqNxgUAz87zH8kM9RsbrColX_FvIy_rYo,53784
 rapidata/rapidata_client/__init__.py,sha256=yU0cRoX-RmOHQv0Qj3yJpHaDET4DHZWSO6w2cAApQhQ,920
 rapidata/rapidata_client/assets/__init__.py,sha256=hKgrOSn8gJcBSULaf4auYhH1S1N5AfcwIhBSq1BOKwQ,323
@@ -438,23 +438,24 @@ rapidata/rapidata_client/metadata/_select_words_metadata.py,sha256=-MK5yQDi_G3BK
 rapidata/rapidata_client/order/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 rapidata/rapidata_client/order/_rapidata_dataset.py,sha256=KyEDb4NvlY4uzN3GjiX331VSGOq_I6FhvTnyKrpo4Sg,19400
 rapidata/rapidata_client/order/_rapidata_order_builder.py,sha256=De-gNvOnuXHz6QaRdNpr0_3zOkIEKu8hgtZkatzxZQ4,13155
-rapidata/rapidata_client/order/rapidata_order.py,sha256=1uhU19uo2xGypd2aJClLRFLWqon-JQN8RTAkTGIg1iQ,10481
+rapidata/rapidata_client/order/rapidata_order.py,sha256=cz1bXkoS7y_MuJDmjdXao7nbCFd8UF4bb5SHZtOYDjI,10483
 rapidata/rapidata_client/order/rapidata_order_manager.py,sha256=HQzZAnwKnwWl2wLcZPqxQRnth1xWhVoGNkMkzzYbjqA,30615
 rapidata/rapidata_client/order/rapidata_results.py,sha256=0y8EOiqUV7XuwpRJyV53mfo-ddRV1cUPdWexbPEVOHM,8044
-rapidata/rapidata_client/rapidata_client.py,sha256=oZrO8KpXQxK4FXDS20FkgdVbE_rBeTztyNYDKbf4ihM,2137
+rapidata/rapidata_client/rapidata_client.py,sha256=fcWlmmNCZahK40Ox4aY153tEihIpwkUxYTuiypKF2SY,2857
 rapidata/rapidata_client/referee/__init__.py,sha256=q0Hv9nmfEpyChejtyMLT8hWKL0vTTf_UgUXPYNJ-H6M,153
 rapidata/rapidata_client/referee/_base_referee.py,sha256=MdFOhdxt3sRnWXLDKLJZKFdVpjBGn9jypPnWWQ6msQA,496
 rapidata/rapidata_client/referee/_early_stopping_referee.py,sha256=ULbokQZ91wc9D_20qHUhe55D28D9eTY1J1cMp_-oIDc,2088
 rapidata/rapidata_client/referee/_naive_referee.py,sha256=PVR8uy8hfRjr2DBzdOFyvou6S3swNc-4UvgjhO-09TU,1209
-rapidata/rapidata_client/selection/__init__.py,sha256=LbafUzvKgKbykbvHZJ7S9aYU82HQl71Y7jAbj_HTZ8c,382
+rapidata/rapidata_client/selection/__init__.py,sha256=OimK44ig39A3kHCR_JGNO4FiUYJ6JUY0ZT0J8dz32Rs,475
 rapidata/rapidata_client/selection/_base_selection.py,sha256=tInbWOgxT_4CHkr5QHoG55ZcUi1ZmfcEGIwLKKCnN20,147
-rapidata/rapidata_client/selection/ab_test_selection.py,sha256=yNVRpYbRBBde_vg_Pf_xSKMemmJSrbHiL7XXKZW6kWw,1301
+rapidata/rapidata_client/selection/ab_test_selection.py,sha256=fymubkVMawqJmYp9FKzWXTki9tgBgoj3cOP8rG9oOd0,1284
 rapidata/rapidata_client/selection/capped_selection.py,sha256=iWhbM1LcayhgFm7oKADXCaKHGdiQIupI0jbYuuEVM2A,1184
 rapidata/rapidata_client/selection/conditional_validation_selection.py,sha256=4etkO5p-wBoI8Wh8vBhNrXm7a_ioFvVmCANJmP8kIwI,2561
 rapidata/rapidata_client/selection/demographic_selection.py,sha256=l4vnNbzlf9ED6BKqN4k5cZXShkXu9L1C5DtO78Vwr5M,1454
 rapidata/rapidata_client/selection/labeling_selection.py,sha256=v26QogjmraFfRoSIgWZl6NMIW_TqbGeuCI2p4HxCeOM,657
 rapidata/rapidata_client/selection/rapidata_selections.py,sha256=Azh0ntBZp9EQNL19imIItotQ8QW3B1gEs5YmuTvUn6U,1526
 rapidata/rapidata_client/selection/shuffling_selection.py,sha256=FzOp7mnBLxNzM5at_-935wd77IHyWnFR1f8uqokiMOg,1201
+rapidata/rapidata_client/selection/static_selection.py,sha256=POhVLjzHcUIuU_GCvRxMuCb27m7CkLxaQPwgf20Xo9o,681
 rapidata/rapidata_client/selection/validation_selection.py,sha256=sedeIa8lpXVXKtFJA9IDeRvo9A1Ne4ZGcepaWDUGhCU,851
 rapidata/rapidata_client/settings/__init__.py,sha256=DTEAT_YykwodZJXqKYOtWRwimLCA-Jxo0F0d-H6A3vM,458
 rapidata/rapidata_client/settings/_rapidata_setting.py,sha256=MD5JhhogSLLrjFKjvL3JhMszOMCygyqLF-st0EwMSkw,352
@@ -468,12 +469,12 @@ rapidata/rapidata_client/settings/play_video_until_the_end.py,sha256=LLHx2_72k5Z
 rapidata/rapidata_client/settings/rapidata_settings.py,sha256=r6eDGo5YHMekOtWqPHD50uI8vEE9VoBJfaWEDFZ78RU,1430
 rapidata/rapidata_client/settings/translation_behaviour.py,sha256=i9n_H0eKJyKW6m3MKH_Cm1XEKWVEWsAV_79xGmGIC-4,742
 rapidata/rapidata_client/validation/__init__.py,sha256=s5wHVtcJkncXSFuL9I0zNwccNOKpWAqxqUjkeohzi2E,24
-rapidata/rapidata_client/validation/rapidata_validation_set.py,sha256=Px_tpFOc5rSvHaDRtN3prNTuXUqz7Y_ZjmbEFLzDciY,2680
+rapidata/rapidata_client/validation/rapidata_validation_set.py,sha256=GaatGGuJCHRvdPbjzI0NRQhxzRftKzf-FkaFv25iB78,2649
 rapidata/rapidata_client/validation/rapids/__init__.py,sha256=WU5PPwtTJlte6U90MDakzx4I8Y0laj7siw9teeXj5R0,21
 rapidata/rapidata_client/validation/rapids/box.py,sha256=t3_Kn6doKXdnJdtbwefXnYKPiTKHneJl9E2inkDSqL8,589
 rapidata/rapidata_client/validation/rapids/rapids.py,sha256=uCKnoSn1RykNHgTFbrvCFlfzU8lF42cff-2I-Pd48w0,4620
 rapidata/rapidata_client/validation/rapids/rapids_manager.py,sha256=F00lPYBUx5fTPRw50iZuobtdbjFo6ZHevPMk101JdaY,14271
-rapidata/rapidata_client/validation/validation_set_manager.py,sha256=OeiJ_WX8qSGdQAOVejnSYxZJUmttNsW71E0MLWUQHXY,26529
+rapidata/rapidata_client/validation/validation_set_manager.py,sha256=Ecm17xf4SFsKdWAGNPdjecuAqsEJ27cHPdns7f8GXhs,26699
 rapidata/rapidata_client/workflow/__init__.py,sha256=7nXcY91xkxjHudBc9H0fP35eBBtgwHGWTQKbb-M4h7Y,477
 rapidata/rapidata_client/workflow/_base_workflow.py,sha256=XyIZFKS_RxAuwIHS848S3AyLEHqd07oTD_5jm2oUbsw,762
 rapidata/rapidata_client/workflow/_classify_workflow.py,sha256=9bT54wxVJgxC-zLk6MVNbseFpzYrvFPjt7DHvxqYfnk,1736
@@ -486,11 +487,10 @@ rapidata/rapidata_client/workflow/_ranking_workflow.py,sha256=XBIifokhu9Kaqo0qSw
 rapidata/rapidata_client/workflow/_select_words_workflow.py,sha256=juTW4TPnnSeBWP3K2QfO092t0u1W8I1ksY6aAtPZOi0,1225
 rapidata/rapidata_client/workflow/_timestamp_workflow.py,sha256=tPi2zu1-SlE_ppbGbMz6MM_2LUSWxM-GA0CZRlB0qFo,1176
 rapidata/service/__init__.py,sha256=s9bS1AJZaWIhLtJX_ZA40_CK39rAAkwdAmymTMbeWl4,68
-rapidata/service/credential_manager.py,sha256=_DIP665fpl4fkqj1l-wjRrBp-8fy2Db7tnXDDx4kWyc,8605
+rapidata/service/credential_manager.py,sha256=3x-Fb6tyqmgtpjI1MSOtXWW_SkzTK8Lo7I0SSL2YD7E,8602
 rapidata/service/local_file_service.py,sha256=pgorvlWcx52Uh3cEG6VrdMK_t__7dacQ_5AnfY14BW8,877
-rapidata/service/openapi_service.py,sha256=fUQGLQzezjJbLqHVq7o6pQZUrK2Y12sfV23OstV4lOk,3234
-rapidata/service/token_manager.py,sha256=C-8dN6P5TXCLANZCHWusmwAful5YBpKjKg0StQtajF0,6547
-rapidata-2.17.0.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-rapidata-2.17.0.dist-info/METADATA,sha256=5yD0-X5uZr9D4YPNk_KKwhXPOVAkcUpGubS61oz3Yfg,1147
-rapidata-2.17.0.dist-info/WHEEL,sha256=XbeZDeTWKc1w7CSIyre5aMDU_-PohRwTQceYnisIYYY,88
-rapidata-2.17.0.dist-info/RECORD,,
+rapidata/service/openapi_service.py,sha256=ORFPfHlb41zOUP5nDjYWZwO-ZcqNF_Mw2r71RitFtS0,4042
+rapidata-2.18.0.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+rapidata-2.18.0.dist-info/METADATA,sha256=RmwN7x5_3aS68UGwC-BTB3FZg8S879NhC32ymMZ8S0A,1187
+rapidata-2.18.0.dist-info/WHEEL,sha256=XbeZDeTWKc1w7CSIyre5aMDU_-PohRwTQceYnisIYYY,88
+rapidata-2.18.0.dist-info/RECORD,,

rapidata/service/token_manager.py

@@ -1,176 +0,0 @@
-import json
-import logging
-import threading
-from datetime import datetime, timedelta
-from typing import Optional, Callable
-
-import requests
-from pydantic import BaseModel
-
-from rapidata.service.credential_manager import CredentialManager
-
-logger = logging.getLogger(__name__)
-
-class TokenInfo(BaseModel):
-    access_token: str
-    expires_in: int
-    issued_at: datetime
-    token_type: str = "Bearer"
-
-    @property
-    def auth_header(self):
-        return f"{self.token_type} {self.access_token}"
-
-    @property
-    def time_remaining(self):
-        remaining = (
-            (self.issued_at + timedelta(seconds=self.expires_in)) - datetime.now()
-        ).total_seconds()
-        return max(0.0, remaining)
-
-
-class TokenManager:
-    def __init__(
-        self,
-        client_id: str | None = None,
-        client_secret: str | None = None,
-        endpoint: str = "https://auth.rapidata.ai",
-        oauth_scope: str = "openid profile email",
-        cert_path: str | None = None,
-        refresh_threshold: float = 0.8,
-        max_sleep_time: float = 30,
-        max_token_lifetime: int = 60,
-    ):
-        self._client_id = client_id
-        self._client_secret = client_secret
-
-        if not client_id or not client_secret:
-            credential_manager = CredentialManager(
-                endpoint=endpoint, cert_path=cert_path
-            )
-            credentials = credential_manager.get_client_credentials()
-            if not credentials:
-                raise ValueError("Failed to fetch client credentials")
-            self._client_id = credentials.client_id
-            self._client_secret = credentials.client_secret
-
-        self._endpoint = endpoint
-        self._oauth_scope = oauth_scope
-        self._cert_path = cert_path
-        self._refresh_threshold = refresh_threshold
-        self._max_sleep_time = max_sleep_time
-        self._max_token_lifetime = max_token_lifetime
-
-        self._token_lock = threading.Lock()
-        self._current_token: Optional[TokenInfo] = None
-        self._refresh_thread: Optional[threading.Thread] = None
-        self._should_stop = threading.Event()
-
-    def fetch_token(self):
-        try:
-            response = requests.post(
-                f"{self._endpoint}/connect/token",
-                data={
-                    "grant_type": "client_credentials",
-                    "client_id": self._client_id,
-                    "client_secret": self._client_secret,
-                    "scope": self._oauth_scope,
-                },
-                verify=self._cert_path,
-            )
-
-            if response.ok:
-                data = response.json()
-                return TokenInfo(
-                    access_token=data["access_token"],
-                    token_type=data["token_type"],
-                    expires_in=min(self._max_token_lifetime, data["expires_in"]),
-                    issued_at=datetime.now(),
-                )
-
-            else:
-                data = response.text
-                error_description = "An unknown error occurred"
-                if "error_description" in data:
-                    error_description = (
-                        data.split("error_description")[1].split("\n")[0].strip()
-                    )
-                raise ValueError(f"Failed to fetch token: {error_description}")
-        except requests.RequestException as e:
-            raise ValueError(f"Failed to fetch token: {e}")
-        except json.JSONDecodeError as e:
-            raise ValueError(f"Failed to parse token response: {e}")
-        except KeyError as e:
-            raise ValueError(f"Failed to extract token from response: {e}")
-
-    def start_token_refresh(self, token_callback: Callable[[TokenInfo], None]) -> None:
-        if self._refresh_thread and self._refresh_thread.is_alive():
-            logger.error("Token refresh thread is already running")
-            return
-
-        def refresh_loop():
-            while not self._should_stop.is_set():
-                try:
-                    with self._token_lock:
-                        if self._should_refresh_token(self._current_token):
-                            logger.debug("Refreshing token")
-                            self._current_token = self.fetch_token()
-                            token_callback(self._current_token)
-
-                    if self._current_token:
-                        time_until_refresh_threshold = (
-                            self._current_token.time_remaining
-                            - (
-                                self._current_token.expires_in
-                                * (1 - self._refresh_threshold)
-                            )
-                        )
-                        logger.debug("Time until refresh threshold: %s", time_until_refresh_threshold)
-                        sleep_time = min(
-                            self._max_sleep_time, time_until_refresh_threshold
-                        )
-                        logger.debug(
-                            f"Sleeping for {sleep_time} until checking the token again"
-                        )
-                        self._should_stop.wait(timeout=max(1.0, sleep_time))
-                    else:
-                        self._should_stop.wait(timeout=self._max_sleep_time)
-                except Exception as e:
-                    logger.error("Failed to refresh token: %s", e)
-                    self._should_stop.wait(timeout=5)
-
-        self._should_stop.clear()
-        self._refresh_thread = threading.Thread(target=refresh_loop, daemon=True)
-        self._refresh_thread.start()
-
-    def stop_token_refresh(self):
-        self._should_stop.set()
-        if self._refresh_thread:
-            self._refresh_thread.join(timeout=1)
-            self._refresh_thread = None
-
-    def get_current_token(self) -> Optional[TokenInfo]:
-        with self._token_lock:
-            return self._current_token
-
-    def _should_refresh_token(self, token: TokenInfo | None) -> bool:
-        if not token:
-            return True
-
-        limit = token.expires_in * (1 - self._refresh_threshold)
-
-        logger.debug(
-            "The token was issued at %s, it expires in %s. It has %s seconds remaining and we refresh the token when it has %s seconds remaining",
-            token.issued_at,
-            token.expires_in,
-            token.time_remaining,
-            limit,
-        )
-        return token.time_remaining < limit
-
-    def __enter__(self):
-        return self
-
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        self.stop_token_refresh()
-        return False