rapidata 2.37.0__py3-none-any.whl → 2.39.0__py3-none-any.whl

rapidata/rapidata_client/filter/user_score_filter.py

@@ -7,28 +7,34 @@ from rapidata.api_client.models.user_score_user_filter_model import (
 
 class UserScoreFilter(RapidataFilter):
     """UserScoreFilter Class
-    
+
     Can be used to filter who to target based on their user score.
-    
+
     Args:
         lower_bound (float): The lower bound of the user score.
         upper_bound (float): The upper bound of the user score.
         dimension (str): The dimension of the userScore to be considerd for the filter.
 
-    Example: 
+    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, dimension: str | None = None):
+
+    def __init__(
+        self,
+        lower_bound: float = 0.0,
+        upper_bound: float = 1.0,
+        dimension: str | None = None,
+    ):
         if lower_bound < 0 or lower_bound > 1:
             raise ValueError("The lower bound must be between 0 and 1.")
         if upper_bound < 0 or upper_bound > 1:
             raise ValueError("The upper bound must be between 0 and 1.")
         if lower_bound >= upper_bound:
             raise ValueError("The lower bound must be less than the upper bound.")
-        
+
         self.upper_bound = upper_bound
         self.lower_bound = lower_bound
         self.dimension = dimension
@@ -40,3 +46,9 @@ class UserScoreFilter(RapidataFilter):
             lowerbound=self.lower_bound,
             dimension=self.dimension,
         )
+
+    def __str__(self) -> str:
+        return f"UserScoreFilter(lower_bound={self.lower_bound}, upper_bound={self.upper_bound}, dimension={self.dimension})"
+
+    def __repr__(self) -> str:
+        return f"UserScoreFilter(lower_bound={self.lower_bound!r}, upper_bound={self.upper_bound!r}, dimension={self.dimension!r})"

rapidata/rapidata_client/order/_rapidata_dataset.py

@@ -6,18 +6,17 @@ from concurrent.futures import ThreadPoolExecutor, as_completed
 from tqdm import tqdm
 
 from typing import Generator
-from rapidata.rapidata_client.logging import (
-    logger,
-    managed_print,
-    RapidataOutputManager,
-)
+from rapidata.rapidata_client.config import logger, managed_print
 import time
 import threading
-from rapidata.rapidata_client.api.rapidata_exception import (
+from rapidata.rapidata_client.api.rapidata_api_client import (
     suppress_rapidata_error_logging,
 )
 from rapidata.rapidata_client.config.rapidata_config import rapidata_config
 
+# Add OpenTelemetry context imports for thread propagation
+from opentelemetry import context as otel_context
+
 
 def chunk_list(lst: list, chunk_size: int) -> Generator:
     for i in range(0, len(lst), chunk_size):
@@ -62,22 +61,37 @@ class RapidataDataset:
             )
             return datapoint
 
+        def upload_with_context(
+            context: otel_context.Context, datapoint: Datapoint, index: int
+        ) -> Datapoint:
+            """Wrapper function that runs upload_text_datapoint with the provided context."""
+            token = otel_context.attach(context)
+            try:
+                return upload_text_datapoint(datapoint, index)
+            finally:
+                otel_context.detach(token)
+
         successful_uploads: list[Datapoint] = []
         failed_uploads: list[Datapoint] = []
 
+        # Capture the current OpenTelemetry context before creating threads
+        current_context = otel_context.get_current()
+
         total_uploads = len(datapoints)
         with ThreadPoolExecutor(
-            max_workers=rapidata_config.maxUploadWorkers
+            max_workers=rapidata_config.upload.maxWorkers
         ) as executor:
             future_to_datapoint = {
-                executor.submit(upload_text_datapoint, datapoint, index=i): datapoint
+                executor.submit(
+                    upload_with_context, current_context, datapoint, i
+                ): datapoint
                 for i, datapoint in enumerate(datapoints)
             }
 
             with tqdm(
                 total=total_uploads,
                 desc="Uploading text datapoints",
-                disable=RapidataOutputManager.silent_mode,
+                disable=rapidata_config.logging.silent_mode,
             ) as pbar:
                 for future in as_completed(future_to_datapoint.keys()):
                     datapoint = future_to_datapoint[future]
@@ -119,7 +133,7 @@ class RapidataDataset:
         urls = datapoint.get_urls()
 
         last_exception = None
-        for attempt in range(rapidata_config.uploadMaxRetries):
+        for attempt in range(rapidata_config.upload.maxRetries):
             try:
                 with suppress_rapidata_error_logging():
                     self.openapi_service.dataset_api.dataset_dataset_id_datapoints_post(
@@ -136,7 +150,7 @@ class RapidataDataset:
 
             except Exception as e:
                 last_exception = e
-                if attempt < rapidata_config.uploadMaxRetries - 1:
+                if attempt < rapidata_config.upload.maxRetries - 1:
                     # Exponential backoff: wait 1s, then 2s, then 4s
                     retry_delay = 2**attempt
                     time.sleep(retry_delay)
@@ -144,13 +158,13 @@ class RapidataDataset:
                     logger.debug(
                         "Retrying %s of %s...",
                         attempt + 1,
-                        rapidata_config.uploadMaxRetries,
+                        rapidata_config.upload.maxRetries,
                     )
 
         # If we get here, all retries failed
         local_failed.append(datapoint)
         tqdm.write(
-            f"Upload failed for {datapoint} after {rapidata_config.uploadMaxRetries} attempts. \nFinal error: \n{str(last_exception)}"
+            f"Upload failed for {datapoint} after {rapidata_config.upload.maxRetries} attempts. \nFinal error: \n{str(last_exception)}"
         )
 
         return local_successful, local_failed
@@ -183,7 +197,7 @@ class RapidataDataset:
                 with tqdm(
                     total=total_uploads,
                     desc="Uploading datapoints",
-                    disable=RapidataOutputManager.silent_mode,
+                    disable=rapidata_config.logging.silent_mode,
                 ) as pbar:
                     prev_ready = 0
                     prev_failed = 0
@@ -291,17 +305,31 @@ class RapidataDataset:
         successful_uploads: list[Datapoint] = []
         failed_uploads: list[Datapoint] = []
 
+        def process_upload_with_context(
+            context: otel_context.Context, datapoint: Datapoint, index: int
+        ) -> tuple[list[Datapoint], list[Datapoint]]:
+            """Wrapper function that runs _process_single_upload with the provided context."""
+            token = otel_context.attach(context)
+            try:
+                return self._process_single_upload(datapoint, index)
+            finally:
+                otel_context.detach(token)
+
+        # Capture the current OpenTelemetry context before creating threads
+        current_context = otel_context.get_current()
+
         try:
             with ThreadPoolExecutor(
-                max_workers=rapidata_config.maxUploadWorkers
+                max_workers=rapidata_config.upload.maxWorkers
             ) as executor:
                 # Process uploads in chunks to avoid overwhelming the system
                 for chunk_idx, chunk in enumerate(chunk_list(datapoints, chunk_size)):
                     futures = [
                         executor.submit(
-                            self._process_single_upload,
+                            process_upload_with_context,
+                            current_context,
                             datapoint,
-                            index=(chunk_idx * chunk_size + i),
+                            chunk_idx * chunk_size + i,
                         )
                         for i, datapoint in enumerate(chunk)
                     ]

rapidata/rapidata_client/order/_rapidata_order_builder.py

@@ -19,10 +19,11 @@ from rapidata.rapidata_client.exceptions.failed_upload_exception import (
     _parse_failed_uploads,
 )
 from rapidata.rapidata_client.filter import RapidataFilter
-from rapidata.rapidata_client.logging import (
+from rapidata.rapidata_client.config import (
     logger,
     managed_print,
-    RapidataOutputManager,
+    rapidata_config,
+    tracer,
 )
 from rapidata.rapidata_client.validation.validation_set_manager import (
     ValidationSetManager,
@@ -35,9 +36,9 @@ from rapidata.rapidata_client.selection._base_selection import RapidataSelection
 from rapidata.rapidata_client.settings import RapidataSetting
 from rapidata.rapidata_client.workflow import Workflow
 from rapidata.service.openapi_service import OpenAPIService
-from rapidata.rapidata_client.config.rapidata_config import rapidata_config
-from rapidata.rapidata_client.api.rapidata_exception import (
+from rapidata.rapidata_client.api.rapidata_api_client import (
     suppress_rapidata_error_logging,
+    RapidataApiClient,
 )
 
 
@@ -153,14 +154,19 @@ class RapidataOrderBuilder:
                 "Using recommended validation set with ID: %s", self.__validation_set_id
             )
         except Exception as e:
-            logger.info("No recommended validation set found, creating new one.")
+            logger.debug("No recommended validation set found, error: %s", e)
 
-        if len(self.__datapoints) < rapidata_config.minOrderDatapointsForValidation:
+        if (
+            len(self.__datapoints)
+            < rapidata_config.order.minOrderDatapointsForValidation
+        ):
             logger.debug(
                 "No recommended validation set found, dataset too small to create one."
             )
             return
 
+        logger.info("No recommended validation set found, creating new one.")
+
         managed_print()
         managed_print(
             f"No recommended validation set found, new one will be created.\nWe recommend adding some truths to ensure the order is accurate."
@@ -170,7 +176,9 @@ class RapidataOrderBuilder:
             order_name=self._name,
             datapoints=random.sample(
                 self.__datapoints,
-                min(rapidata_config.autoValidationSetSize, len(self.__datapoints)),
+                min(
+                    rapidata_config.order.autoValidationSetSize, len(self.__datapoints)
+                ),
             ),
             settings=self.__settings,
         )
@@ -182,9 +190,6 @@ class RapidataOrderBuilder:
         """
         Create the Rapidata order by making the necessary API calls based on the builder's configuration.
 
-        Args:
-            max_upload_workers (int, optional): The maximum number of worker threads for processing media paths. Defaults to 10.
-
         Raises:
             ValueError: If both media paths and texts are provided, or if neither is provided.
             AssertionError: If the workflow is a CompareWorkflow and media paths are not in pairs.
@@ -225,10 +230,11 @@ class RapidataOrderBuilder:
         logger.debug("Adding media to the order.")
 
         if self.__dataset:
-            _, failed_uploads = self.__dataset.add_datapoints(self.__datapoints)
+            with tracer.start_as_current_span("add_datapoints"):
+                _, failed_uploads = self.__dataset.add_datapoints(self.__datapoints)
 
-            if failed_uploads:
-                raise FailedUploadException(self.__dataset, order, failed_uploads)
+                if failed_uploads:
+                    raise FailedUploadException(self.__dataset, order, failed_uploads)
 
         else:
             raise RuntimeError(

rapidata/rapidata_client/order/rapidata_order.py

@@ -20,12 +20,13 @@ from rapidata.api_client.models.preliminary_download_model import (
 from rapidata.api_client.models.workflow_artifact_model import WorkflowArtifactModel
 from rapidata.rapidata_client.order.rapidata_results import RapidataResults
 from rapidata.service.openapi_service import OpenAPIService
-from rapidata.rapidata_client.logging import (
+from rapidata.rapidata_client.config import (
     logger,
     managed_print,
-    RapidataOutputManager,
+    rapidata_config,
+    tracer,
 )
-from rapidata.rapidata_client.api.rapidata_exception import (
+from rapidata.rapidata_client.api.rapidata_api_client import (
     suppress_rapidata_error_logging,
 )
 
@@ -74,36 +75,40 @@ class RapidataOrder:
 
     def run(self) -> "RapidataOrder":
         """Runs the order to start collecting responses."""
-        logger.info("Starting order '%s'", self)
-        self.__openapi_service.order_api.order_order_id_submit_post(
-            self.id, SubmitOrderModel(ignoreFailedDatapoints=True)
-        )
-        logger.debug("Order '%s' has been started.", self)
-        managed_print(
-            f"Order '{self.name}' is now viewable under: {self.order_details_page}"
-        )
-        return self
+        with tracer.start_as_current_span("RapidataOrder.run"):
+            logger.info("Starting order '%s'", self)
+            self.__openapi_service.order_api.order_order_id_submit_post(
+                self.id, SubmitOrderModel(ignoreFailedDatapoints=True)
+            )
+            logger.debug("Order '%s' has been started.", self)
+            managed_print(
+                f"Order '{self.name}' is now viewable under: {self.order_details_page}"
+            )
+            return self
 
     def pause(self) -> None:
         """Pauses the order."""
-        logger.info("Pausing order '%s'", self)
-        self.__openapi_service.order_api.order_order_id_pause_post(self.id)
-        logger.debug("Order '%s' has been paused.", self)
-        managed_print(f"Order '{self}' has been paused.")
+        with tracer.start_as_current_span("RapidataOrder.pause"):
+            logger.info("Pausing order '%s'", self)
+            self.__openapi_service.order_api.order_order_id_pause_post(self.id)
+            logger.debug("Order '%s' has been paused.", self)
+            managed_print(f"Order '{self}' has been paused.")
 
     def unpause(self) -> None:
         """Unpauses/resumes the order."""
-        logger.info("Unpausing order '%s'", self)
-        self.__openapi_service.order_api.order_order_id_resume_post(self.id)
-        logger.debug("Order '%s' has been unpaused.", self)
-        managed_print(f"Order '{self}' has been unpaused.")
+        with tracer.start_as_current_span("RapidataOrder.unpause"):
+            logger.info("Unpausing order '%s'", self)
+            self.__openapi_service.order_api.order_order_id_resume_post(self.id)
+            logger.debug("Order '%s' has been unpaused.", self)
+            managed_print(f"Order '{self}' has been unpaused.")
 
     def delete(self) -> None:
         """Deletes the order."""
-        logger.info("Deleting order '%s'", self)
-        self.__openapi_service.order_api.order_order_id_delete(self.id)
-        logger.debug("Order '%s' has been deleted.", self)
-        managed_print(f"Order '{self}' has been deleted.")
+        with tracer.start_as_current_span("RapidataOrder.delete"):
+            logger.info("Deleting order '%s'", self)
+            self.__openapi_service.order_api.order_order_id_delete(self.id)
+            logger.debug("Order '%s' has been deleted.", self)
+            managed_print(f"Order '{self}' has been deleted.")
 
     def get_status(self) -> str:
         """
@@ -119,7 +124,8 @@ class RapidataOrder:
             Completed: The order has been completed.\n
             Failed: The order has failed.
         """
-        return self.__openapi_service.order_api.order_order_id_get(self.id).state
+        with tracer.start_as_current_span("RapidataOrder.get_status"):
+            return self.__openapi_service.order_api.order_order_id_get(self.id).state
 
     def display_progress_bar(self, refresh_rate: int = 5) -> None:
         """
@@ -150,7 +156,7 @@ class RapidataOrder:
             desc="Processing order",
             unit="%",
             bar_format="{desc}: {percentage:3.0f}%|{bar}| completed [{elapsed}<{remaining}, {rate_fmt}]",
-            disable=RapidataOutputManager.silent_mode,
+            disable=rapidata_config.logging.silent_mode,
         ) as pbar:
             last_percentage = 0
             while True:
@@ -193,31 +199,37 @@ class RapidataOrder:
             preliminary_results: If True, returns the preliminary results of the order. Defaults to False.
                 Note that preliminary results are not final and may not contain all the datapoints & responses. Only the onese that are already available.
         """
-        logger.info("Getting results for order '%s'...", self)
-        if preliminary_results and self.get_status() not in [OrderState.COMPLETED]:
-            return self.__get_preliminary_results()
-
-        elif preliminary_results and self.get_status() in [OrderState.COMPLETED]:
-            managed_print("Order is already completed. Returning final results.")
-
-        while self.get_status() not in [
-            OrderState.COMPLETED,
-            OrderState.PAUSED,
-            OrderState.MANUALREVIEW,
-            OrderState.FAILED,
-        ]:
-            sleep(5)
+        with tracer.start_as_current_span("RapidataOrder.get_results"):
+            logger.info("Getting results for order '%s'...", self)
+            if preliminary_results and self.get_status() not in [OrderState.COMPLETED]:
+                return self.__get_preliminary_results()
+
+            elif preliminary_results and self.get_status() in [OrderState.COMPLETED]:
+                managed_print("Order is already completed. Returning final results.")
+
+            while (state := self.get_status()) not in [
+                OrderState.COMPLETED,
+                OrderState.PAUSED,
+                OrderState.MANUALREVIEW,
+                OrderState.FAILED,
+            ]:
+                sleep(5)
+                logger.debug(
+                    "Order '%s' is in state %s not yet completed. Waiting...",
+                    self,
+                    state,
+                )
 
-        try:
-            return RapidataResults(
-                json.loads(
-                    self.__openapi_service.order_api.order_order_id_download_results_get(
-                        order_id=self.id
+            try:
+                return RapidataResults(
+                    json.loads(
+                        self.__openapi_service.order_api.order_order_id_download_results_get(
+                            order_id=self.id
+                        )
                     )
                 )
-            )
-        except (ApiException, json.JSONDecodeError) as e:
-            raise Exception(f"Failed to get order results: {str(e)}") from e
+            except (ApiException, json.JSONDecodeError) as e:
+                raise Exception(f"Failed to get order results: {str(e)}") from e
 
     def view(self) -> None:
         """