rapidata 2.41.2__py3-none-any.whl → 2.42.0__py3-none-any.whl

rapidata/rapidata_client/datapoints/_datapoint_uploader.py

@@ -0,0 +1,95 @@
+from rapidata.api_client.models.text_asset_input import TextAssetInput
+from rapidata.rapidata_client.datapoints._datapoint import Datapoint
+from rapidata.service.openapi_service import OpenAPIService
+from rapidata.api_client.models.multi_asset_input_assets_inner import (
+    MultiAssetInput,
+    MultiAssetInputAssetsInner,
+)
+from rapidata.api_client.models.create_datapoint_model import CreateDatapointModel
+from rapidata.api_client.models.create_datapoint_model_asset import (
+    CreateDatapointModelAsset,
+)
+from rapidata.api_client.models.create_datapoint_result import CreateDatapointResult
+from rapidata.api_client.models.create_datapoint_from_files_model_metadata_inner import (
+    CreateDatapointFromFilesModelMetadataInner,
+)
+from rapidata.api_client.models.existing_asset_input import ExistingAssetInput
+from rapidata.rapidata_client.datapoints._asset_uploader import AssetUploader
+from rapidata.rapidata_client.datapoints.metadata import (
+    PromptMetadata,
+    MediaAssetMetadata,
+    PrivateTextMetadata,
+    SelectWordsMetadata,
+    Metadata,
+)
+
+
+class DatapointUploader:
+    def __init__(self, openapi_service: OpenAPIService):
+        self.openapi_service = openapi_service
+        self.asset_uploader = AssetUploader(openapi_service)
+
+    def upload_datapoint(
+        self, datapoint: Datapoint, dataset_id: str, index: int
+    ) -> CreateDatapointResult:
+        metadata = self._get_metadata(datapoint)
+
+        uploaded_asset = (
+            self._handle_media_datapoint(datapoint)
+            if datapoint.data_type == "media"
+            else self._handle_text_datapoint(datapoint)
+        )
+        return self.openapi_service.dataset_api.dataset_dataset_id_datapoint_post(
+            dataset_id=dataset_id,
+            create_datapoint_model=CreateDatapointModel(
+                asset=uploaded_asset,
+                metadata=metadata,
+                sortIndex=index,
+            ),
+        )
+
+    def _get_metadata(
+        self, datapoint: Datapoint
+    ) -> list[CreateDatapointFromFilesModelMetadataInner]:
+        datapoint_metadata: list[Metadata] = []
+        if datapoint.context:
+            datapoint_metadata.append(PromptMetadata(prompt=datapoint.context))
+        if datapoint.sentence:
+            datapoint_metadata.append(
+                SelectWordsMetadata(select_words=datapoint.sentence)
+            )
+        if datapoint.media_context:
+            datapoint_metadata.append(
+                MediaAssetMetadata(
+                    internal_file_name=self.asset_uploader.upload_asset(
+                        datapoint.media_context
+                    )
+                )
+            )
+        if datapoint.private_note:
+            datapoint_metadata.append(PrivateTextMetadata(text=datapoint.private_note))
+
+        metadata = [
+            CreateDatapointFromFilesModelMetadataInner(
+                actual_instance=metadata.to_model()
+            )
+            for metadata in datapoint_metadata
+        ]
+
+        return metadata
+
+    def _handle_text_datapoint(self, datapoint: Datapoint) -> CreateDatapointModelAsset:
+        return CreateDatapointModelAsset(
+            actual_instance=self.asset_uploader.get_uploaded_text_input(
+                datapoint.asset
+            ),
+        )
+
+    def _handle_media_datapoint(
+        self, datapoint: Datapoint
+    ) -> CreateDatapointModelAsset:
+        return CreateDatapointModelAsset(
+            actual_instance=self.asset_uploader.get_uploaded_asset_input(
+                datapoint.asset
+            ),
+        )

rapidata/rapidata_client/datapoints/assets/__init__.py

@@ -1,11 +0,0 @@
-"""Assets Package
-
-This package provides classes for different types of assets, including MediaAsset, TextAsset, and MultiAsset.
-"""
-
-from ._base_asset import BaseAsset
-from ._media_asset import MediaAsset
-from ._text_asset import TextAsset
-from ._multi_asset import MultiAsset
-from .data_type_enum import RapidataDataTypes
-from ._sessions import SessionManager

rapidata/rapidata_client/datapoints/metadata/_media_asset_metadata.py

@@ -1,23 +1,26 @@
 from rapidata.api_client.models.prompt_asset_metadata_input import (
     PromptAssetMetadataInput,
 )
-from rapidata.api_client.models.url_asset_input import UrlAssetInput
 from rapidata.rapidata_client.datapoints.metadata._base_metadata import Metadata
-from rapidata.api_client.models.prompt_asset_metadata_input_asset import (
-    PromptAssetMetadataInputAsset,
+from rapidata.api_client.models.multi_asset_input_assets_inner import (
+    ExistingAssetInput,
+    MultiAssetInputAssetsInner,
 )
 
 
 class MediaAssetMetadata(Metadata):
 
-    def __init__(self, url: str):
+    def __init__(self, internal_file_name: str):
         super().__init__()
-        self._url = url
+        self._internal_file_name = internal_file_name
 
     def to_model(self):
         return PromptAssetMetadataInput(
             _t="PromptAssetMetadataInput",
-            asset=PromptAssetMetadataInputAsset(
-                actual_instance=UrlAssetInput(_t="UrlAssetInput", url=self._url)
+            asset=MultiAssetInputAssetsInner(
+                actual_instance=ExistingAssetInput(
+                    _t="ExistingAssetInput",
+                    name=self._internal_file_name,
+                ),
             ),
         )

rapidata/rapidata_client/demographic/demographic_manager.py

@@ -1,34 +1,47 @@
-from rapidata.service.openapi_service import OpenAPIService
-from rapidata.rapidata_client.datapoints.assets import MediaAsset
-from rapidata.api_client.models.create_demographic_rapid_model import (
-    CreateDemographicRapidModel,
+from argparse import Action
+from rapidata.api_client import ExistingAssetInput
+from rapidata.api_client.models.create_demographic_rapid_model_asset import (
+    CreateDemographicRapidModelAsset,
 )
+from rapidata.service.openapi_service import OpenAPIService
 from rapidata.api_client.models.classify_payload import ClassifyPayload
 from rapidata.rapidata_client.config import logger
+from rapidata.api_client.models.create_demographic_rapid_model_new import (
+    CreateDemographicRapidModelNew,
+)
+from rapidata.rapidata_client.datapoints._asset_uploader import AssetUploader
 
 
 class DemographicManager:
     def __init__(self, openapi_service: OpenAPIService):
         self._openapi_service = openapi_service
+        self._asset_uploader = AssetUploader(openapi_service)
         logger.debug("DemographicManager initialized")
 
     def create_demographic_rapid(
         self, instruction: str, answer_options: list[str], datapoint: str, key: str
     ):
 
-        media = MediaAsset(path=datapoint)
-        model = CreateDemographicRapidModel(
+        model = CreateDemographicRapidModelNew(
             key=key,
             payload=ClassifyPayload(
                 _t="ClassifyPayload",
                 possibleCategories=answer_options,
                 title=instruction,
             ),
+            asset=CreateDemographicRapidModelAsset(
+                actual_instance=ExistingAssetInput(
+                    _t="ExistingAssetInput",
+                    name=self._asset_uploader.upload_asset(datapoint),
+                ),
+            ),
         )
 
-        self._openapi_service.rapid_api.rapid_demographic_post(
-            model=model, file=[media.to_file()]
+        result = self._openapi_service.rapid_api.rapid_demographic_new_post(
+            create_demographic_rapid_model_new=model
         )
+        logger.info(f"Demographic Rapid created: {result.rapid_id}")
+        return result.rapid_id
 
     def __str__(self) -> str:
         return "DemographicManager"

rapidata/rapidata_client/exceptions/failed_upload_exception.py

@@ -1,14 +1,3 @@
-from typing import cast
-from rapidata.api_client.models.file_asset_model import FileAssetModel
-from rapidata.api_client.models.get_failed_datapoints_result import (
-    GetFailedDatapointsResult,
-)
-from rapidata.api_client.models.multi_asset_model import MultiAssetModel
-from rapidata.api_client.models.original_filename_metadata_model import (
-    OriginalFilenameMetadataModel,
-)
-from rapidata.api_client.models.source_url_metadata_model import SourceUrlMetadataModel
-from rapidata.rapidata_client.datapoints.assets import MediaAsset, MultiAsset
 from rapidata.rapidata_client.datapoints._datapoint import Datapoint
 from rapidata.rapidata_client.order.dataset._rapidata_dataset import RapidataDataset
 from rapidata.rapidata_client.order.rapidata_order import RapidataOrder
@@ -29,54 +18,3 @@ class FailedUploadException(Exception):
 
     def __str__(self) -> str:
         return f"Failed to upload {self.failed_uploads}"
-
-
-def _parse_failed_uploads(failed_uploads: GetFailedDatapointsResult) -> list[Datapoint]:
-    failed_datapoints = failed_uploads.datapoints
-    if not failed_datapoints:
-        return []
-    if isinstance(failed_datapoints[0].asset.actual_instance, FileAssetModel):
-        failed_assets = [
-            MediaAsset(
-                __get_asset_name(cast(FileAssetModel, datapoint.asset.actual_instance))
-            )
-            for datapoint in failed_datapoints
-        ]
-    elif isinstance(failed_datapoints[0].asset.actual_instance, MultiAssetModel):
-        failed_assets = []
-        backend_assets = [
-            cast(MultiAssetModel, failed_upload.asset.actual_instance).assets
-            for failed_upload in failed_datapoints
-        ]
-        for assets in backend_assets:
-            failed_assets.append(
-                MultiAsset(
-                    [
-                        MediaAsset(
-                            __get_asset_name(
-                                cast(FileAssetModel, asset.actual_instance)
-                            )
-                        )
-                        for asset in assets
-                        if isinstance(asset.actual_instance, FileAssetModel)
-                    ]
-                )
-            )
-    else:
-        raise ValueError(
-            f"Unsupported asset type: {type(failed_datapoints[0].asset.actual_instance)}"
-        )
-
-    return [Datapoint(asset=asset) for asset in failed_assets]
-
-
-def __get_asset_name(failed_datapoint: FileAssetModel) -> str:
-    metadata = failed_datapoint.metadata
-    if "sourceUrl" in metadata:
-        return cast(SourceUrlMetadataModel, metadata["sourceUrl"].actual_instance).url
-    elif "originalFilename" in metadata:
-        return cast(
-            OriginalFilenameMetadataModel, metadata["originalFilename"].actual_instance
-        ).original_filename
-    else:
-        return ""

rapidata/rapidata_client/order/_rapidata_order_builder.py

@@ -20,7 +20,6 @@ from rapidata.api_client.models.sticky_state import StickyState
 from rapidata.rapidata_client.datapoints._datapoint import Datapoint
 from rapidata.rapidata_client.exceptions.failed_upload_exception import (
     FailedUploadException,
-    _parse_failed_uploads,
 )
 from rapidata.rapidata_client.filter import RapidataFilter
 from rapidata.rapidata_client.config import (
@@ -280,15 +279,6 @@ class RapidataOrderBuilder:
         try:
             self.__openapi_service.order_api.order_order_id_preview_post(self.order_id)
         except Exception:
-            failed_uploads = _parse_failed_uploads(
-                self.__openapi_service.dataset_api.dataset_dataset_id_datapoints_failed_get(
-                    self.__dataset.id
-                )
-            )
-            logger.error(
-                "Internal download error for datapoints: %s\nWARNING: Failed Datapoints in error do not contain metadata.",
-                failed_uploads,
-            )
             raise FailedUploadException(self.__dataset, order, failed_uploads)
         return order
 

rapidata/rapidata_client/order/dataset/_rapidata_dataset.py

@@ -1,6 +1,4 @@
 from rapidata.rapidata_client.datapoints._datapoint import Datapoint
-from rapidata.rapidata_client.datapoints.assets import TextAsset, MediaAsset
-from rapidata.service import LocalFileService
 from rapidata.service.openapi_service import OpenAPIService
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from tqdm import tqdm
@@ -8,12 +6,11 @@ from tqdm import tqdm
 from typing import Generator
 from rapidata.rapidata_client.config import logger
 import time
-import threading
 from rapidata.rapidata_client.api.rapidata_api_client import (
     suppress_rapidata_error_logging,
 )
 from rapidata.rapidata_client.config.rapidata_config import rapidata_config
-from rapidata.rapidata_client.order.dataset._progress_tracker import ProgressTracker
+from rapidata.rapidata_client.datapoints._datapoint_uploader import DatapointUploader
 
 # Add OpenTelemetry context imports for thread propagation
 from opentelemetry import context as otel_context
@@ -28,81 +25,80 @@ class RapidataDataset:
     def __init__(self, dataset_id: str, openapi_service: OpenAPIService):
         self.id = dataset_id
         self.openapi_service = openapi_service
-        self.local_file_service = LocalFileService()
+        self.datapoint_uploader = DatapointUploader(openapi_service)
 
     def add_datapoints(
         self,
         datapoints: list[Datapoint],
     ) -> tuple[list[Datapoint], list[Datapoint]]:
-        if not datapoints:
-            return [], []
-
-        effective_asset_type = datapoints[0]._get_effective_asset_type()
-
-        logger.debug(f"Config for datapoint upload: {rapidata_config}")
-
-        if issubclass(effective_asset_type, MediaAsset):
-            return self._add_media_from_paths(
-                datapoints,
-            )
-        elif issubclass(effective_asset_type, TextAsset):
-            return self._add_texts(datapoints)
-        else:
-            raise ValueError(f"Unsupported asset type: {effective_asset_type}")
-
-    def _add_texts(
-        self, datapoints: list[Datapoint]
-    ) -> tuple[list[Datapoint], list[Datapoint]]:
-
-        def upload_text_datapoint(datapoint: Datapoint, index: int) -> Datapoint:
-            model = datapoint.create_text_upload_model(index)
-
-            self.openapi_service.dataset_api.dataset_dataset_id_datapoints_texts_post(
-                dataset_id=self.id, create_datapoint_from_text_sources_model=model
-            )
-            return datapoint
+        """
+        Process uploads in chunks with a ThreadPoolExecutor.
 
-        def upload_with_context(
-            context: otel_context.Context, datapoint: Datapoint, index: int
-        ) -> Datapoint:
-            """Wrapper function that runs upload_text_datapoint with the provided context."""
-            token = otel_context.attach(context)
-            try:
-                return upload_text_datapoint(datapoint, index)
-            finally:
-                otel_context.detach(token)
+        Args:
+            media_paths: List of assets to upload
+            multi_metadata: Optional sequence of sequences of metadata
+            chunk_size: Number of items to process in each batch
 
+        Returns:
+            tuple[list[str], list[str]]: Lists of successful and failed uploads
+        """
         successful_uploads: list[Datapoint] = []
         failed_uploads: list[Datapoint] = []
 
-        # Capture the current OpenTelemetry context before creating threads
-        current_context = otel_context.get_current()
-
-        total_uploads = len(datapoints)
-        with ThreadPoolExecutor(
-            max_workers=rapidata_config.upload.maxWorkers
-        ) as executor:
-            future_to_datapoint = {
-                executor.submit(
-                    upload_with_context, current_context, datapoint, i
-                ): datapoint
-                for i, datapoint in enumerate(datapoints)
-            }
+        with tqdm(
+            total=len(datapoints),
+            desc="Uploading datapoints",
+            disable=rapidata_config.logging.silent_mode,
+        ) as progress_bar:
+
+            def process_upload_with_context(
+                context: otel_context.Context, datapoint: Datapoint, index: int
+            ) -> tuple[list[Datapoint], list[Datapoint]]:
+                """Wrapper function that runs _process_single_upload with the provided context."""
+                token = otel_context.attach(context)
+                try:
+                    return self._process_single_upload(datapoint, index)
+                finally:
+                    otel_context.detach(token)
+
+            # Capture the current OpenTelemetry context before creating threads
+            current_context = otel_context.get_current()
+
+            with ThreadPoolExecutor(
+                max_workers=rapidata_config.upload.maxWorkers
+            ) as executor:
+                # Process uploads in chunks to avoid overwhelming the system
+                for chunk_idx, chunk in enumerate(
+                    chunk_list(datapoints, rapidata_config.upload.chunkSize)
+                ):
+                    futures = [
+                        executor.submit(
+                            process_upload_with_context,
+                            current_context,
+                            datapoint,
+                            chunk_idx * rapidata_config.upload.chunkSize + i,
+                        )
+                        for i, datapoint in enumerate(chunk)
+                    ]
+
+                    # Wait for this chunk to complete before starting the next one
+                    for future in as_completed(futures):
+                        try:
+                            chunk_successful, chunk_failed = future.result()
+                            successful_uploads.extend(chunk_successful)
+                            failed_uploads.extend(chunk_failed)
+                            progress_bar.update(
+                                len(chunk_successful) + len(chunk_failed)
+                            )
+                        except Exception as e:
+                            logger.error("Future execution failed: %s", str(e))
 
-            with tqdm(
-                total=total_uploads,
-                desc="Uploading text datapoints",
-                disable=rapidata_config.logging.silent_mode,
-            ) as pbar:
-                for future in as_completed(future_to_datapoint.keys()):
-                    datapoint = future_to_datapoint[future]
-                    try:
-                        result = future.result()
-                        pbar.update(1)
-                        successful_uploads.append(result)
-                    except Exception as e:
-                        failed_uploads.append(datapoint)
-                        logger.error("Upload failed for %s: %s", datapoint, str(e))
+        if failed_uploads:
+            logger.error(
+                "Upload failed for %s datapoints: %s",
+                len(failed_uploads),
+                failed_uploads,
+            )
 
         return successful_uploads, failed_uploads
 
@@ -128,21 +124,14 @@ class RapidataDataset:
         local_successful: list[Datapoint] = []
         local_failed: list[Datapoint] = []
 
-        metadata = datapoint.get_prepared_metadata()
-
-        local_paths = datapoint.get_local_file_paths()
-        urls = datapoint.get_urls()
-
         last_exception = None
         for attempt in range(rapidata_config.upload.maxRetries):
             try:
                 with suppress_rapidata_error_logging():
-                    self.openapi_service.dataset_api.dataset_dataset_id_datapoints_post(
+                    self.datapoint_uploader.upload_datapoint(
                         dataset_id=self.id,
-                        file=local_paths,
-                        url=urls,
-                        metadata=metadata,
-                        sort_index=index,
+                        datapoint=datapoint,
+                        index=index,
                     )
 
                 local_successful.append(datapoint)
@@ -170,115 +159,6 @@ class RapidataDataset:
 
         return local_successful, local_failed
 
-    def _process_uploads_in_chunks(
-        self,
-        datapoints: list[Datapoint],
-    ) -> tuple[list[Datapoint], list[Datapoint]]:
-        """
-        Process uploads in chunks with a ThreadPoolExecutor.
-
-        Args:
-            media_paths: List of assets to upload
-            multi_metadata: Optional sequence of sequences of metadata
-            chunk_size: Number of items to process in each batch
-
-        Returns:
-            tuple[list[str], list[str]]: Lists of successful and failed uploads
-        """
-        successful_uploads: list[Datapoint] = []
-        failed_uploads: list[Datapoint] = []
-
-        def process_upload_with_context(
-            context: otel_context.Context, datapoint: Datapoint, index: int
-        ) -> tuple[list[Datapoint], list[Datapoint]]:
-            """Wrapper function that runs _process_single_upload with the provided context."""
-            token = otel_context.attach(context)
-            try:
-                return self._process_single_upload(datapoint, index)
-            finally:
-                otel_context.detach(token)
-
-        # Capture the current OpenTelemetry context before creating threads
-        current_context = otel_context.get_current()
-
-        with ThreadPoolExecutor(
-            max_workers=rapidata_config.upload.maxWorkers
-        ) as executor:
-            # Process uploads in chunks to avoid overwhelming the system
-            for chunk_idx, chunk in enumerate(
-                chunk_list(datapoints, rapidata_config.upload.chunkSize)
-            ):
-                futures = [
-                    executor.submit(
-                        process_upload_with_context,
-                        current_context,
-                        datapoint,
-                        chunk_idx * rapidata_config.upload.chunkSize + i,
-                    )
-                    for i, datapoint in enumerate(chunk)
-                ]
-
-                # Wait for this chunk to complete before starting the next one
-                for future in as_completed(futures):
-                    try:
-                        chunk_successful, chunk_failed = future.result()
-                        successful_uploads.extend(chunk_successful)
-                        failed_uploads.extend(chunk_failed)
-                    except Exception as e:
-                        logger.error("Future execution failed: %s", str(e))
-
-        return successful_uploads, failed_uploads
-
-    def _add_media_from_paths(
-        self,
-        datapoints: list[Datapoint],
-        progress_poll_interval: float = 0.5,
-    ) -> tuple[list[Datapoint], list[Datapoint]]:
-        """
-        Upload media paths in chunks with managed resources.
-
-        Args:
-            datapoints: List of Datapoint objects to upload
-            chunk_size: Number of items to process in each batch
-            progress_poll_interval: Time in seconds between progress checks
-        Returns:
-            tuple[list[Datapoint], list[Datapoint]]: Lists of successful and failed datapoints
-
-        Raises:
-            ValueError: If multi_metadata lengths don't match media_paths length
-        """
-
-        # Setup tracking variables
-        total_uploads = len(datapoints)
-
-        # Create and start progress tracking thread
-        progress_tracker = ProgressTracker(
-            dataset_id=self.id,
-            openapi_service=self.openapi_service,
-            total_uploads=total_uploads,
-            progress_poll_interval=progress_poll_interval,
-        )
-        progress_thread = progress_tracker.create_thread()
-        progress_thread.start()
-
-        # Process uploads in chunks
-        try:
-            successful_uploads, failed_uploads = self._process_uploads_in_chunks(
-                datapoints,
-            )
-        finally:
-            progress_tracker.complete()
-            progress_thread.join(10)
-
-        if failed_uploads:
-            logger.error(
-                "Upload failed for %s datapoints: %s",
-                len(failed_uploads),
-                failed_uploads,
-            )
-
-        return successful_uploads, failed_uploads
-
     def __str__(self) -> str:
         return f"RapidataDataset(id={self.id})"