rapidata 2.37.0__py3-none-any.whl → 3.0.0a8__py3-none-any.whl

rapidata/rapidata_client/order/rapidata_order.py

@@ -1,34 +1,40 @@
+from __future__ import annotations
+
+
 # Standard library imports
 import json
 import urllib.parse
 import webbrowser
 from time import sleep
-from typing import cast
+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.api_client.exceptions import ApiException
-from rapidata.api_client.models.campaign_artifact_model import CampaignArtifactModel
-from rapidata.api_client.models.order_state import OrderState
-from rapidata.api_client.models.preview_order_model import PreviewOrderModel
-from rapidata.api_client.models.submit_order_model import SubmitOrderModel
-from rapidata.api_client.models.preliminary_download_model import (
-    PreliminaryDownloadModel,
-)
-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,
 )
 
+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.results.rapidata_results import RapidataResults
+
+T = TypeVar("T")
+
 
 class RapidataOrder:
     """
@@ -52,58 +58,203 @@ class RapidataOrder:
         self.id = order_id
         self.name = name
         self.__created_at: datetime | None = None
-        self.__openapi_service = openapi_service
+        self._openapi_service = openapi_service
         self.__workflow_id: str = ""
         self.__campaign_id: str = ""
         self.__pipeline_id: str = ""
-        self._max_retries = 10
-        self._retry_delay = 2
         self.order_details_page = (
-            f"https://app.{self.__openapi_service.environment}/order/detail/{self.id}"
+            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
+
+        Raises:
+            Exception: If the operation fails after all retries, includes order failure message if available
+        """
+        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:
+        """
+        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.__created_at = self._openapi_service.order_api.order_order_id_get(
                 self.id
             ).order_date
         return self.__created_at
 
-    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}"
+    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,
         )
-        return self
+
+    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."""
-        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 +270,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:
         """
@@ -128,12 +280,15 @@ class RapidataOrder:
         Args:
             refresh_rate: How often to refresh the progress bar, in seconds.
         """
+        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"
@@ -142,7 +297,9 @@ class RapidataOrder:
 
         if self.get_status() == OrderState.MANUALREVIEW:
             raise Exception(
-                f"Order '{self}' is in manual review. It might take some time to start. To speed up the process, contact support (info@rapidata.ai).\nOnce started, run this method again to display the progress bar."
+                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(
@@ -150,11 +307,14 @@ 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:
-                current_percentage = self._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
@@ -164,26 +324,6 @@ class RapidataOrder:
 
                 sleep(refresh_rate)
 
-    @property
-    def _workflow_progress(self):
-        """
-        Gets the workflow progress.
-        """
-        progress = None
-        for _ in range(self._max_retries // 2):
-            try:
-                with suppress_rapidata_error_logging():
-                    workflow_id = self.__get_workflow_id()
-                    progress = self.__openapi_service.workflow_api.workflow_workflow_id_progress_get(
-                        workflow_id
-                    )
-                break
-            except Exception:
-                sleep(self._retry_delay * 2)
-        if not progress:
-            raise Exception("Failed to get progress. Please try again later.")
-        return progress
-
     def get_results(self, preliminary_results: bool = False) -> RapidataResults:
         """
         Gets the results of the order.
@@ -191,44 +331,74 @@ class RapidataOrder:
 
         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 onese that are already available.
+                Note that preliminary results are not final and may not contain all the datapoints & responses. Only the ones 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"):
+            from rapidata.api_client.models.order_state import OrderState
+            from rapidata.api_client.exceptions import ApiException
+            from rapidata.rapidata_client.results.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:
-            return RapidataResults(
-                json.loads(
-                    self.__openapi_service.order_api.order_order_id_download_results_get(
+            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.results.rapidata_results import RapidataResults
+
+        try:
+            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 order results: {str(e)}") from 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.
-
-        Raises:
-            Exception: If the order is not in processing state.
-        """
+        """Opens the order details page in the browser."""
         logger.info("Opening order details page in browser...")
-        could_open_browser = webbrowser.open(self.order_details_page)
-        if not could_open_browser:
+        if not webbrowser.open(self.order_details_page):
             encoded_url = urllib.parse.quote(
                 self.order_details_page, safe="%/:=&?~#+!$,;'@()*[]"
             )
@@ -239,24 +409,23 @@ class RapidataOrder:
             )
 
     def preview(self) -> None:
-        """
-        Opens a preview of the order in the browser.
+        """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
 
-        Raises:
-            Exception: If the order is not in processing state.
-        """
         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._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}"
-        could_open_browser = webbrowser.open(auth_url)
-        if not could_open_browser:
+        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
@@ -264,80 +433,6 @@ class RapidataOrder:
                 + Fore.RESET
             )
 
-    def __get_pipeline_id(self) -> str:
-        """Internal method to fetch and cache the pipeline ID."""
-        if not self.__pipeline_id:
-            for _ in range(self._max_retries):
-                try:
-                    self.__pipeline_id = (
-                        self.__openapi_service.order_api.order_order_id_get(
-                            self.id
-                        ).pipeline_id
-                    )
-                    break
-                except Exception:
-                    sleep(self._retry_delay)
-            else:
-                raise Exception("Failed to fetch pipeline ID after retries.")
-        return self.__pipeline_id
-
-    def __get_workflow_id(self) -> str:
-        """Internal method to fetch and cache the workflow ID."""
-        if not self.__workflow_id:
-            self.__fetch_workflow_and_campaign_ids()
-        return self.__workflow_id
-
-    def __get_campaign_id(self) -> str:
-        """Internal method to fetch and cache the campaign ID."""
-        if not self.__campaign_id:
-            self.__fetch_workflow_and_campaign_ids()
-        return self.__campaign_id
-
-    def __fetch_workflow_and_campaign_ids(self) -> None:
-        """Fetches workflow and campaign IDs from the pipeline."""
-        if self.__workflow_id and self.__campaign_id:
-            return
-        pipeline_id = self.__get_pipeline_id()
-        for _ in range(self._max_retries):
-            try:
-                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
-                return
-            except Exception:
-                sleep(self._retry_delay)
-        raise Exception("Failed to fetch workflow and campaign IDs after retries.")
-
-    def __get_preliminary_results(self) -> RapidataResults:
-        """Internal method to fetch preliminary results."""
-        try:
-            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
-
-            elapsed = 0
-            timeout = 60
-            while elapsed < timeout:
-                preliminary_results = self.__openapi_service.pipeline_api.pipeline_preliminary_download_preliminary_download_id_get(
-                    preliminary_download_id=download_id
-                )
-                if preliminary_results:
-                    return RapidataResults(json.loads(preliminary_results.decode()))
-                sleep(1)
-                elapsed += 1
-            raise Exception("Timeout waiting for preliminary results.")
-        except (ApiException, json.JSONDecodeError) as e:
-            raise Exception(f"Failed to get preliminary results: {str(e)}") from e
-
     def __str__(self) -> str:
         return f"RapidataOrder(name='{self.name}', order_id='{self.id}')"