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

rapidata/rapidata_client/order/rapidata_order.py

@@ -1,24 +1,39 @@
+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 logger, managed_print, RapidataOutputManager
+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.results.rapidata_results import RapidataResults
+
+T = TypeVar("T")
 
 
 class RapidataOrder:
@@ -43,50 +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}"
+        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
+
+        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.id).order_date
+            self.__created_at = self._openapi_service.order_api.order_order_id_get(
+                self.id
+            ).order_date
         return self.__created_at
-    
-    def run(self) -> "RapidataOrder":
+
+    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."""
-        logger.info(f"Starting order '{self}'")
-        self.__openapi_service.order_api.order_order_id_submit_post(self.id, SubmitOrderModel(ignoreFailedDatapoints=True))
-        logger.debug(f"Order '{self}' has been started.")
-        managed_print(f"Order '{self.name}' is now viewable under: {self.order_details_page}")
-        return self
+        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(f"Pausing order '{self}'")
-        self.__openapi_service.order_api.order_order_id_pause_post(self.id)
-        logger.debug(f"Order '{self}' has been paused.")
-        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(f"Unpausing order '{self}'")
-        self.__openapi_service.order_api.order_order_id_resume_post(self.id)
-        logger.debug(f"Order '{self}' has been unpaused.")
-        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(f"Deleting order '{self}'")
-        self.__openapi_service.order_api.order_order_id_delete(self.id)
-        logger.debug(f"Order '{self}' has been deleted.")
-        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:
         """
@@ -102,36 +270,51 @@ 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:
+    def display_progress_bar(self, refresh_rate: int = 5) -> None:
         """
         Displays a progress bar for the order processing using tqdm.
-        
-        Args: 
+
+        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")
+            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. "
-                "To speed up the process, contact support (info@rapidata.ai).\n"
-                "Once started, run this method again to display the progress bar."
+                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=RapidataOutputManager.silent_mode) as pbar:
+
+        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._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
@@ -141,142 +324,115 @@ class RapidataOrder:
 
                 sleep(refresh_rate)
 
-    @property
-    def _workflow_progress(self):
-        """
-        Gets the workflow progress.
-        """
-        progress = None
-        for _ in range(self._max_retries // 2):
-            try:
-                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. 
+        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.
 
         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.
+            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.
         """
-        logger.info(f"Getting results for order '{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.")
+        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
 
-        while self.get_status() not in [OrderState.COMPLETED, OrderState.PAUSED, OrderState.MANUALREVIEW, OrderState.FAILED]:
-            sleep(5)
+            logger.info("Getting results for order '%s'...", self)
 
-        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
-        
-    def view(self) -> None:
-        """
-        Opens the order details page in the browser.
-        
-        Raises:
-            Exception: If the order is not in processing state.
-        """
-        logger.info("Opening order details page in browser...")
-        could_open_browser = webbrowser.open(self.order_details_page)
-        if not could_open_browser:
-            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.
-        
-        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.id, PreviewOrderModel(ignoreFailedDatapoints=True))
-            logger.info("Order is now in preview state.")
+            if preliminary_results and self.get_status() not in [OrderState.COMPLETED]:
+                return self._get_preliminary_results()
 
-        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:
-            encoded_url = urllib.parse.quote(auth_url, safe="%/:=&?~#+!$,;'@()*[]")
-            managed_print(Fore.RED + f'Please open this URL in your browser: "{encoded_url}"' + Fore.RESET)
+            if preliminary_results and self.get_status() == OrderState.COMPLETED:
+                managed_print("Order is already completed. Returning final results.")
 
-    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
+            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...",
+            )
 
-    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
+            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_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 _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
 
-    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(
+            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(
+
+            def check_results():
+                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.")
+                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
+            )
+
     def __str__(self) -> str:
         return f"RapidataOrder(name='{self.name}', order_id='{self.id}')"