Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,6 @@
tracer,
)
from rapidata.rapidata_client.config._qr_preview import (
print_campaign_preview_qr_for_pipeline,
print_order_preview_link,
)
from rapidata.rapidata_client.validation.validation_set_manager import (
Expand Down Expand Up @@ -181,13 +180,6 @@ def _create(self) -> RapidataOrder:
except Exception as e:
logger.error("Failed to set order to preview: %s", e)

# The campaign is materialised once the order enters preview, so the
# campaign artifact only shows up on the pipeline a moment later — let
# the helper poll for it.
print_campaign_preview_qr_for_pipeline(
openapi_service=self._openapi_service,
pipeline_id=result.pipeline_id,
)
print_order_preview_link(
environment=self._openapi_service.environment,
order_id=order.id,
Expand Down
278 changes: 6 additions & 272 deletions src/rapidata/rapidata_client/order/rapidata_order.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,35 +6,22 @@
import urllib.parse
import webbrowser
from time import sleep
from typing import cast, Callable, TypeVar, TYPE_CHECKING
from typing import TYPE_CHECKING
from colorama import Fore
from datetime import datetime
from tqdm.auto 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_endpoint_output import (
GetWorkflowProgressEndpointOutput,
)
from rapidata.rapidata_client.results.rapidata_results import RapidataResults

T = TypeVar("T")


class RapidataOrder:
"""
Expand All @@ -59,61 +46,11 @@ def __init__(
self.name = name
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.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],
Expand Down Expand Up @@ -149,74 +86,6 @@ def created_at(self) -> datetime:
)
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.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.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) -> GetWorkflowProgressEndpointOutput:
"""Gets the workflow progress (internal use only)."""

def get_progress():
with suppress_rapidata_error_logging():
workflow_id = self.__get_workflow_id()
return self._openapi_service.workflow.workflow_api.workflow_workflow_id_progress_get(
workflow_id
)

return self._retry_operation(
get_progress,
max_retries=5,
retry_delay=4,
)

def run(self, after: RapidataOrder | str | None = None) -> RapidataOrder:
"""Runs the order to start collecting responses.
Args:
Expand Down Expand Up @@ -302,75 +171,6 @@ def get_status(self) -> str:
self.id
).state

def display_progress_bar(self, refresh_rate: int = 5) -> None:
"""
Displays a progress bar for the order processing using tqdm.

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"
)
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."
)

# Terminal states that should break the progress-bar loop
# regardless of completion_percentage. Without this check a
# Failed or Paused order would pin the caller's thread forever.
terminal_states = {
OrderState.COMPLETED,
OrderState.PAUSED,
OrderState.FAILED,
OrderState.MANUALREVIEW,
}

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
)

if current_percentage > last_percentage:
pbar.update(current_percentage - last_percentage)
last_percentage = current_percentage

if current_percentage >= 100:
break

current_state = self.get_status()
if current_state in terminal_states:
logger.info(
"Progress bar exiting early: order is in state %s",
current_state,
)
break

sleep(refresh_rate)

def _regenerate_results(self) -> None:
"""Triggers regeneration of an order whose results have gone stale.

Expand All @@ -385,18 +185,14 @@ def _regenerate_results(self) -> None:
)
self._openapi_service.order.order_api.order_order_id_retry_post(self.id)

def get_results(self, preliminary_results: bool = False) -> RapidataResults:
def get_results(self) -> RapidataResults:
"""
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.
If the order's results have gone stale, regeneration is triggered automatically and this method blocks until the fresh results are ready.

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.

Info:
Currently the SDK does not support streaming. The preliminary results are simply a snapshot of the results at the time of the request.
If the order is still processing, this method will block until the order
is completed and then return the results. If the order's results have
gone stale, regeneration is triggered automatically and this method
blocks until the fresh results are ready.
"""
with tracer.start_as_current_span("RapidataOrder.get_results"):
from rapidata.api_client.models.order_state import OrderState
Expand All @@ -407,12 +203,6 @@ def get_results(self, preliminary_results: bool = False) -> 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.")

# Stale results have no downloadable file until the pipeline is re-run;
# trigger that automatically before waiting for the re-completion.
if self.get_status() == OrderState.STALERESULTS:
Expand All @@ -436,35 +226,6 @@ def get_results(self, preliminary_results: bool = False) -> RapidataResults:
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.start_preliminary_download_endpoint_input import (
StartPreliminaryDownloadEndpointInput,
)
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.pipeline_api.pipeline_pipeline_id_preliminary_download_post(
pipeline_id, StartPreliminaryDownloadEndpointInput(sendEmail=False)
).download_id

def check_results():
results = self._openapi_service.pipeline.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...")
Expand All @@ -478,33 +239,6 @@ def view(self) -> None:
+ 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_endpoint_input import (
PreviewOrderEndpointInput,
)

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.order_api.order_order_id_preview_post(
self.id, PreviewOrderEndpointInput(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}')"

Expand Down
Loading