rapidata 2.21.5__py3-none-any.whl → 2.23.0__py3-none-any.whl

rapidata/rapidata_client/order/_rapidata_dataset.py

@@ -17,7 +17,7 @@ from tqdm import tqdm
 
 from pydantic import StrictStr
 from typing import cast, Sequence, Generator
-from logging import Logger
+from rapidata.rapidata_client.logging import logger
 import time
 import threading
 
@@ -33,11 +33,11 @@ class RapidataDataset:
         self.dataset_id = dataset_id
         self.openapi_service = openapi_service
         self.local_file_service = LocalFileService()
-        self._logger = Logger(__name__)
 
     def _add_texts(
         self,
         text_assets: list[TextAsset] | list[MultiAsset],
+        metadata_list: Sequence[Sequence[Metadata]] | None = None,
         max_workers: int = 10,
     ):
         for text_asset in text_assets:
@@ -46,17 +46,25 @@ class RapidataDataset:
                     isinstance(asset, TextAsset) for asset in text_asset.assets
                 ), "All assets in a MultiAsset must be of type TextAsset."
 
-        def upload_text_datapoint(text_asset: TextAsset | MultiAsset, index: int) -> None:
+        def upload_text_datapoint(text_asset: TextAsset | MultiAsset, metadata_per_datapoint: Sequence[Metadata] | None, index: int) -> None:
             if isinstance(text_asset, TextAsset):
                 texts = [text_asset.text]
             elif isinstance(text_asset, MultiAsset):
                 texts = [asset.text for asset in text_asset.assets if isinstance(asset, TextAsset)]
             else:
                 raise ValueError(f"Unsupported asset type: {type(text_asset)}")
+            
+            metadata = []
+            if metadata_per_datapoint:
+                for meta in metadata_per_datapoint:
+                    meta_model = meta.to_model() if meta else None
+                    if meta_model:
+                        metadata.append(CreateDatapointFromFilesModelMetadataInner(meta_model))
 
             model = CreateDatapointFromTextSourcesModel(
                 textSources=texts,
                 sortIndex=index,
+                metadata=metadata,
             )
 
             upload_response = self.openapi_service.dataset_api.dataset_dataset_id_datapoints_texts_post(dataset_id=self.dataset_id, create_datapoint_from_text_sources_model=model)
@@ -67,8 +75,8 @@ class RapidataDataset:
         total_uploads = len(text_assets)
         with ThreadPoolExecutor(max_workers=max_workers) as executor:
             futures = [
-                executor.submit(upload_text_datapoint, text_asset, index=i)
-                for i, text_asset in enumerate(text_assets)
+                executor.submit(upload_text_datapoint, text_asset, metadata, index=i)
+                for i, (text_asset, metadata) in enumerate(zip_longest(text_assets, metadata_list or []))
             ]
 
             with tqdm(total=total_uploads, desc="Uploading text datapoints") as pbar:
@@ -79,7 +87,7 @@ class RapidataDataset:
     def _process_single_upload(
         self,
         media_asset: MediaAsset | MultiAsset, 
-        meta: Metadata | None, 
+        meta_list: Sequence[Metadata] | None, 
         index: int,
     ) -> tuple[list[str], list[str]]:
         """
@@ -87,7 +95,7 @@ class RapidataDataset:
         
         Args:
             media_asset: MediaAsset or MultiAsset to upload
-            meta: Optional metadata for the asset
+            meta_list: Optional sequence of metadata for the asset
             index: Sort index for the upload
             
         Returns:
@@ -105,16 +113,20 @@ class RapidataDataset:
                 identifiers_to_track = [identifier] if identifier else []
             elif isinstance(media_asset, MultiAsset):
                 assets = cast(list[MediaAsset], media_asset.assets)
-                identifiers_to_track: list[str] = [
+                identifiers_to_track = [
                     (asset._url if asset._url else cast(str, asset.path)) 
                     for asset in assets
                 ]
             else:
                 raise ValueError(f"Unsupported asset type: {type(media_asset)}")
 
-            meta_model = meta.to_model() if meta else None
-
-            metadata = [CreateDatapointFromFilesModelMetadataInner(meta_model)] if meta_model else []
+            # Convert multiple metadata to models
+            metadata = []
+            if meta_list:
+                for meta in meta_list:
+                    meta_model = meta.to_model() if meta else None
+                    if meta_model:
+                        metadata.append(CreateDatapointFromFilesModelMetadataInner(meta_model))
 
             local_paths: bool = assets[0].is_local()
             files: list[StrictStr] = []
@@ -144,14 +156,14 @@ class RapidataDataset:
                 
             if upload_response.errors:
                 error_msg = f"Error uploading datapoint: {upload_response.errors}"
-                self._logger.error(error_msg)
+                logger.error(error_msg)
                 local_failed.extend(identifiers_to_track)
                 raise ValueError(error_msg)
 
             local_successful.extend(identifiers_to_track)
 
         except Exception as e:
-            self._logger.error(f"\nUpload failed for {identifiers_to_track}: {str(e)}") # \n to avoid same line as tqdm
+            logger.error(f"\nUpload failed for {identifiers_to_track}: {str(e)}") # \n to avoid same line as tqdm
             local_failed.extend(identifiers_to_track)
 
         return local_successful, local_failed
@@ -233,11 +245,11 @@ class RapidataDataset:
                                     # If we're not at 100% but it's been a while with no progress
                                     if stall_count > 5:
                                         # We've polled several times with no progress, assume we're done
-                                        self._logger.warning(f"\nProgress seems stalled at {total_completed}/{total_uploads}. Please try again.")
+                                        logger.warning(f"\nProgress seems stalled at {total_completed}/{total_uploads}. Please try again.")
                                         break
                                 
                         except Exception as e:
-                            self._logger.error(f"\nError checking progress: {str(e)}")
+                            logger.error(f"\nError checking progress: {str(e)}")
                             stall_count += 1
                             
                             if stall_count > 10:  # Too many consecutive errors
@@ -248,7 +260,7 @@ class RapidataDataset:
                         time.sleep(progress_poll_interval)
                 
             except Exception as e:
-                self._logger.error(f"Progress tracking thread error: {str(e)}")
+                logger.error(f"Progress tracking thread error: {str(e)}")
                 progress_error_event.set()
                 
         # Create and return the thread
@@ -259,7 +271,7 @@ class RapidataDataset:
     def _process_uploads_in_chunks(
         self,
         media_paths: list[MediaAsset] | list[MultiAsset],
-        metadata: Sequence[Metadata] | None,
+        multi_metadata: Sequence[Sequence[Metadata]] | None,
         max_workers: int,
         chunk_size: int,
         stop_progress_tracking: threading.Event,
@@ -270,7 +282,7 @@ class RapidataDataset:
         
         Args:
             media_paths: List of assets to upload
-            metadata: Optional sequence of metadata
+            multi_metadata: Optional sequence of sequences of metadata
             max_workers: Maximum number of concurrent workers
             chunk_size: Number of items to process in each batch
             stop_progress_tracking: Event to signal progress tracking to stop
@@ -286,16 +298,16 @@ class RapidataDataset:
             with ThreadPoolExecutor(max_workers=max_workers) as executor:
                 # Process uploads in chunks to avoid overwhelming the system
                 for chunk_idx, chunk in enumerate(chunk_list(media_paths, chunk_size)):
-                    chunk_metadata = metadata[chunk_idx * chunk_size:(chunk_idx + 1) * chunk_size] if metadata else None
+                    chunk_multi_metadata = multi_metadata[chunk_idx * chunk_size:(chunk_idx + 1) * chunk_size] if multi_metadata else None
                     
                     futures = [
                         executor.submit(
                             self._process_single_upload, 
                             media_asset, 
-                            meta, 
+                            meta_list, 
                             index=(chunk_idx * chunk_size + i)
                         )
-                        for i, (media_asset, meta) in enumerate(zip_longest(chunk, chunk_metadata or []))
+                        for i, (media_asset, meta_list) in enumerate(zip_longest(chunk, chunk_multi_metadata or []))
                     ]
                     
                     # Wait for this chunk to complete before starting the next one
@@ -308,7 +320,7 @@ class RapidataDataset:
                             successful_uploads.extend(chunk_successful)
                             failed_uploads.extend(chunk_failed)
                         except Exception as e:
-                            self._logger.error(f"Future execution failed: {str(e)}")
+                            logger.error(f"Future execution failed: {str(e)}")
         finally:
             # Signal to the progress tracking thread that all uploads have been submitted
             stop_progress_tracking.set()
@@ -349,19 +361,18 @@ class RapidataDataset:
             
             success_rate = (total_ready / total_uploads * 100) if total_uploads > 0 else 0
             
-            self._logger.info(f"Upload complete: {total_ready} ready, {total_uploads-total_ready} failed ({success_rate:.1f}% success rate)")
-            print(f"Upload complete, {total_ready} ready, {total_uploads-total_ready} failed ({success_rate:.1f}% success rate)")
+            logger.info(f"Upload complete: {total_ready} ready, {total_uploads-total_ready} failed ({success_rate:.1f}% success rate)")
         except Exception as e:
-            self._logger.error(f"Error getting final progress: {str(e)}")
-            self._logger.info(f"Upload summary from local tracking: {len(successful_uploads)} succeeded, {len(failed_uploads)} failed")
+            logger.error(f"Error getting final progress: {str(e)}")
+            logger.info(f"Upload summary from local tracking: {len(successful_uploads)} succeeded, {len(failed_uploads)} failed")
 
         if failed_uploads:
-            print(f"Failed uploads: {failed_uploads}")
+            logger.error(f"Failed uploads: {failed_uploads}")
 
     def _add_media_from_paths(
         self,
         media_paths: list[MediaAsset] | list[MultiAsset],
-        metadata: Sequence[Metadata] | None = None,
+        multi_metadata: Sequence[Sequence[Metadata]] | None = None,
         max_workers: int = 5,
         chunk_size: int = 50,
         progress_poll_interval: float = 0.5,
@@ -371,7 +382,7 @@ class RapidataDataset:
         
         Args:
             media_paths: List of MediaAsset or MultiAsset objects to upload
-            metadata: Optional sequence of metadata matching media_paths length
+            multi_metadata: Optional sequence of sequences of metadata matching media_paths length
             max_workers: Maximum number of concurrent upload workers
             chunk_size: Number of items to process in each batch
             progress_poll_interval: Time in seconds between progress checks
@@ -380,10 +391,14 @@ class RapidataDataset:
             tuple[list[str], list[str]]: Lists of successful and failed URLs
             
         Raises:
-            ValueError: If metadata length doesn't match media_paths length
+            ValueError: If multi_metadata lengths don't match media_paths length
         """
-        if metadata is not None and len(metadata) != len(media_paths):
-            raise ValueError("metadata must be None or have the same length as media_paths")
+
+        if multi_metadata and not len(multi_metadata) == len(media_paths):
+            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.")
         
         # Setup tracking variables
         total_uploads = len(media_paths)
@@ -405,7 +420,7 @@ class RapidataDataset:
         try:
             successful_uploads, failed_uploads = self._process_uploads_in_chunks(
                 media_paths,
-                metadata,
+                multi_metadata,
                 max_workers,
                 chunk_size,
                 stop_progress_tracking,

rapidata/rapidata_client/order/_rapidata_order_builder.py

@@ -29,6 +29,8 @@ from rapidata.rapidata_client.assets import MediaAsset, TextAsset, MultiAsset, B
 
 from typing import Optional, cast, Sequence
 
+from rapidata.rapidata_client.logging import logger, managed_print
+
 
 class RapidataOrderBuilder:
     """Builder object for creating Rapidata orders.
@@ -51,7 +53,7 @@ class RapidataOrderBuilder:
         self.__dataset: Optional[RapidataDataset]
         self.__workflow: Workflow | None = None
         self.__referee: Referee | None = None
-        self.__metadata: Sequence[Metadata] | 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] = []
@@ -73,7 +75,7 @@ class RapidataOrderBuilder:
             raise ValueError("You must provide a workflow to create an order.")
 
         if self.__referee is None:
-            print("No referee provided, using default NaiveReferee.")
+            managed_print("No referee provided, using default NaiveReferee.")
             self.__referee = NaiveReferee()
 
         return CreateOrderModel(
@@ -113,6 +115,7 @@ class RapidataOrderBuilder:
             RapidataOrder: The created RapidataOrder instance.
         """
         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
@@ -125,12 +128,17 @@ class RapidataOrderBuilder:
         )
 
         self.order_id = str(result.order_id)
+        logger.debug(f"Order created with ID: {self.order_id}")
 
         self.__dataset = (
             RapidataDataset(result.dataset_id, self.__openapi_service)
             if result.dataset_id
             else None
         )
+        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,
@@ -138,15 +146,18 @@ class RapidataOrderBuilder:
             name=self._name,
         )
 
+        logger.debug(f"Order created: {order}")
+        logger.debug("Adding media to the order.")
+
         if all(isinstance(item, MediaAsset) for item in self.__assets) and self.__dataset:
             assets = cast(list[MediaAsset], self.__assets)
-            self.__dataset._add_media_from_paths(assets, self.__metadata, max_upload_workers)
+            self.__dataset._add_media_from_paths(assets, self.__multi_metadata, max_upload_workers)
 
         elif (
             all(isinstance(item, TextAsset) for item in self.__assets) and self.__dataset
         ):
             assets = cast(list[TextAsset], self.__assets)
-            self.__dataset._add_texts(assets)
+            self.__dataset._add_texts(assets, self.__multi_metadata)
 
         elif (
             all(isinstance(item, MultiAsset) for item in self.__assets) and self.__dataset
@@ -167,11 +178,11 @@ class RapidataOrderBuilder:
             # Process based on the asset type
             if issubclass(first_asset_type, MediaAsset):
                 self.__dataset._add_media_from_paths(
-                    multi_assets, self.__metadata, max_upload_workers
+                    multi_assets, self.__multi_metadata, max_upload_workers
                 )
 
             elif issubclass(first_asset_type, TextAsset):
-                self.__dataset._add_texts(multi_assets)
+                self.__dataset._add_texts(multi_assets, self.__multi_metadata)
 
             else:
                 raise ValueError(
@@ -183,6 +194,8 @@ class RapidataOrderBuilder:
                 "Media paths must all be of the same type: MediaAsset, TextAsset, or MultiAsset."
             )
         
+        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)
 
         return order
@@ -221,35 +234,44 @@ class RapidataOrderBuilder:
 
     def _media(
         self,
-        asset: Sequence[BaseAsset],
-        metadata: Sequence[Metadata] | None = None,
+        assets: Sequence[BaseAsset],
+        multi_metadata: Sequence[Sequence[Metadata]] | None = None,
     ) -> "RapidataOrderBuilder":
         """
         Set the media assets for the order.
 
         Args:
-            asset: (list[MediaAsset] | list[TextAsset] | list[MultiAsset]): The paths of the media assets to be set.
-            metadata: (list[Metadata] | None, optional): Metadata for the media assets. Defaults to None.
+            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.
 
         Returns:
             RapidataOrderBuilder: The updated RapidataOrderBuilder instance.
         """
-        if not isinstance(asset, list):
+        if not isinstance(assets, list):
             raise TypeError("Media paths must be provided as a list of paths.")
 
-        for a in asset:
+        for a in assets:
             if not isinstance(a, (MediaAsset, TextAsset, MultiAsset)):
                 raise TypeError(
                     "Media paths must be of type MediaAsset, TextAsset, or MultiAsset."
                 )
 
-        if metadata:
-            for data in metadata:
-                if not isinstance(data, Metadata):
-                    raise TypeError("Metadata must be of type Metadata.")
+        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 = asset
-        self.__metadata = metadata
+        self.__assets = assets
+        self.__multi_metadata = multi_metadata
         return self
 
     def _settings(self, settings: Sequence[RapidataSetting]) -> "RapidataOrderBuilder":
@@ -291,7 +313,7 @@ class RapidataOrderBuilder:
                 raise TypeError("Filters must be of type Filter.")
 
         if len(self.__user_filters) > 0:
-            print("Overwriting existing user filters.")
+            managed_print("Overwriting existing user filters.")
 
         self.__user_filters = filters
         return self

rapidata/rapidata_client/order/rapidata_order.py

@@ -17,6 +17,7 @@ from rapidata.api_client.models.preliminary_download_model import PreliminaryDow
 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
 
 
 class RapidataOrder:
@@ -47,23 +48,29 @@ class RapidataOrder:
         self._max_retries = 10  
         self._retry_delay = 2
         self.order_details_page = f"https://app.{self.__openapi_service.environment}/order/detail/{self.order_id}"
+        logger.debug("RapidataOrder initialized")
 
-    def run(self, print_link: bool = True) -> "RapidataOrder":
+    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.order_id)
-        if print_link:
-            print(f"Order '{self.name}' is now viewable under: {self.order_details_page}")
+        logger.debug(f"Order '{self}' has been started.")
+        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_pause_post(self.order_id)
-        print(f"Order '{self}' has been paused.")
+        logger.debug(f"Order '{self}' has been paused.")
+        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_resume_post(self.order_id)
-        print(f"Order '{self}' has been unpaused.")
+        logger.debug(f"Order '{self}' has been unpaused.")
+        managed_print(f"Order '{self}' has been unpaused.")
 
     def get_status(self) -> str:
         """
@@ -95,12 +102,12 @@ class RapidataOrder:
             raise Exception("Order has not been started yet. Please start it first.")
         
         while self.get_status() == OrderState.SUBMITTED:
-            print(f"Order '{self.name}' 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.name}' is in manual review. It might take some time to start. "
+                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."
             )
@@ -145,12 +152,12 @@ class RapidataOrder:
                 Note that preliminary results are not final and may not contain all the datapoints & responses. Only the onese that are already available.
                 This will throw an exception if there are no responses available yet.
         """
-
+        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]:
-            print("Order is already completed. Returning final results.")
+            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)
@@ -167,10 +174,11 @@ class RapidataOrder:
         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="%/:=&?~#+!$,;'@()*[]")
-            print(Fore.RED + f'Please open this URL in your browser: "{encoded_url}"' + Fore.RESET)
+            managed_print(Fore.RED + f'Please open this URL in your browser: "{encoded_url}"' + Fore.RESET)
         
     def preview(self) -> None:
         """
@@ -178,13 +186,14 @@ class RapidataOrder:
         
         Raises:
             Exception: If the order is not in processing state.
-        """        
+        """
+        logger.info("Opening order preview in browser...")
         campaign_id = self.__get_campaign_id()
         auth_url = f"https://app.{self.__openapi_service.environment}/order/detail/{self.order_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="%/:=&?~#+!$,;'@()*[]")
-            print(Fore.RED + f'Please open this URL in your browser: "{encoded_url}"' + Fore.RESET)
+            managed_print(Fore.RED + f'Please open this URL in your browser: "{encoded_url}"' + Fore.RESET)
 
     def __get_pipeline_id(self) -> str:
         """Internal method to fetch and cache the pipeline ID."""
@@ -249,7 +258,7 @@ class RapidataOrder:
             raise Exception(f"Failed to get preliminary results: {str(e)}") from e
 
     def __str__(self) -> str:
-        return f"name: '{self.name}' order id: {self.order_id}"
+        return f"RapidataOrder(name='{self.name}', order_id='{self.order_id}')"
 
     def __repr__(self) -> str:
         return f"RapidataOrder(name='{self.name}', order_id='{self.order_id}')"