rapidata 1.8.3__py3-none-any.whl → 2.44.4__py3-none-any.whl

rapidata/rapidata_client/order/dataset/_rapidata_dataset.py

@@ -0,0 +1,153 @@
+from rapidata.rapidata_client.datapoints._datapoint import Datapoint
+from rapidata.service.openapi_service import OpenAPIService
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from tqdm import tqdm
+
+from rapidata.rapidata_client.config import logger
+import time
+from rapidata.rapidata_client.api.rapidata_api_client import (
+    suppress_rapidata_error_logging,
+)
+from rapidata.rapidata_client.config.rapidata_config import rapidata_config
+from rapidata.rapidata_client.datapoints._datapoint_uploader import DatapointUploader
+
+# Add OpenTelemetry context imports for thread propagation
+from opentelemetry import context as otel_context
+
+
+class RapidataDataset:
+    def __init__(self, dataset_id: str, openapi_service: OpenAPIService):
+        self.id = dataset_id
+        self.openapi_service = openapi_service
+        self.datapoint_uploader = DatapointUploader(openapi_service)
+
+    def add_datapoints(
+        self,
+        datapoints: list[Datapoint],
+    ) -> tuple[list[Datapoint], list[Datapoint]]:
+        """
+        Process uploads concurrently with a ThreadPoolExecutor.
+
+        Args:
+            datapoints: List of datapoints to upload
+
+        Returns:
+            tuple[list[Datapoint], list[Datapoint]]: Lists of successful and failed uploads
+        """
+        successful_uploads: list[Datapoint] = []
+        failed_uploads: list[Datapoint] = []
+
+        with tqdm(
+            total=len(datapoints),
+            desc="Uploading datapoints",
+            disable=rapidata_config.logging.silent_mode,
+        ) as progress_bar:
+
+            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()
+
+            with ThreadPoolExecutor(
+                max_workers=rapidata_config.upload.maxWorkers
+            ) as executor:
+                # Submit all uploads at once - executor manages concurrency via max_workers
+                futures = [
+                    executor.submit(
+                        process_upload_with_context,
+                        current_context,
+                        datapoint,
+                        i,
+                    )
+                    for i, datapoint in enumerate(datapoints)
+                ]
+
+                # Process results as they complete (no blocking on chunks)
+                for future in as_completed(futures):
+                    try:
+                        upload_successful, upload_failed = future.result()
+                        successful_uploads.extend(upload_successful)
+                        failed_uploads.extend(upload_failed)
+                        progress_bar.update(len(upload_successful))
+                    except Exception as e:
+                        logger.error("Future execution failed: %s", str(e))
+
+        if failed_uploads:
+            logger.error(
+                "Upload failed for %s datapoints: %s",
+                len(failed_uploads),
+                failed_uploads,
+            )
+
+        return successful_uploads, failed_uploads
+
+    def _process_single_upload(
+        self,
+        datapoint: Datapoint,
+        index: int,
+    ) -> tuple[list[Datapoint], list[Datapoint]]:
+        """
+        Process single upload with retry logic and error tracking.
+
+        Args:
+            media_asset: MediaAsset or MultiAsset to upload
+            meta_list: Optional sequence of metadata for the asset
+            index: Sort index for the upload
+            max_retries: Maximum number of retry attempts (default: 3)
+
+        Returns:
+            tuple[list[Datapoint], list[Datapoint]]: Lists of successful and failed datapoints
+        """
+        logger.debug("Processing single upload for %s with index %s", datapoint, index)
+
+        local_successful: list[Datapoint] = []
+        local_failed: list[Datapoint] = []
+
+        last_exception = None
+        for attempt in range(rapidata_config.upload.maxRetries):
+            try:
+                with suppress_rapidata_error_logging():
+                    self.datapoint_uploader.upload_datapoint(
+                        dataset_id=self.id,
+                        datapoint=datapoint,
+                        index=index,
+                    )
+
+                local_successful.append(datapoint)
+
+                return local_successful, local_failed
+
+            except Exception as e:
+                last_exception = e
+                if attempt < rapidata_config.upload.maxRetries - 1:
+                    # Exponential backoff: wait 1s, then 2s, then 4s
+                    retry_delay = 2**attempt
+                    time.sleep(retry_delay)
+                    logger.debug("Error: %s", str(last_exception))
+                    logger.debug(
+                        "Retrying %s of %s...",
+                        attempt + 1,
+                        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.upload.maxRetries} attempts. \nFinal error: \n{str(last_exception)}"
+        )
+
+        return local_successful, local_failed
+
+    def __str__(self) -> str:
+        return f"RapidataDataset(id={self.id})"
+
+    def __repr__(self) -> str:
+        return self.__str__()

rapidata/rapidata_client/order/rapidata_order.py

@@ -1,62 +1,320 @@
-from time import sleep
-from rapidata.rapidata_client.dataset.rapidata_dataset import RapidataDataset
-from rapidata.service.openapi_service import OpenAPIService
+from __future__ import annotations
+
+
+# Standard library imports
 import json
-from rapidata.api_client.exceptions import ApiException
-from typing import Optional, cast, Any
-from rapidata.api_client.models.workflow_artifact_model import WorkflowArtifactModel
+import urllib.parse
+import webbrowser
+from time import sleep
+from typing import cast, Callable, TypeVar, TYPE_CHECKING
+from colorama import Fore
+from datetime import datetime
 from tqdm import tqdm
 
+# Local/application imports
+from rapidata.service.openapi_service import OpenAPIService
+from rapidata.rapidata_client.config import (
+    logger,
+    managed_print,
+    rapidata_config,
+    tracer,
+)
+from rapidata.rapidata_client.api.rapidata_api_client import (
+    suppress_rapidata_error_logging,
+)
+
+if TYPE_CHECKING:
+    from rapidata.api_client.models.campaign_artifact_model import CampaignArtifactModel
+    from rapidata.api_client.models.file_stream_result import FileStreamResult
+    from rapidata.api_client.models.order_state import OrderState
+    from rapidata.api_client.models.workflow_artifact_model import WorkflowArtifactModel
+    from rapidata.api_client.models.get_workflow_progress_result import (
+        GetWorkflowProgressResult,
+    )
+    from rapidata.rapidata_client.order.rapidata_results import RapidataResults
+
+T = TypeVar("T")
+
+
 class RapidataOrder:
     """
-    Represents a Rapidata order.
+    An instance of a Rapidata order.
+
+    Used to interact with a specific order in the Rapidata system, such as starting,
+    pausing, and retrieving results.
 
     Args:
-        The ID of the order.
-        The optional Dataset associated with the order.
-        The OpenAPIService instance used to interact with the Rapidata API.
-        The name of the order.
+        name: The name of the order.
+        order_id: The ID of the order.
+        openapi_service: The OpenAPIService instance for API interaction.
     """
 
     def __init__(
         self,
+        name: str,
         order_id: str,
-        dataset: Optional[RapidataDataset],
         openapi_service: OpenAPIService,
-        name: str,
     ):
-        self.openapi_service = openapi_service
-        self.order_id = order_id
-        self._dataset = dataset
+        self.id = order_id
         self.name = name
-        self._workflow_id = None
+        self.__created_at: datetime | None = None
+        self._openapi_service = openapi_service
+        self.__workflow_id: str = ""
+        self.__campaign_id: str = ""
+        self.__pipeline_id: str = ""
+        self.order_details_page = (
+            f"https://app.{self._openapi_service.environment}/order/detail/{self.id}"
+        )
+        logger.debug("RapidataOrder initialized")
+
+    def _get_order_failure_message(self) -> str | None:
+        """Retrieves the failure message from the order if available."""
+        try:
+            order = self._openapi_service.order_api.order_order_id_get(self.id)
+            return order.failure_message
+        except Exception:
+            logger.debug("Failed to get order failure message", self, exc_info=True)
+            return None
+
+    def _retry_operation(
+        self,
+        operation: Callable[[], T],
+        max_retries: int = 10,
+        retry_delay: float = 2,
+    ) -> T:
+        """
+        Unified retry logic for all operations with failure message handling.
+
+        Args:
+            operation: The operation to retry
+            max_retries: Maximum number of retry attempts
+            retry_delay: Delay between retries in seconds
+
+        Returns:
+            The result of the operation
 
-    def submit(self):
+        Raises:
+            Exception: If the operation fails after all retries, includes order failure message if available
         """
-        Submits the order for processing.
+        last_exception = None
+
+        for attempt in range(max_retries):
+            try:
+                return operation()
+            except Exception as e:
+                last_exception = e
+                if attempt < max_retries - 1:
+                    sleep(retry_delay)
+
+        failure_message = self._get_order_failure_message()
+        if failure_message:
+            raise Exception(failure_message) from last_exception
+
+        raise Exception(
+            f"Operation failed after {max_retries} retries: {str(last_exception)}"
+        ) from last_exception
+
+    def _wait_for_state(
+        self,
+        target_states: list[OrderState],
+        check_interval: float = 5,
+        status_message: str | None = None,
+    ) -> str:
         """
-        self.openapi_service.order_api.order_submit_post(self.order_id)
+        Wait until the order reaches one of the target states.
+
+        Args:
+            target_states: List of states to wait for
+            check_interval: How often to check the state in seconds
+            status_message: Optional message to display while waiting
+
+        Returns:
+            The final state reached
+        """
+        while (current_state := self.get_status()) not in target_states:
+            if status_message:
+                logger.debug(status_message, self, current_state)
+            sleep(check_interval)
+
+        return current_state
+
+    @property
+    def created_at(self) -> datetime:
+        """Returns the creation date of the order."""
+        if not self.__created_at:
+            self.__created_at = self._openapi_service.order_api.order_order_id_get(
+                self.id
+            ).order_date
+        return self.__created_at
+
+    def __get_pipeline_id(self) -> str:
+        """Gets the pipeline ID for this order (cached, internal use only)."""
+        if not self.__pipeline_id:
+            self.__pipeline_id = self._retry_operation(
+                lambda: self._openapi_service.order_api.order_order_id_get(
+                    self.id
+                ).pipeline_id,
+            )
+        return self.__pipeline_id
+
+    def __get_workflow_id(self) -> str:
+        """Gets the workflow ID for this order (cached, internal use only)."""
+        if not self.__workflow_id:
+            self.__load_workflow_and_campaign_ids()
+        return self.__workflow_id
+
+    def __get_campaign_id(self) -> str:
+        """Gets the campaign ID for this order (cached, internal use only)."""
+        if not self.__campaign_id:
+            self.__load_workflow_and_campaign_ids()
+        return self.__campaign_id
+
+    def __load_workflow_and_campaign_ids(self) -> None:
+        """Loads workflow and campaign IDs from the pipeline (with retry logic)."""
+        if self.__workflow_id and self.__campaign_id:
+            return
+
+        def fetch_ids():
+            from rapidata.api_client.models.workflow_artifact_model import (
+                WorkflowArtifactModel,
+            )
+            from rapidata.api_client.models.campaign_artifact_model import (
+                CampaignArtifactModel,
+            )
+
+            pipeline_id = self.__get_pipeline_id()
+            pipeline = self._openapi_service.pipeline_api.pipeline_pipeline_id_get(
+                pipeline_id
+            )
+            self.__workflow_id = cast(
+                WorkflowArtifactModel,
+                pipeline.artifacts["workflow-artifact"].actual_instance,
+            ).workflow_id
+            self.__campaign_id = cast(
+                CampaignArtifactModel,
+                pipeline.artifacts["campaign-artifact"].actual_instance,
+            ).campaign_id
+
+        self._retry_operation(fetch_ids)
+
+    def __get_workflow_progress(self) -> GetWorkflowProgressResult:
+        """Gets the workflow progress (internal use only)."""
+
+        def get_progress():
+            from rapidata.api_client.models.get_workflow_progress_result import (
+                GetWorkflowProgressResult,
+            )
+
+            with suppress_rapidata_error_logging():
+                workflow_id = self.__get_workflow_id()
+                return self._openapi_service.workflow_api.workflow_workflow_id_progress_get(
+                    workflow_id
+                )
+
+        return self._retry_operation(
+            get_progress,
+            max_retries=5,
+            retry_delay=4,
+        )
+
+    def run(self) -> RapidataOrder:
+        """Runs the order to start collecting responses."""
+        with tracer.start_as_current_span("RapidataOrder.run"):
+            from rapidata.api_client.models.submit_order_model import SubmitOrderModel
+
+            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."""
+        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."""
+        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."""
+        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:
         """
         Gets the status of the order.
-        
-        Returns: 
-            The status of the order.
+
+        States:
+            Created: The order has been created but not started yet.\n
+            Preview: The order has been set up and ready but not collecting responses yet.\n
+            Submitted: The order has been submitted and is being reviewed.\n
+            ManualReview: The order is in manual review - something went wrong with the automatic approval.\n
+            Processing: The order is actively being processed.\n
+            Paused: The order has been paused.\n
+            Completed: The order has been completed.\n
+            Failed: The order has failed.
         """
-        return self.openapi_service.order_api.order_get_by_id_get(self.order_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=5):
+    def display_progress_bar(self, refresh_rate: int = 5) -> None:
         """
         Displays a progress bar for the order processing using tqdm.
-        
-        Prameter: 
-            How often to refresh the progress bar, in seconds.
+
+        Args:
+            refresh_rate: How often to refresh the progress bar, in seconds.
         """
-        with tqdm(total=100, desc="Processing order", unit="%", bar_format="{desc}: {percentage:3.0f}%|{bar}| completed [{elapsed}<{remaining}, {rate_fmt}]") as pbar:
+        from rapidata.api_client.models.order_state import OrderState
+
+        if refresh_rate < 1:
+            raise ValueError("refresh_rate must be at least 1")
+
+        if self.get_status() == OrderState.CREATED:
+            raise Exception("Order has not been started yet. Please start it first.")
+
+        # Wait for submission review
+        while self.get_status() == OrderState.SUBMITTED:
+            managed_print(
+                f"Order '{self}' is submitted and being reviewed. Standby...", end="\r"
+            )
+            sleep(1)
+
+        if self.get_status() == OrderState.MANUALREVIEW:
+            raise Exception(
+                f"Order '{self}' is in manual review. It might take some time to start. "
+                f"To speed up the process, contact support (info@rapidata.ai).\n"
+                f"Once started, run this method again to display the progress bar."
+            )
+
+        with tqdm(
+            total=100,
+            desc="Processing order",
+            unit="%",
+            bar_format="{desc}: {percentage:3.0f}%|{bar}| completed [{elapsed}<{remaining}, {rate_fmt}]",
+            disable=rapidata_config.logging.silent_mode,
+        ) as pbar:
             last_percentage = 0
             while True:
-                current_percentage = self._get_workflow_progress().completion_percentage
+                current_percentage = (
+                    self.__get_workflow_progress().completion_percentage
+                )
+
                 if current_percentage > last_percentage:
                     pbar.update(current_percentage - last_percentage)
                     last_percentage = current_percentage
@@ -66,72 +324,117 @@ class RapidataOrder:
 
                 sleep(refresh_rate)
 
-    def _get_workflow_id(self):
-        if self._workflow_id:
-            return self._workflow_id
-
-        for _ in range(10):
-            try:
-                order_result = self.openapi_service.order_api.order_get_by_id_get(self.order_id)
-                pipeline = self.openapi_service.pipeline_api.pipeline_id_get(order_result.pipeline_id)
-                self._workflow_id = cast(WorkflowArtifactModel, pipeline.artifacts["workflow-artifact"].actual_instance).workflow_id
-                break
-            except Exception:
-                sleep(2)
-        if not self._workflow_id:
-            raise Exception("Order has not started yet. Please start it or wait for a few seconds and try again.")
-        return self._workflow_id
-    
-    def _get_workflow_progress(self):
-        workflow_id = self._get_workflow_id()
-        progress = None
-        for _ in range(2):
-            try:
-                progress = self.openapi_service.workflow_api.workflow_get_progress_get(workflow_id)
-                break
-            except Exception:
-                sleep(5)
-
-        if not progress:
-            raise Exception(f"Failed to get progress. Please try again in a few seconds.")
-        
-        return progress
-
-        
-    def get_results(self) -> dict[str, Any]:
+    def get_results(self, preliminary_results: bool = False) -> RapidataResults:
         """
-        Gets the results of the order. 
+        Gets the results of the order.
         If the order is still processing, this method will block until the order is completed and then return the results.
 
-        Returns: 
-            The results of the order.
+        Args:
+            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 ones that are already available.
         """
-        while self.get_status() not in ["Completed", "Paused", "ManuelReview"]:
-            sleep(5)
+        with tracer.start_as_current_span("RapidataOrder.get_results"):
+            from rapidata.api_client.models.order_state import OrderState
+            from rapidata.api_client.exceptions import ApiException
+            from rapidata.rapidata_client.order.rapidata_results import RapidataResults
+
+            logger.info("Getting results for order '%s'...", self)
+
+            if preliminary_results and self.get_status() not in [OrderState.COMPLETED]:
+                return self._get_preliminary_results()
+
+            if preliminary_results and self.get_status() == OrderState.COMPLETED:
+                managed_print("Order is already completed. Returning final results.")
+
+            self._wait_for_state(
+                target_states=[
+                    OrderState.COMPLETED,
+                    OrderState.PAUSED,
+                    OrderState.MANUALREVIEW,
+                    OrderState.FAILED,
+                ],
+                status_message="Order '%s' is in state %s, waiting for completion...",
+            )
+
+            try:
+                results = (
+                    self._openapi_service.order_api.order_order_id_download_results_get(
+                        order_id=self.id
+                    )
+                )
+                return RapidataResults(json.loads(results))
+            except (ApiException, json.JSONDecodeError) as e:
+                raise Exception(f"Failed to get order results: {str(e)}") from e
+
+    def _get_preliminary_results(self) -> RapidataResults:
+        """Fetches preliminary results for an in-progress order."""
+        from rapidata.api_client.models.preliminary_download_model import (
+            PreliminaryDownloadModel,
+        )
+        from rapidata.api_client.exceptions import ApiException
+        from rapidata.rapidata_client.order.rapidata_results import RapidataResults
 
         try:
-            # Get the raw result string
-            result_str = self.openapi_service.order_api.order_result_get(id=self.order_id)
-            # Parse the result string as JSON
-            return json.loads(result_str)
-        except ApiException as e:
-            # Handle API exceptions
-            raise Exception(f"Failed to get order results: {str(e)}") from e
-        except json.JSONDecodeError as e:
-            # Handle JSON parsing errors
-            raise Exception(f"Failed to parse order results: {str(e)}") from e
+            pipeline_id = self.__get_pipeline_id()
+            download_id = self._openapi_service.pipeline_api.pipeline_pipeline_id_preliminary_download_post(
+                pipeline_id, PreliminaryDownloadModel(sendEmail=False)
+            ).download_id
+
+            def check_results():
+                results = self._openapi_service.pipeline_api.pipeline_preliminary_download_preliminary_download_id_get(
+                    preliminary_download_id=download_id
+                )
+                return RapidataResults(json.loads(results))
+
+            return self._retry_operation(
+                check_results,
+                max_retries=60,
+                retry_delay=1,
+            )
+
+        except (ApiException, json.JSONDecodeError) as e:
+            raise Exception(f"Failed to get preliminary results: {str(e)}") from e
+
+    def view(self) -> None:
+        """Opens the order details page in the browser."""
+        logger.info("Opening order details page in browser...")
+        if not webbrowser.open(self.order_details_page):
+            encoded_url = urllib.parse.quote(
+                self.order_details_page, safe="%/:=&?~#+!$,;'@()*[]"
+            )
+            managed_print(
+                Fore.RED
+                + f"Please open this URL in your browser: '{encoded_url}'"
+                + Fore.RESET
+            )
+
+    def preview(self) -> None:
+        """Opens a preview of the order in the browser."""
+        from rapidata.api_client.models.order_state import OrderState
+        from rapidata.api_client.models.preview_order_model import PreviewOrderModel
+
+        logger.info("Opening order preview in browser...")
+
+        if self.get_status() == OrderState.CREATED:
+            logger.info("Order is still in state created. Setting it to preview.")
+            self._openapi_service.order_api.order_order_id_preview_post(
+                self.id, PreviewOrderModel(ignoreFailedDatapoints=True)
+            )
+            logger.info("Order is now in preview state.")
+
+        campaign_id = self.__get_campaign_id()
+        auth_url = f"https://app.{self._openapi_service.environment}/order/detail/{self.id}/preview?campaignId={campaign_id}"
+
+        if not webbrowser.open(auth_url):
+            encoded_url = urllib.parse.quote(auth_url, safe="%/:=&?~#+!$,;'@()*[]")
+            managed_print(
+                Fore.RED
+                + f"Please open this URL in your browser: '{encoded_url}'"
+                + Fore.RESET
+            )
 
-    @property
-    def dataset(self) -> RapidataDataset | None:
-        """
-        The dataset associated with the order.
-        Returns: 
-            The RapidataDataset instance.
-        """
-        return self._dataset
-    
     def __str__(self) -> str:
-        return f"name: '{self.name}' order id: {self.order_id}"
-    
+        return f"RapidataOrder(name='{self.name}', order_id='{self.id}')"
+
     def __repr__(self) -> str:
-        return f"name: '{self.name}' order id: {self.order_id}"
+        return f"RapidataOrder(name='{self.name}', order_id='{self.id}')"