rapidata 2.28.2__py3-none-any.whl → 2.44.3__py3-none-any.whl

rapidata/rapidata_client/order/_rapidata_order_builder.py

@@ -1,35 +1,51 @@
-from rapidata.api_client.models.ab_test_selection_a_inner import (
-    AbTestSelectionAInner,
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Literal, Optional, Sequence, get_args
+import random
+import secrets
+
+from rapidata.rapidata_client.datapoints._datapoint import Datapoint
+
+if TYPE_CHECKING:
+    from rapidata.api_client.models.create_order_model import CreateOrderModel
+from rapidata.rapidata_client.exceptions.failed_upload_exception import (
+    FailedUploadException,
 )
-from rapidata.api_client.models.create_order_model import CreateOrderModel
-from rapidata.api_client.models.create_order_model_referee import (
-    CreateOrderModelReferee,
+from rapidata.rapidata_client.filter import RapidataFilter, UserScoreFilter
+from rapidata.rapidata_client.config import (
+    logger,
+    managed_print,
+    rapidata_config,
+    tracer,
 )
-from rapidata.api_client.models.create_order_model_user_filters_inner import (
-    CreateOrderModelUserFiltersInner,
+from rapidata.rapidata_client.validation.validation_set_manager import (
+    ValidationSetManager,
 )
-from rapidata.api_client.models.create_order_model_workflow import (
-    CreateOrderModelWorkflow,
+from rapidata.rapidata_client.validation.rapidata_validation_set import (
+    RapidataValidationSet,
 )
-
-from rapidata.rapidata_client.settings import RapidataSetting
-from rapidata.rapidata_client.metadata._base_metadata import Metadata
-from rapidata.rapidata_client.order._rapidata_dataset import RapidataDataset
-from rapidata.rapidata_client.referee._naive_referee import NaiveReferee
-from rapidata.rapidata_client.selection._base_selection import RapidataSelection
-from rapidata.rapidata_client.filter import RapidataFilter
-from rapidata.rapidata_client.workflow import Workflow
+from rapidata.rapidata_client.order.dataset._rapidata_dataset import RapidataDataset
 from rapidata.rapidata_client.order.rapidata_order import RapidataOrder
 from rapidata.rapidata_client.referee import Referee
+from rapidata.rapidata_client.referee._naive_referee import NaiveReferee
+from rapidata.rapidata_client.selection._base_selection import RapidataSelection
+from rapidata.rapidata_client.settings import RapidataSetting
+from rapidata.rapidata_client.workflow import (
+    Workflow,
+    CompareWorkflow,
+    ClassifyWorkflow,
+)
+from rapidata.rapidata_client.selection import (
+    ConditionalValidationSelection,
+    LabelingSelection,
+    CappedSelection,
+)
 from rapidata.service.openapi_service import OpenAPIService
+from rapidata.rapidata_client.api.rapidata_api_client import (
+    suppress_rapidata_error_logging,
+)
 
-from rapidata.rapidata_client.workflow._compare_workflow import CompareWorkflow
-
-from rapidata.rapidata_client.assets import MediaAsset, TextAsset, MultiAsset, BaseAsset
-
-from typing import Optional, cast, Sequence
-
-from rapidata.rapidata_client.logging import logger, managed_print
+StickyStateLiteral = Literal["Temporary", "Permanent", "Passive"]
 
 
 class RapidataOrderBuilder:
@@ -48,18 +64,21 @@ class RapidataOrderBuilder:
         openapi_service: OpenAPIService,
     ):
         self._name = name
-        self.order_id: str | None = None
-        self.__openapi_service = openapi_service
-        self.__dataset: Optional[RapidataDataset]
-        self.__workflow: Workflow | None = None
-        self.__referee: Referee | None = None
-        self.__multi_metadata: Sequence[Sequence[Metadata]] | None = None
-        self.__validation_set_id: str | None = None
-        self.__settings: Sequence[RapidataSetting] | None = None
-        self.__user_filters: list[RapidataFilter] = []
-        self.__selections: list[RapidataSelection] = []
-        self.__priority: int = 50
-        self.__assets: Sequence[BaseAsset] = []
+        self._openapi_service = openapi_service
+        self._dataset: Optional[RapidataDataset] = None
+        self._workflow: Workflow | None = None
+        self._referee: Referee | None = None
+        self._validation_set: RapidataValidationSet | None = None
+        self._settings: Sequence[RapidataSetting] | None = None
+        self._user_filters: list[RapidataFilter] = []
+        self._selections: Sequence[RapidataSelection] = []
+        self._priority: int | None = None
+        self._datapoints: list[Datapoint] = []
+        self._sticky_state_value: StickyStateLiteral | None = None
+        self._temporary_sticky_enabled: bool = False
+        self._validation_set_manager: ValidationSetManager = ValidationSetManager(
+            self._openapi_service
+        )
 
     def _to_model(self) -> CreateOrderModel:
         """
@@ -71,41 +90,162 @@ class RapidataOrderBuilder:
         Returns:
             CreateOrderModel: The model representing the order configuration.
         """
-        if self.__workflow is None:
+        if self._workflow is None:
             raise ValueError("You must provide a workflow to create an order.")
 
-        if self.__referee is None:
+        if self._referee is None:
             managed_print("No referee provided, using default NaiveReferee.")
-            self.__referee = NaiveReferee()
+            self._referee = NaiveReferee()
+
+        sticky_state = self._sticky_state_value
+        if not sticky_state and self._temporary_sticky_enabled:
+            sticky_state = "Temporary"
+            logger.debug(
+                "Setting sticky state to Temporary due to temporary sticky enabled."
+            )
+
+        validation_set_id = (
+            self._validation_set.id
+            if (self._validation_set and not self._selections)
+            else None
+        )
+        from rapidata.api_client.models.create_order_model import CreateOrderModel
+        from rapidata.api_client.models.sticky_state import StickyState
 
         return CreateOrderModel(
-            _t="CreateOrderModel",
             orderName=self._name,
-            workflow=CreateOrderModelWorkflow(self.__workflow._to_model()),
-            userFilters=[
-                CreateOrderModelUserFiltersInner(user_filter._to_model())
-                for user_filter in self.__user_filters
-            ],
-            referee=CreateOrderModelReferee(self.__referee._to_model()),
-            validationSetId=self.__validation_set_id,
+            workflow=self._workflow._to_model(),
+            userFilters=[user_filter._to_model() for user_filter in self._user_filters],
+            referee=self._referee._to_model(),
+            validationSetId=validation_set_id,
             featureFlags=(
-                [setting._to_feature_flag() for setting in self.__settings]
-                if self.__settings is not None
+                [setting._to_feature_flag() for setting in self._settings]
+                if self._settings is not None
+                else None
+            ),
+            selections=(
+                [selection._to_model() for selection in self._selections]
+                if self._selections
                 else None
             ),
-            selections=[
-                AbTestSelectionAInner(selection._to_model())
-                for selection in self.__selections
-            ],
-            priority=self.__priority,
+            priority=self._priority,
+            stickyState=(StickyState(sticky_state) if sticky_state else None),
         )
 
-    def _create(self, max_upload_workers: int = 10) -> RapidataOrder:
+    def _generate_id(self, length=9):
+        """Optimal balance: fast comparisons + low collision risk"""
+        alphabet = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"
+        return "".join(secrets.choice(alphabet) for _ in range(length))
+
+    def _attach_validation_set_id(self) -> None:
         """
-        Create the Rapidata order by making the necessary API calls based on the builder's configuration.
+        Sets the validation set for the order.
+        """
+        assert self._workflow is not None
+
+        required_amount = min(int(len(self._datapoints) * 0.01) or 3, 15)
+
+        if self._validation_set is None:
+            try:
+                with suppress_rapidata_error_logging():
+                    val_set_id = (
+                        (
+                            self._openapi_service.validation_api.validation_set_recommended_get(
+                                asset_type=[self._datapoints[0].get_asset_type()],
+                                modality=[self._workflow.modality],
+                                instruction=self._workflow._get_instruction(),
+                                prompt_type=[
+                                    t.value
+                                    for t in self._datapoints[0].get_prompt_type()
+                                ],
+                            )
+                        )
+                        .validation_sets[0]
+                        .id
+                    )
+                    self._validation_set = (
+                        self._validation_set_manager.get_validation_set_by_id(
+                            val_set_id
+                        )
+                    )
+
+            except Exception as e:
+                logger.debug("No recommended validation set found, error: %s", e)
+
+        sufficient_rapids_count = False
+        if self._validation_set is not None:
+            sufficient_rapids_count = (
+                self._validation_set_manager._get_total_and_labeled_rapids_count(
+                    self._validation_set.id
+                )[1]
+                >= required_amount
+            )
 
-        Args:
-            max_upload_workers (int, optional): The maximum number of worker threads for processing media paths. Defaults to 10.
+        if self._validation_set is None or not sufficient_rapids_count:
+            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."
+            )
+
+            new_dimension = self._generate_id()
+            logger.debug("New dimension created: %s", new_dimension)
+            rng = random.Random(42)
+            self._validation_set = (
+                self._validation_set_manager._create_order_validation_set(
+                    workflow=self._workflow,
+                    name=self._name,
+                    datapoints=rng.sample(
+                        self._datapoints, len(self._datapoints)
+                    ),  # shuffle the datapoints with a specific seed
+                    required_amount=required_amount,
+                    settings=self._settings,
+                    dimensions=[new_dimension],
+                )
+            )
+
+        self._validation_set.update_should_alert(False)
+        self._validation_set.update_can_be_flagged(False)
+
+        logger.debug("New validation set created for order: %s", self._validation_set)
+
+        self._selections = [
+            CappedSelection(
+                selections=[
+                    ConditionalValidationSelection(
+                        validation_set_id=self._validation_set.id,
+                        dimensions=self._validation_set.dimensions,
+                        thresholds=[0, 0.5, 0.7],
+                        chances=[1, 1, 0.2],
+                        rapid_counts=[10, 1, 1],
+                    ),
+                    LabelingSelection(amount=1),
+                ],
+                max_rapids=3,
+            )
+        ]
+
+        for dimension in self._validation_set.dimensions:
+            self._user_filters.append(
+                UserScoreFilter(
+                    lower_bound=0.3,
+                    upper_bound=1,
+                    dimension=dimension,
+                )
+            )
+
+    def _create(self) -> RapidataOrder:
+        """
+        Create the Rapidata order by making the necessary API calls based on the builder's configuration.
 
         Raises:
             ValueError: If both media paths and texts are provided, or if neither is provided.
@@ -114,51 +254,56 @@ class RapidataOrderBuilder:
         Returns:
             RapidataOrder: The created RapidataOrder instance.
         """
+        if (
+            rapidata_config.order.autoValidationSetCreation
+            and isinstance(self._workflow, (CompareWorkflow, ClassifyWorkflow))
+            and not self._selections
+            and rapidata_config.enableBetaFeatures
+        ):
+            self._attach_validation_set_id()
+            self._temporary_sticky_enabled = True
+            logger.debug("Temporary sticky enabled for order creation.")
+
         order_model = self._to_model()
-        logger.debug(f"Creating order with model: {order_model}")
-        if isinstance(
-            self.__workflow, CompareWorkflow
-        ):  # Temporary fix; will be handled by backend in the future
-            assert all(
-                isinstance(item, MultiAsset) for item in self.__assets
-            ), "The media paths must be of type MultiAsset for comparison tasks."
-
-        result = self.__openapi_service.order_api.order_post(
-            create_order_model=order_model
-        )
+        logger.debug("Creating order with model: %s", order_model)
 
-        self.order_id = str(result.order_id)
-        logger.debug(f"Order created with ID: {self.order_id}")
+        self._temporary_sticky_enabled = False
+        logger.debug("Disabling temporary sticky after order creation.")
 
-        self.__dataset = (
-            RapidataDataset(result.dataset_id, self.__openapi_service)
-            if result.dataset_id
-            else None
+        result = self._openapi_service.order_api.order_post(
+            create_order_model=order_model
         )
-        if self.__dataset:
-            logger.debug(f"Dataset created with ID: {self.__dataset.dataset_id}")
-        else:
-            logger.warning("No dataset created for this order.")
 
         order = RapidataOrder(
-            order_id=self.order_id,
-            openapi_service=self.__openapi_service,
+            order_id=str(result.order_id),
+            openapi_service=self._openapi_service,
             name=self._name,
         )
+        logger.debug("Order created: %s", order)
 
-        logger.debug(f"Order created: {order}")
-        logger.debug("Adding media to the order.")
+        if not result.dataset_id:
+            logger.error("No Dataset was created for order %s", order.id)
+            return order
 
-        if self.__dataset:
-            self.__dataset._add_datapoints(self.__assets, self.__multi_metadata, max_upload_workers)
-        
-        logger.debug("Media added to the order.")
-        logger.debug("Setting order to preview")
-        self.__openapi_service.order_api.order_order_id_preview_post(self.order_id)
+        self._dataset = RapidataDataset(result.dataset_id, self._openapi_service)
 
+        logger.debug("Dataset created with ID: %s", self._dataset.id)
+        logger.debug("Adding datapoints to the order.")
+        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)
+
+        logger.debug("Datapoints added to the order.")
+        logger.debug("Setting order to preview")
+        try:
+            self._openapi_service.order_api.order_order_id_preview_post(order.id)
+        except Exception:
+            raise FailedUploadException(self._dataset, order, failed_uploads)
         return order
 
-    def _workflow(self, workflow: Workflow) -> "RapidataOrderBuilder":
+    def _set_workflow(self, workflow: Workflow) -> RapidataOrderBuilder:
         """
         Set the workflow for the order.
 
@@ -171,10 +316,10 @@ class RapidataOrderBuilder:
         if not isinstance(workflow, Workflow):
             raise TypeError("Workflow must be of type Workflow.")
 
-        self.__workflow = workflow
+        self._workflow = workflow
         return self
 
-    def _referee(self, referee: Referee) -> "RapidataOrderBuilder":
+    def _set_referee(self, referee: Referee) -> "RapidataOrderBuilder":
         """
         Set the referee for the order.
 
@@ -187,52 +332,33 @@ class RapidataOrderBuilder:
         if not isinstance(referee, Referee):
             raise TypeError("Referee must be of type Referee.")
 
-        self.__referee = referee
+        self._referee = referee
         return self
 
-    def _media(
+    def _set_datapoints(
         self,
-        assets: Sequence[BaseAsset],
-        multi_metadata: Sequence[Sequence[Metadata]] | None = None,
-    ) -> "RapidataOrderBuilder":
+        datapoints: list[Datapoint],
+    ) -> RapidataOrderBuilder:
         """
-        Set the media assets for the order.
+        Set the datapoints for the order.
 
         Args:
-            assets: (list[MediaAsset] | list[TextAsset] | list[MultiAsset]): The paths of the media assets to be set.
-            multi_metadata: (list[list[Metadata]] | None, optional): Metadatas for the media assets. Defaults to None.
+            datapoints: (Sequence[Datapoint]): The datapoints to be set.
 
         Returns:
             RapidataOrderBuilder: The updated RapidataOrderBuilder instance.
         """
-        if not isinstance(assets, list):
-            raise TypeError("Media paths must be provided as a list of paths.")
-
-        for a in assets:
-            if not isinstance(a, (MediaAsset, TextAsset, MultiAsset)):
-                raise TypeError(
-                    "Media paths must be of type MediaAsset, TextAsset, or MultiAsset."
-                )
+        if not isinstance(datapoints, list):
+            raise TypeError(
+                "Datapoints must be provided as a list of Datapoint objects."
+            )
 
-        if multi_metadata:
-            for data in multi_metadata:
-                if not isinstance(data, list):
-                    raise TypeError("Metadata must be provided as a list of Metadata objects.")
-                for d in data:
-                    if not isinstance(d, Metadata):
-                        raise TypeError("Metadata must be of type Metadata.")
-
-        if multi_metadata and not len(multi_metadata) == len(assets):
-            raise ValueError("The number of assets must match the number of metadatas.")
-        
-        if multi_metadata and not all(len(data) == len(multi_metadata[0]) for data in multi_metadata):
-            raise ValueError("All metadatas must have the same length.")
-
-        self.__assets = assets
-        self.__multi_metadata = multi_metadata
+        self._datapoints = datapoints
         return self
 
-    def _settings(self, settings: Sequence[RapidataSetting]) -> "RapidataOrderBuilder":
+    def _set_settings(
+        self, settings: Sequence[RapidataSetting]
+    ) -> RapidataOrderBuilder:
         """
         Set the settings for the order.
 
@@ -245,15 +371,15 @@ class RapidataOrderBuilder:
 
         if not isinstance(settings, list):
             raise TypeError("Settings must be provided as a list of Setting objects.")
-        
+
         for s in settings:
             if not isinstance(s, RapidataSetting):
                 raise TypeError("The settings list must only contain Setting objects.")
 
-        self.__settings = settings
+        self._settings = settings
         return self
 
-    def _filters(self, filters: Sequence[RapidataFilter]) -> "RapidataOrderBuilder":
+    def _set_filters(self, filters: Sequence[RapidataFilter]) -> RapidataOrderBuilder:
         """
         Set the filters for the order, e.g., country, language, userscore, etc.
 
@@ -270,13 +396,15 @@ class RapidataOrderBuilder:
             if not isinstance(f, RapidataFilter):
                 raise TypeError("Filters must be of type Filter.")
 
-        if len(self.__user_filters) > 0:
+        if len(self._user_filters) > 0:
             managed_print("Overwriting existing user filters.")
 
-        self.__user_filters = filters
+        self._user_filters = filters
         return self
 
-    def _validation_set_id(self, validation_set_id: str) -> "RapidataOrderBuilder":
+    def _set_validation_set_id(
+        self, validation_set_id: str | None = None
+    ) -> RapidataOrderBuilder:
         """
         Set the validation set ID for the order.
 
@@ -286,28 +414,19 @@ class RapidataOrderBuilder:
         Returns:
             RapidataOrderBuilder: The updated RapidataOrderBuilder instance.
         """
-        if not isinstance(validation_set_id, str):
+        if validation_set_id is not None and not isinstance(validation_set_id, str):
             raise TypeError("Validation set ID must be of type str.")
-
-        self.__validation_set_id = validation_set_id
+        if validation_set_id is not None:
+            self._validation_set = (
+                self._validation_set_manager.get_validation_set_by_id(validation_set_id)
+            )
+            return self
+        self._validation_set = validation_set_id
         return self
 
-    def _rapids_per_bag(self, amount: int) -> "RapidataOrderBuilder":
-        """
-        Define the number of tasks a user sees in a single session.
-
-        Args:
-            amount (int): The number of tasks a user sees in a single session.
-
-        Returns:
-            RapidataOrderBuilder: The updated RapidataOrderBuilder instance.
-
-        Raises:
-            NotImplementedError: This method is not implemented yet.
-        """
-        raise NotImplementedError("Not implemented yet.")
-
-    def _selections(self, selections: Sequence[RapidataSelection]) -> "RapidataOrderBuilder":
+    def _set_selections(
+        self, selections: Sequence[RapidataSelection]
+    ) -> RapidataOrderBuilder:
         """
         Set the selections for the order.
 
@@ -326,10 +445,10 @@ class RapidataOrderBuilder:
             if not isinstance(selection, RapidataSelection):
                 raise TypeError("Selections must be of type Selection.")
 
-        self.__selections = selections  # type: ignore
+        self._selections = selections  # type: ignore
         return self
 
-    def _priority(self, priority: int) -> "RapidataOrderBuilder":
+    def _set_priority(self, priority: int | None = None) -> RapidataOrderBuilder:
         """
         Set the priority for the order.
 
@@ -339,8 +458,24 @@ class RapidataOrderBuilder:
         Returns:
             RapidataOrderBuilder: The updated RapidataOrderBuilder instance.
         """
-        if not isinstance(priority, int):
+        if priority is not None and not isinstance(priority, int):
             raise TypeError("Priority must be of type int.")
 
-        self.__priority = priority
+        self._priority = priority
+        return self
+
+    def _set_sticky_state(
+        self, sticky_state: StickyStateLiteral | None = None
+    ) -> RapidataOrderBuilder:
+        """
+        Set the sticky state for the order.
+        """
+        sticky_state_valid_values = get_args(StickyStateLiteral)
+
+        if sticky_state is not None and sticky_state not in sticky_state_valid_values:
+            raise ValueError(
+                f"Sticky state must be one of {sticky_state_valid_values} or None"
+            )
+
+        self._sticky_state_value = sticky_state
         return self