rapidata 2.10.1__tar.gz → 2.11.0__tar.gz

{rapidata-2.10.1 → rapidata-2.11.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: rapidata
-Version: 2.10.1
+Version: 2.11.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

{rapidata-2.10.1 → rapidata-2.11.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "rapidata"
-version = "2.10.1"
+version = "2.11.0"
 description = "Rapidata package containing the Rapidata Python Client to interact with the Rapidata Web API in an easy way."
 authors = ["Rapidata AG <info@rapidata.ai>"]
 license = "Apache-2.0"

{rapidata-2.10.1 → rapidata-2.11.0}/src/rapidata/rapidata_client/assets/_media_asset.py

@@ -4,6 +4,7 @@ Defines the MediaAsset class for handling media file paths within assets.
 Implements lazy loading for URL-based media to prevent unnecessary downloads.
 """
 
+from typing import Optional, cast, Sequence
 import os
 from io import BytesIO
 from rapidata.rapidata_client.assets._base_asset import BaseAsset
@@ -13,10 +14,10 @@ from PIL import Image
 from tinytag import TinyTag
 import tempfile
 from pydantic import StrictStr, StrictBytes
-from typing import Optional, cast
 import logging
 from functools import cached_property
-
+from requests.adapters import HTTPAdapter
+from urllib3.util.retry import Retry
 
 class MediaAsset(BaseAsset):
     """MediaAsset Class with Lazy Loading
@@ -77,6 +78,7 @@ class MediaAsset(BaseAsset):
         
         self._url = None
         self._content = None
+        self._session: None | requests.Session  = None
         
         if re.match(r'^https?://', path):
             self._url = path
@@ -223,7 +225,7 @@ class MediaAsset(BaseAsset):
 
     def __get_media_bytes(self, url: str) -> bytes:
         """
-        Downloads and validates media files from URL.
+        Downloads and validates media files from URL with retry logic and session reuse.
         
         Args:
             url: URL of the media file
@@ -233,13 +235,21 @@ class MediaAsset(BaseAsset):
             
         Raises:
             ValueError: If media type is unsupported or content validation fails
-            requests.exceptions.RequestException: If download fails
+            requests.exceptions.RequestException: If download fails after all retries
         """
+        # Use existing session or throw error if not set
+        if self._session is None:
+            raise RuntimeError("HTTP session not configured")
+
         try:
-            response = requests.get(url, stream=False)
+            response = self._session.get(
+                url, 
+                stream=False,
+                timeout=(5, 30)  # (connect timeout, read timeout)
+            )
             response.raise_for_status()
         except requests.exceptions.RequestException as e:
-            self._logger.error(f"Failed to download media from {url}: {str(e)}")
+            self._logger.error(f"Failed to download media from {url} after retries: {str(e)}")
             raise
 
         content = response.content

{rapidata-2.10.1 → rapidata-2.11.0}/src/rapidata/rapidata_client/filter/language_filter.py

@@ -9,11 +9,16 @@ class LanguageFilter(RapidataFilter):
     """LanguageFilter Class
     
     Can be used to filter who to target based on language codes.
-    
-    example: LanguageFilter(["en", "de"]) -> will limit the order to be shown to only people who have their phone set to english or german
 
     Args:
-        language_codes (list[str]): List of language codes to filter by."""
+        language_codes (list[str]): List of language codes to filter by.
+            
+    Example: 
+        ```python
+        LanguageFilter(["en", "de"])
+        ```
+        This will limit the order to be shown to only people who have their phone set to english or german
+    """
     def __init__(self, language_codes: list[str]):
         if not isinstance(language_codes, list):
             raise ValueError("Language codes must be a list")

{rapidata-2.10.1 → rapidata-2.11.0}/src/rapidata/rapidata_client/filter/user_score_filter.py

@@ -10,11 +10,15 @@ class UserScoreFilter(RapidataFilter):
     
     Can be used to filter who to target based on their user score.
     
-    example: UserScoreFilter(0.5, 0.9) -> will only show the order to users that have a UserScore of >=0.5 and <=0.9
-
     Args:
         lower_bound (float): The lower bound of the user score.
         upper_bound (float): The upper bound of the user score.
+
+    Example: 
+        ```python
+        UserScoreFilter(0.5, 0.9)
+        ```
+        This will only show the order to users that have a UserScore of >=0.5 and <=0.9
     """
     def __init__(self, lower_bound: float = 0.0, upper_bound: float = 1.0):
         if lower_bound < 0 or lower_bound > 1:

{rapidata-2.10.1 → rapidata-2.11.0}/src/rapidata/rapidata_client/metadata/_base_metadata.py

@@ -9,5 +9,5 @@ class Metadata:
         self._identifier = identifier
 
     @abstractmethod
-    def _to_model(self) -> Any:
+    def to_model(self) -> Any:
         pass

{rapidata-2.10.1 → rapidata-2.11.0}/src/rapidata/rapidata_client/metadata/_private_text_metadata.py

@@ -10,7 +10,7 @@ class PrivateTextMetadata(Metadata):
         super().__init__(identifier=identifier)
         self._text = text
 
-    def _to_model(self):
+    def to_model(self):
         return PrivateTextMetadataInput(
             _t="PrivateTextMetadataInput", identifier=self._identifier, text=self._text
         )

{rapidata-2.10.1 → rapidata-2.11.0}/src/rapidata/rapidata_client/metadata/_prompt_metadata.py

@@ -14,5 +14,5 @@ class PromptMetadata(Metadata):
         self._prompt = prompt
     
 
-    def _to_model(self):
+    def to_model(self):
         return PromptMetadataInput(_t="PromptMetadataInput", identifier=self._identifier, prompt=self._prompt)

{rapidata-2.10.1 → rapidata-2.11.0}/src/rapidata/rapidata_client/metadata/_public_text_metadata.py

@@ -10,7 +10,7 @@ class PublicTextMetadata(Metadata):
         super().__init__(identifier=identifier)
         self._text = text
 
-    def _to_model(self):
+    def to_model(self):
         return PublicTextMetadataInput(
             _t="PublicTextMetadataInput", identifier=self._identifier, text=self._text
         )

{rapidata-2.10.1 → rapidata-2.11.0}/src/rapidata/rapidata_client/metadata/_select_words_metadata.py

@@ -12,7 +12,7 @@ class SelectWordsMetadata(Metadata):
         self.identifier = identifier
         self.select_words = select_words
 
-    def _to_model(self):
+    def to_model(self):
         return TranscriptionMetadataInput(
             _t="TranscriptionMetadataInput",
             identifier=self.identifier,

rapidata-2.11.0/src/rapidata/rapidata_client/order/_rapidata_dataset.py

@@ -0,0 +1,225 @@
+from itertools import zip_longest
+
+from rapidata.api_client.models.datapoint_metadata_model import DatapointMetadataModel
+from rapidata.api_client.models.datapoint_metadata_model_metadata_inner import (
+    DatapointMetadataModelMetadataInner,
+)
+from rapidata.api_client.models.upload_text_sources_to_dataset_model import (
+    UploadTextSourcesToDatasetModel,
+)
+from rapidata.rapidata_client.metadata._base_metadata import Metadata
+from rapidata.rapidata_client.assets import TextAsset, MediaAsset, MultiAsset
+from rapidata.service import LocalFileService
+from rapidata.service.openapi_service import OpenAPIService
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from tqdm import tqdm
+
+from pydantic import StrictBytes, StrictStr
+from typing import Optional, cast, Sequence, Generator
+from logging import Logger
+from requests.adapters import HTTPAdapter, Retry
+import time
+import requests
+
+
+def chunk_list(lst: list, chunk_size: int) -> Generator:
+    """Split list into chunks to prevent resource exhaustion"""
+    for i in range(0, len(lst), chunk_size):
+        yield lst[i:i + chunk_size]
+
+class RapidataDataset:
+
+    def __init__(self, dataset_id: str, openapi_service: OpenAPIService):
+        self.dataset_id = dataset_id
+        self.openapi_service = openapi_service
+        self.local_file_service = LocalFileService()
+        self._logger = Logger(__name__)
+
+    def _add_texts(
+        self,
+        text_assets: list[TextAsset] | list[MultiAsset],
+        max_workers: int = 10,
+    ):
+        for text_asset in text_assets:
+            if isinstance(text_asset, MultiAsset):
+                assert all(
+                    isinstance(asset, TextAsset) for asset in text_asset.assets
+                ), "All assets in a MultiAsset must be of type TextAsset."
+
+        def upload_text_datapoint(text_asset: TextAsset | MultiAsset, index: int) -> None:
+            if isinstance(text_asset, TextAsset):
+                texts = [text_asset.text]
+            elif isinstance(text_asset, MultiAsset):
+                texts = [asset.text for asset in text_asset.assets if isinstance(asset, TextAsset)]
+            else:
+                raise ValueError(f"Unsupported asset type: {type(text_asset)}")
+
+            model = UploadTextSourcesToDatasetModel(
+                datasetId=self.dataset_id,
+                textSources=texts,
+                sortIndex=index,
+            )
+
+            upload_response = self.openapi_service.dataset_api.dataset_creat_text_datapoint_post(model)
+
+            if upload_response.errors:
+                raise ValueError(f"Error uploading text datapoint: {upload_response.errors}")
+
+        total_uploads = len(text_assets)
+        with ThreadPoolExecutor(max_workers=max_workers) as executor:
+            futures = [
+                executor.submit(upload_text_datapoint, text_asset, index=i)
+                for i, text_asset in enumerate(text_assets)
+            ]
+
+            with tqdm(total=total_uploads, desc="Uploading text datapoints") as pbar:
+                for future in as_completed(futures):
+                    future.result()  # This will raise any exceptions that occurred during execution
+                    pbar.update(1)
+
+    def _add_media_from_paths(
+        self,
+        media_paths: list[MediaAsset] | list[MultiAsset],
+        metadata: Sequence[Metadata] | None = None,
+        max_workers: int = 10,
+        max_retries: int = 5,
+        chunk_size: int = 50,
+    ) -> tuple[list[str], list[str]]:
+        """
+        Upload media paths in chunks with managed resources.
+        
+        Args:
+            media_paths: List of MediaAsset or MultiAsset objects to upload
+            metadata: Optional sequence of metadata matching media_paths length
+            max_workers: Maximum number of concurrent upload workers
+            max_retries: Maximum number of retry attempts per failed request
+            chunk_size: Number of items to process in each batch
+            
+        Returns:
+            tuple[list[str], list[str]]: Lists of successful and failed URLs
+            
+        Raises:
+            ValueError: If metadata length doesn't match media_paths length
+        """
+        if metadata is not None and len(metadata) != len(media_paths):
+            raise ValueError("metadata must be None or have the same length as media_paths")
+
+        # Configure session with retry logic
+        session = requests.Session()
+        retries = Retry(
+            total=max_retries,
+            backoff_factor=1,
+            status_forcelist=[500, 502, 503, 504],
+            allowed_methods=["GET"],
+            respect_retry_after_header=True
+        )
+        
+        adapter = HTTPAdapter(
+            pool_connections=max_workers * 2,
+            pool_maxsize=max_workers * 4,
+            max_retries=retries
+        )
+        session.mount('http://', adapter)
+        session.mount('https://', adapter)
+
+        def upload_datapoint(
+            media_asset: MediaAsset | MultiAsset, 
+            meta: Metadata | None, 
+            index: int,
+            session: requests.Session
+        ) -> tuple[list[str], list[str]]:
+            """Process single upload with error tracking."""
+            local_successful: list[str] = []
+            local_failed: list[str] = []
+            identifiers_to_track: list[str] = []
+            
+            try:
+                # Get identifier for this upload (URL or file path)
+                if isinstance(media_asset, MediaAsset):
+                    media_asset._session = session
+                    assets = [media_asset]
+                    identifier = media_asset._url if media_asset._url else media_asset.path
+                    identifiers_to_track = [identifier] if identifier else []
+                elif isinstance(media_asset, MultiAsset):
+                    assets = cast(list[MediaAsset], media_asset.assets)
+                    for asset in assets:
+                        asset._session = session
+                    identifiers_to_track: list[str] = [
+                        (asset._url if asset._url else cast(str, asset.path)) 
+                        for asset in assets
+                    ]
+                else:
+                    raise ValueError(f"Unsupported asset type: {type(media_asset)}")
+
+                meta_model = meta.to_model() if meta else None
+                model = DatapointMetadataModel(
+                    datasetId=self.dataset_id,
+                    metadata=([DatapointMetadataModelMetadataInner(meta_model)] if meta_model else []),
+                    sortIndex=index,
+                )
+
+                files: list[tuple[StrictStr, StrictBytes] | StrictStr | StrictBytes] = []
+                for asset in assets:
+                    if isinstance(asset, MediaAsset):
+                        files.append(asset.to_file())
+
+                upload_response = self.openapi_service.dataset_api.dataset_create_datapoint_post(
+                    model=model,
+                    files=files
+                )
+
+                if upload_response.errors:
+                    error_msg = f"Error uploading datapoint: {upload_response.errors}"
+                    self._logger.error(error_msg)
+                    local_failed.extend(identifiers_to_track)
+                    raise ValueError(error_msg)
+
+                local_successful.extend(identifiers_to_track)
+
+            except Exception as e:
+                self._logger.error(f"\nUpload failed for {identifiers_to_track}: {str(e)}") # \n to avoid same line as tqdm
+                local_failed.extend(identifiers_to_track)
+
+            return local_successful, local_failed
+
+
+        # Process uploads in chunks
+        successful_uploads: list[str] = []
+        failed_uploads: list[str] = []
+        total_uploads = len(media_paths)
+
+        with tqdm(total=total_uploads, desc="Uploading datapoints") as pbar:
+            for chunk_idx, chunk in enumerate(chunk_list(media_paths, chunk_size)):
+                chunk_metadata = metadata[chunk_idx * chunk_size:(chunk_idx + 1) * chunk_size] if metadata else None
+                
+                with ThreadPoolExecutor(max_workers=max_workers) as executor:
+                    futures = [
+                        executor.submit(
+                            upload_datapoint, 
+                            media_asset, 
+                            meta, 
+                            index=(chunk_idx * chunk_size + i),
+                            session=session
+                        )
+                        for i, (media_asset, meta) in enumerate(zip_longest(chunk, chunk_metadata or []))
+                    ]
+
+                    for future in as_completed(futures):
+                        try:
+                            chunk_successful, chunk_failed = future.result()
+                            successful_uploads.extend(chunk_successful)
+                            failed_uploads.extend(chunk_failed)
+                        except Exception as e:
+                            self._logger.error(f"Future execution failed: {str(e)}")
+                        finally:
+                            pbar.update(1)
+
+        # Log summary statistics
+        success_rate = len(successful_uploads) / total_uploads * 100 if total_uploads > 0 else 0
+        self._logger.info(f"Upload complete: {len(successful_uploads)} successful, {len(failed_uploads)} failed ({success_rate:.1f}% success rate)")
+
+        if failed_uploads:
+            print(f"Failed uploads: {failed_uploads}")
+
+        return successful_uploads, failed_uploads
+

{rapidata-2.10.1 → rapidata-2.11.0}/src/rapidata/rapidata_client/order/rapidata_results.py

@@ -193,7 +193,10 @@ class RapidataResults(dict):
                 key: value for key, value in result.items() 
                 if not isinstance(value, dict)
             }
-            
+
+            row["assetA"] = asset_a
+            row["assetB"] = asset_b
+
             # Handle comparative metrics
             for key, values in result.items():
                 if isinstance(values, dict) and len(values) == 2:

{rapidata-2.10.1 → rapidata-2.11.0}/src/rapidata/rapidata_client/selection/conditional_validation_selection.py

@@ -17,6 +17,19 @@ class ConditionalValidationSelection(RapidataSelection):
         thresholds (list[float]): The thresholds to use for the user score.
         chances (list[float]): The chances of showing a validation rapid for each threshold.
         rapid_counts (list[int]): The amount of validation rapids that will be shown per session of this validation set for each threshold if selected by probability. (all or nothing)
+
+    Example:
+        ```python
+        ConditionalValidationSelection(
+            validation_set_id="validation_set_id",
+            thresholds=[0, 0.7], # (0 must be the first threshold)
+            chances=[1, 0.2],
+            rapid_counts=[1, 1]
+        )
+        ```
+        This means that there's a 100% chance of showing a validation rapid if the user score is between 0 and 0.7, 
+        and a 20% chance of showing a validation rapid if the user score is between 0.7 and 1.
+
     """
 
     def __init__(

{rapidata-2.10.1 → rapidata-2.11.0}/src/rapidata/rapidata_client/selection/demographic_selection.py

@@ -14,7 +14,16 @@ class DemographicSelection(RapidataSelection):
     Args:
         keys (list[str]): List of keys for the demographic rapids to be shown. As an example: "age"
         max_rapids (int): The maximum number of rapids to run.\n
-            Allows to provide more keys, in case some of the earlier ones are not selected because of high confidence."""
+            Allows to provide more keys, in case some of the earlier ones are not selected because of high confidence.
+                 
+    Example:
+        ```python
+        DemographicSelection(["age", "gender"], 1)
+        ```
+        This will try to ask the user about their age, if that is not selected due to an already high confidence, it will try asking about their gender.
+        The gender question may also be skipped if the confidence is high enough.
+    """
+    
 
     def __init__(self, keys: list[str], max_rapids: int):
         self.keys = keys

{rapidata-2.10.1 → rapidata-2.11.0}/src/rapidata/rapidata_client/validation/rapids/rapids.py

@@ -66,7 +66,7 @@ class Rapid():
             payload=AddValidationRapidModelPayload(self.payload),
             truth=AddValidationRapidModelTruth(self.truth),
             metadata=[
-              DatapointMetadataModelMetadataInner(meta._to_model())
+              DatapointMetadataModelMetadataInner(meta.to_model())
               for meta in self.metadata
             ],
             randomCorrectProbability=self.randomCorrectProbability,
@@ -93,7 +93,7 @@ class Rapid():
           payload=AddValidationRapidModelPayload(self.payload),
           truth=AddValidationRapidModelTruth(self.truth),
           metadata=[
-              DatapointMetadataModelMetadataInner(meta._to_model())
+              DatapointMetadataModelMetadataInner(meta.to_model())
               for meta in self.metadata
           ],
           randomCorrectProbability=self.randomCorrectProbability,

{rapidata-2.10.1 → rapidata-2.11.0}/src/rapidata/rapidata_client/validation/rapids/rapids_manager.py

@@ -238,10 +238,10 @@ class RapidsManager:
 
         model_truth = BoundingBoxTruth(
             _t="BoundingBoxTruth", 
-            xMax=truths[0].x_max / img_dimensions[0],
-            xMin=truths[0].x_min / img_dimensions[0],
-            yMax=truths[0].y_max / img_dimensions[1],
-            yMin=truths[0].y_min / img_dimensions[1],
+            xMax=truths[0].x_max / img_dimensions[0] * 100,
+            xMin=truths[0].x_min / img_dimensions[0] * 100,
+            yMax=truths[0].y_max / img_dimensions[1] * 100,
+            yMin=truths[0].y_min / img_dimensions[1] * 100,
         )
 
         coverage = self._calculate_boxes_coverage(truths, img_dimensions[0], img_dimensions[1])

{rapidata-2.10.1 → rapidata-2.11.0}/src/rapidata/rapidata_client/validation/validation_set_manager.py

@@ -58,6 +58,14 @@ class ValidationSetManager:
                 Will be match up with the datapoints using the list index.
             print_confirmation (bool, optional): Whether to print a confirmation message that validation set has been created. Defaults to True.
             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.
+
+        Example:
+            ```python
+            options: ["yes", "no", "maybe"]
+            datapoints: ["datapoint1", "datapoint2"]
+            truths: [["yes"], ["no", "maybe"]]
+            ```
+            This would mean: first datapoint correct answer is "yes", second datapoint is "no" or "maybe"
         """
         
         if len(datapoints) != len(truths):
@@ -116,6 +124,14 @@ class ValidationSetManager:
                 Will be match up with the datapoints using the list index.
             print_confirmation (bool, optional): Whether to print a confirmation message that validation set has been created. Defaults to True.
             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.
+
+        Example:
+            ```python
+            instruction: "Which image has a cat?"
+            datapoints = [["image1.jpg", "image2.jpg"], ["image3.jpg", "image4.jpg"]]
+            truths: ["image1.jpg", "image4.jpg"]
+            ```
+            This would mean: first comparison image1.jpg has a cat, second comparison image4.jpg has a cat
         """
         
         if len(datapoints) != len(truths):
@@ -164,8 +180,8 @@ class ValidationSetManager:
             truths (list[list[int]]): The truths for each datapoint. Outer list is for each datapoint, inner list is for each truth.\n
                 example:
                     datapoints: ["datapoint1", "datapoint2"]
-                    sentences: ["this example 1", "this example 2"]
-                    truths: [[0, 1], [2]] -> first datapoint correct words are "this" and "example", second datapoint is "2"
+                    sentences: ["this example 1", "this example with another text"]
+                    truths: [[0, 1], [2]] -> first datapoint correct words are "this" and "example", second datapoint is "with"
             datapoints (list[str]): The datapoints that will be used for validation.
             sentences (list[str]): The sentences that will be used for validation. The sentece will be split up by spaces to be selected by the labeler.
                 Must be the same length as datapoints.
@@ -173,6 +189,14 @@ class ValidationSetManager:
             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).
             print_confirmation (bool, optional): Whether to print a confirmation message that validation set has been created. Defaults to True.
             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.
+
+        Example:
+            ```python
+            datapoints: ["datapoint1", "datapoint2"]
+            sentences: ["this example 1", "this example with another text"]
+            truths: [[0, 1], [2]]
+            ```
+            This would mean: first datapoint the correct words are "this" and "example", second datapoint is "with"
             """
         
         if not all([isinstance(truth, (list, tuple)) for truth in truths]):
@@ -222,6 +246,13 @@ class ValidationSetManager:
             contexts (list[str], optional): The contexts for each datapoint. Defaults to None.
             print_confirmation (bool, optional): Whether to print a confirmation message that validation set has been created. Defaults to True.
             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.
+
+        Example:
+            ```python
+            datapoints: ["datapoint1", "datapoint2"]
+            truths: [[Box(0, 0, 100, 100)], [Box(50, 50, 150, 150)]]
+            ```
+            This would mean: first datapoint the object is in the top left corner, second datapoint the object is in the center
         """
         
         if len(datapoints) != len(truths):
@@ -274,6 +305,13 @@ class ValidationSetManager:
             contexts (list[str], optional): The contexts for each datapoint. Defaults to None.
             print_confirmation (bool, optional): Whether to print a confirmation message that validation set has been created. Defaults to True.
             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.
+
+        Example:
+            ```python
+            datapoints: ["datapoint1", "datapoint2"]
+            truths: [[Box(0, 0, 100, 100)], [Box(50, 50, 150, 150)]]
+            ```
+            This would mean: first datapoint the object is in the top left corner, second datapoint the object is in the center
         """
         
         if len(datapoints) != len(truths):
@@ -326,6 +364,13 @@ class ValidationSetManager:
             contexts (list[str], optional): The contexts for each datapoint. Defaults to None.
             print_confirmation (bool, optional): Whether to print a confirmation message that validation set has been created. Defaults to True.
             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.
+
+        Example:
+            ```python
+            datapoints: ["datapoint1", "datapoint2"]
+            truths: [[(0, 10)], [(20, 30)]]
+            ```
+            This would mean: first datapoint the correct interval is from 0 to 10, second datapoint the correct interval is from 20 to 30
         """
         
         if len(datapoints) != len(truths):

{rapidata-2.10.1 → rapidata-2.11.0}/src/rapidata/service/token_manager.py

@@ -11,7 +11,6 @@ from rapidata.service.credential_manager import CredentialManager
 
 logger = logging.getLogger(__name__)
 
-
 class TokenInfo(BaseModel):
     access_token: str
     expires_in: int
@@ -40,6 +39,7 @@ class TokenManager:
         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
@@ -59,6 +59,7 @@ class TokenManager:
         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
@@ -83,7 +84,7 @@ class TokenManager:
                 return TokenInfo(
                     access_token=data["access_token"],
                     token_type=data["token_type"],
-                    expires_in=data["expires_in"],
+                    expires_in=min(self._max_token_lifetime, data["expires_in"]),
                     issued_at=datetime.now(),
                 )