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

rapidata/rapidata_client/config/upload_config.py

@@ -0,0 +1,64 @@
+from pathlib import Path
+from typing import Callable
+from pydantic import BaseModel, Field, field_validator
+from rapidata.rapidata_client.config import logger
+
+# Type alias for config update handlers
+UploadConfigUpdateHandler = Callable[["UploadConfig"], None]
+
+# Global list to store registered handlers
+_upload_config_handlers: list[UploadConfigUpdateHandler] = []
+
+
+def register_upload_config_handler(handler: UploadConfigUpdateHandler) -> None:
+    """Register a handler to be called when the upload configuration updates."""
+    _upload_config_handlers.append(handler)
+
+
+def unregister_upload_config_handler(handler: UploadConfigUpdateHandler) -> None:
+    """Unregister a previously registered handler."""
+    if handler in _upload_config_handlers:
+        _upload_config_handlers.remove(handler)
+
+
+class UploadConfig(BaseModel):
+    """
+    Holds the configuration for the upload process.
+
+    Attributes:
+        maxWorkers (int): The maximum number of worker threads for processing media paths. Defaults to 25.
+        maxRetries (int): The maximum number of retries for failed uploads. Defaults to 3.
+    """
+
+    maxWorkers: int = Field(default=25)
+    maxRetries: int = Field(default=3)
+    cacheUploads: bool = Field(default=True)
+    cacheTimeout: float = Field(default=0.1)
+    cacheLocation: Path = Field(default=Path.home() / ".rapidata" / "upload_cache")
+    cacheSizeLimit: int = Field(default=100_000_000)  # 100MB
+
+    @field_validator("maxWorkers")
+    @classmethod
+    def validate_max_workers(cls, v: int) -> int:
+        if v > 200:
+            logger.warning(
+                f"maxWorkers is set to {v}, which is above the recommended limit of 200 and default optimum of 25. "
+                "This may lead to suboptimal performance due to connection pool constraints."
+            )
+        return v
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self._notify_handlers()
+
+    def __setattr__(self, name: str, value) -> None:
+        super().__setattr__(name, value)
+        self._notify_handlers()
+
+    def _notify_handlers(self) -> None:
+        """Notify all registered handlers that the configuration has updated."""
+        for handler in _upload_config_handlers:
+            try:
+                handler(self)
+            except Exception as e:
+                logger.warning(f"Warning: UploadConfig handler failed: {e}")

rapidata/rapidata_client/datapoints/__init__.py

@@ -0,0 +1,10 @@
+from ._datapoint import Datapoint
+from .metadata import (
+    Metadata,
+    PromptMetadata,
+    PrivateTextMetadata,
+    PublicTextMetadata,
+    SelectWordsMetadata,
+    MediaAssetMetadata,
+    PromptIdentifierMetadata,
+)

rapidata/rapidata_client/datapoints/_asset_uploader.py

@@ -0,0 +1,169 @@
+from __future__ import annotations
+
+import re
+import os
+from typing import TYPE_CHECKING, cast
+
+from rapidata.rapidata_client.config.upload_config import register_upload_config_handler
+from rapidata.rapidata_client.config.upload_config import UploadConfig
+from rapidata.service.openapi_service import OpenAPIService
+from rapidata.rapidata_client.config import logger
+from rapidata.rapidata_client.config import tracer
+from rapidata.rapidata_client.config import rapidata_config
+from diskcache import FanoutCache
+
+if TYPE_CHECKING:
+    from rapidata.api_client.models.i_asset_input import IAssetInput
+
+
+class AssetUploader:
+    _shared_upload_cache: FanoutCache = FanoutCache(
+        rapidata_config.upload.cacheLocation,
+        shards=rapidata_config.upload.maxWorkers,
+        timeout=rapidata_config.upload.cacheTimeout,
+        size_limit=rapidata_config.upload.cacheSizeLimit,
+    )
+
+    def __init__(self, openapi_service: OpenAPIService):
+        self.openapi_service = openapi_service
+        # Register the handler for config update changes
+        register_upload_config_handler(self._handle_config_update)
+
+    @classmethod
+    def _handle_config_update(cls, config: UploadConfig):
+        """
+        Handle updates to the upload config (e.g., cache location, size, workers, timeout).
+        Re-instantiates the shared cache based on the new configuration.
+        """
+        logger.debug("Updating AssetUploader shared upload cache with new config")
+        try:
+            # Dispose of the old cache (closes open resources)
+            old_cache = getattr(cls, "_shared_upload_cache", None)
+            if old_cache:
+                try:
+                    old_cache.close()
+                except Exception:
+                    pass  # ignore close errors
+
+            # Re-initialize with the updated config
+            cls._shared_upload_cache = FanoutCache(
+                config.cacheLocation,
+                shards=config.maxWorkers,
+                timeout=config.cacheTimeout,
+                size_limit=config.cacheSizeLimit,
+            )
+            logger.info(
+                "AssetUploader shared upload cache updated: location=%s, shards=%s, timeout=%s, size_limit=%s",
+                config.cacheLocation,
+                config.maxWorkers,
+                config.cacheTimeout,
+                config.cacheSizeLimit,
+            )
+        except Exception as e:
+            logger.warning(f"Failed to update AssetUploader shared upload cache: {e}")
+
+    def _get_cache_key(self, asset: str) -> str:
+        """Generate cache key for an asset, including environment."""
+        env = self.openapi_service.environment
+        if not os.path.exists(asset):
+            raise FileNotFoundError(f"File not found: {asset}")
+
+        stat = os.stat(asset)
+        return f"{env}@{asset}:{stat.st_size}:{stat.st_mtime_ns}"
+
+    def upload_asset(self, asset: str) -> str:
+        with tracer.start_as_current_span("AssetUploader.upload_asset"):
+            logger.debug("Uploading asset: %s", asset)
+            assert isinstance(asset, str), "Asset must be a string"
+            if re.match(r"^https?://", asset):
+                response = self.openapi_service.asset_api.asset_url_post(
+                    url=asset,
+                )
+                return response.file_name
+
+            asset_key = self._get_cache_key(asset)
+            cached_value = self._shared_upload_cache.get(asset_key)
+            if cached_value is not None:
+                logger.debug("Asset found in cache")
+                return cast(str, cached_value)  # Type hint for the linter
+
+            response = self.openapi_service.asset_api.asset_file_post(
+                file=asset,
+            )
+            if rapidata_config.upload.cacheUploads:
+                self._shared_upload_cache[asset_key] = response.file_name
+                logger.debug("Asset added to cache")
+            logger.info("Asset uploaded: %s", response.file_name)
+            return response.file_name
+
+    def get_uploaded_text_input(self, assets: list[str] | str) -> IAssetInput:
+        from rapidata.api_client.models.i_asset_input import IAssetInput
+        from rapidata.api_client.models.i_asset_input_text_asset_input import (
+            IAssetInputTextAssetInput,
+        )
+        from rapidata.api_client.models.i_asset_input_multi_asset_input import (
+            IAssetInputMultiAssetInput,
+        )
+
+        if isinstance(assets, list):
+            return IAssetInput(
+                actual_instance=IAssetInputMultiAssetInput(
+                    _t="MultiAssetInput",
+                    assets=[
+                        IAssetInput(
+                            actual_instance=IAssetInputTextAssetInput(
+                                _t="TextAssetInput", text=asset
+                            )
+                        )
+                        for asset in assets
+                    ],
+                )
+            )
+        else:
+            return IAssetInput(
+                actual_instance=IAssetInputTextAssetInput(
+                    _t="TextAssetInput", text=assets
+                )
+            )
+
+    def get_uploaded_asset_input(self, assets: list[str] | str) -> IAssetInput:
+        from rapidata.api_client.models.i_asset_input import IAssetInput
+        from rapidata.api_client.models.i_asset_input_existing_asset_input import (
+            IAssetInputExistingAssetInput,
+        )
+        from rapidata.api_client.models.i_asset_input_multi_asset_input import (
+            IAssetInputMultiAssetInput,
+        )
+
+        if isinstance(assets, list):
+            return IAssetInput(
+                actual_instance=IAssetInputMultiAssetInput(
+                    _t="MultiAssetInput",
+                    assets=[
+                        IAssetInput(
+                            actual_instance=IAssetInputExistingAssetInput(
+                                _t="ExistingAssetInput",
+                                name=self.upload_asset(asset),
+                            )
+                        )
+                        for asset in assets
+                    ],
+                )
+            )
+        else:
+            return IAssetInput(
+                actual_instance=IAssetInputExistingAssetInput(
+                    _t="ExistingAssetInput",
+                    name=self.upload_asset(assets),
+                )
+            )
+
+    def clear_cache(self):
+        self._shared_upload_cache.clear()
+        logger.info("Upload cache cleared")
+
+    def __str__(self) -> str:
+        return f"AssetUploader(openapi_service={self.openapi_service})"
+
+    def __repr__(self) -> str:
+        return f"AssetUploader(openapi_service={self.openapi_service})"

rapidata/rapidata_client/datapoints/_datapoint.py

@@ -0,0 +1,92 @@
+from __future__ import annotations
+
+from typing import Literal, TYPE_CHECKING
+
+from pydantic import BaseModel, model_validator, field_validator
+from typing_extensions import Self
+from rapidata.rapidata_client.datapoints.assets.constants import (
+    ALLOWED_VIDEO_EXTENSIONS,
+    ALLOWED_IMAGE_EXTENSIONS,
+    ALLOWED_AUDIO_EXTENSIONS,
+)
+from rapidata.rapidata_client.config import logger
+
+if TYPE_CHECKING:
+    from rapidata.api_client.models.asset_type import AssetType
+    from rapidata.api_client.models.prompt_type import PromptType
+
+
+class Datapoint(BaseModel):
+    asset: str | list[str]
+    data_type: Literal["text", "media"]
+    context: str | None = None
+    media_context: str | None = None
+    sentence: str | None = None
+    private_note: str | None = None
+    group: str | None = None
+
+    @field_validator("context")
+    @classmethod
+    def context_not_empty(cls, v: str | None) -> str | None:
+        if v is not None and v == "":
+            raise ValueError(
+                "context cannot be an empty string. If not needed, set to None."
+            )
+        return v
+
+    @field_validator("media_context")
+    @classmethod
+    def media_context_not_empty(cls, v: str | None) -> str | None:
+        if v is not None and v == "":
+            raise ValueError(
+                "media_context cannot be an empty string. If not needed, set to None."
+            )
+        return v
+
+    @field_validator("sentence")
+    @classmethod
+    def sentence_has_space(cls, v: str | None) -> str | None:
+        if v is not None and len(v.split()) <= 1:
+            raise ValueError("sentence must contain at least two words.")
+        return v
+
+    @model_validator(mode="after")
+    def check_sentence_and_context(self) -> Self:
+        if isinstance(self.sentence, str) and isinstance(self.context, str):
+            raise ValueError(
+                "Both 'sentence' and 'context' cannot be strings at the same time."
+            )
+        return self
+
+    def get_asset_type(self) -> AssetType:
+        from rapidata.api_client.models.asset_type import AssetType
+
+        if self.data_type == "text":
+            return AssetType.TEXT
+
+        evaluation_asset = self.asset[0] if isinstance(self.asset, list) else self.asset
+        if any(evaluation_asset.endswith(ext) for ext in ALLOWED_IMAGE_EXTENSIONS):
+            return AssetType.IMAGE
+        elif any(evaluation_asset.endswith(ext) for ext in ALLOWED_VIDEO_EXTENSIONS):
+            return AssetType.VIDEO
+        elif any(evaluation_asset.endswith(ext) for ext in ALLOWED_AUDIO_EXTENSIONS):
+            return AssetType.AUDIO
+        else:
+            logger.debug(
+                f"Cannot get asset type for asset type: {type(evaluation_asset)}"
+            )
+            return AssetType.NONE
+
+    def get_prompt_type(self) -> list[PromptType]:
+        from rapidata.api_client.models.prompt_type import PromptType
+
+        prompt_types = []
+        if self.context:
+            prompt_types.append(PromptType.TEXT)
+        if self.media_context:
+            prompt_types.append(PromptType.ASSET)
+
+        if len(prompt_types) == 0:
+            return [PromptType.NONE]
+
+        return prompt_types

rapidata/rapidata_client/datapoints/_datapoint_uploader.py

@@ -0,0 +1,47 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from rapidata.rapidata_client.datapoints._datapoint import Datapoint
+from rapidata.service.openapi_service import OpenAPIService
+from rapidata.rapidata_client.datapoints._asset_uploader import AssetUploader
+
+if TYPE_CHECKING:
+    from rapidata.api_client.models.create_datapoint_result import CreateDatapointResult
+
+
+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:
+        from rapidata.api_client.models.create_datapoint_model import (
+            CreateDatapointModel,
+        )
+        from rapidata.api_client.models.i_asset_input import IAssetInput
+
+        uploaded_asset: IAssetInput = (
+            self.asset_uploader.get_uploaded_asset_input(datapoint.asset)
+            if datapoint.data_type == "media"
+            else self.asset_uploader.get_uploaded_text_input(datapoint.asset)
+        )
+        return self.openapi_service.dataset_api.dataset_dataset_id_datapoint_post(
+            dataset_id=dataset_id,
+            create_datapoint_model=CreateDatapointModel(
+                asset=uploaded_asset,
+                context=datapoint.context,
+                contextAsset=(
+                    self.asset_uploader.get_uploaded_asset_input(
+                        datapoint.media_context
+                    )
+                    if datapoint.media_context
+                    else None
+                ),
+                transcription=datapoint.sentence,
+                sortIndex=index,
+                group=datapoint.group,
+            ),
+        )

rapidata/rapidata_client/datapoints/_datapoints_validator.py

@@ -0,0 +1,70 @@
+from itertools import zip_longest
+from typing import Literal, cast, Iterable
+from rapidata.rapidata_client.datapoints._datapoint import Datapoint
+
+
+class DatapointsValidator:
+    @staticmethod
+    def validate_datapoints(
+        datapoints: list[str] | list[list[str]],
+        contexts: list[str] | None = None,
+        media_contexts: list[str] | None = None,
+        sentences: list[str] | None = None,
+        private_notes: list[str] | None = None,
+        groups: list[str] | None = None,
+    ) -> None:
+        if contexts and len(contexts) != len(datapoints):
+            raise ValueError("Number of contexts must match number of datapoints")
+        if media_contexts and len(media_contexts) != len(datapoints):
+            raise ValueError("Number of media contexts must match number of datapoints")
+        if sentences and len(sentences) != len(datapoints):
+            raise ValueError("Number of sentences must match number of datapoints")
+        if private_notes and len(private_notes) != len(datapoints):
+            raise ValueError("Number of private notes must match number of datapoints")
+        if groups and (
+            len(groups) != len(datapoints) or len(groups) != len(set(groups))
+        ):
+            raise ValueError(
+                "Number of groups must match number of datapoints and must be unique."
+            )
+
+    @staticmethod
+    def map_datapoints(
+        datapoints: list[str] | list[list[str]],
+        contexts: list[str] | None = None,
+        media_contexts: list[str] | None = None,
+        sentences: list[str] | None = None,
+        private_notes: list[str] | None = None,
+        groups: list[str] | None = None,
+        data_type: Literal["text", "media"] = "media",
+    ) -> list[Datapoint]:
+        DatapointsValidator.validate_datapoints(
+            datapoints=datapoints,
+            contexts=contexts,
+            media_contexts=media_contexts,
+            sentences=sentences,
+            private_notes=private_notes,
+            groups=groups,
+        )
+        return [
+            Datapoint(
+                asset=asset,
+                data_type=data_type,
+                context=context,
+                media_context=media_context,
+                sentence=sentence,
+                private_note=private_note,
+                group=group,
+            )
+            for asset, context, media_context, sentence, private_note, group in cast(
+                "Iterable[tuple[str | list[str], str | None, str | None, str | None, str | None, str | None]]",  # because iterator only supports 5 arguments with specific type casting
+                zip_longest(
+                    datapoints,
+                    contexts or [],
+                    media_contexts or [],
+                    sentences or [],
+                    private_notes or [],
+                    groups or [],
+                ),
+            )
+        ]

rapidata/rapidata_client/datapoints/assets/constants.py

@@ -0,0 +1,7 @@
+ALLOWED_IMAGE_EXTENSIONS = [".jpg", ".jpeg", ".png", ".gif", ".webp"]
+ALLOWED_AUDIO_EXTENSIONS = [".mp3"]
+ALLOWED_VIDEO_EXTENSIONS = [".mp4"]
+
+ALLOWED_MEDIA_EXTENSIONS = (
+    ALLOWED_IMAGE_EXTENSIONS + ALLOWED_AUDIO_EXTENSIONS + ALLOWED_VIDEO_EXTENSIONS
+)

rapidata/rapidata_client/{metadata → datapoints/metadata}/__init__.py

@@ -4,3 +4,4 @@ from ._public_text_metadata import PublicTextMetadata
 from ._prompt_metadata import PromptMetadata
 from ._select_words_metadata import SelectWordsMetadata
 from ._media_asset_metadata import MediaAssetMetadata
+from ._prompt_identifier_metadata import PromptIdentifierMetadata

rapidata/rapidata_client/datapoints/metadata/_media_asset_metadata.py

@@ -0,0 +1,26 @@
+from rapidata.api_client.models.prompt_asset_metadata_input import (
+    PromptAssetMetadataInput,
+)
+from rapidata.rapidata_client.datapoints.metadata._base_metadata import Metadata
+from rapidata.api_client.models.multi_asset_input_assets_inner import (
+    ExistingAssetInput,
+)
+from rapidata.api_client.models.prompt_asset_metadata_input_asset import (
+    PromptAssetMetadataInputAsset,
+)
+from pydantic import BaseModel
+
+
+class MediaAssetMetadata(Metadata, BaseModel):
+    internal_file_name: str
+
+    def to_model(self):
+        return PromptAssetMetadataInput(
+            _t="PromptAssetMetadataInput",
+            asset=PromptAssetMetadataInputAsset(
+                actual_instance=ExistingAssetInput(
+                    _t="ExistingAssetInput",
+                    name=self.internal_file_name,
+                ),
+            ),
+        )

rapidata/rapidata_client/datapoints/metadata/_private_text_metadata.py

@@ -0,0 +1,16 @@
+from rapidata.api_client.models.private_text_metadata_input import (
+    PrivateTextMetadataInput,
+)
+from rapidata.rapidata_client.datapoints.metadata._base_metadata import Metadata
+from pydantic import BaseModel
+
+
+class PrivateTextMetadata(Metadata, BaseModel):
+
+    text: str
+    identifier: str = "private_text"
+
+    def to_model(self):
+        return PrivateTextMetadataInput(
+            _t="PrivateTextMetadataInput", identifier=self.identifier, text=self.text
+        )

rapidata/rapidata_client/datapoints/metadata/_prompt_identifier_metadata.py

@@ -0,0 +1,14 @@
+from rapidata.rapidata_client.datapoints.metadata._base_metadata import Metadata
+from rapidata.api_client.models.private_text_metadata_input import (
+    PrivateTextMetadataInput,
+)
+from pydantic import BaseModel
+
+
+class PromptIdentifierMetadata(Metadata, BaseModel):
+    identifier: str
+
+    def to_model(self):
+        return PrivateTextMetadataInput(
+            _t="PrivateTextMetadataInput", identifier="prompt-id", text=self.identifier
+        )

rapidata/rapidata_client/datapoints/metadata/_prompt_metadata.py

@@ -0,0 +1,12 @@
+from rapidata.rapidata_client.datapoints.metadata._base_metadata import Metadata
+from rapidata.api_client.models.prompt_metadata_input import PromptMetadataInput
+from pydantic import BaseModel
+
+
+class PromptMetadata(Metadata, BaseModel):
+    """The PromptMetadata class is used to display a prompt to the user."""
+
+    prompt: str
+
+    def to_model(self):
+        return PromptMetadataInput(_t="PromptMetadataInput", prompt=self.prompt)

rapidata/rapidata_client/datapoints/metadata/_public_text_metadata.py

@@ -0,0 +1,16 @@
+from rapidata.api_client.models.public_text_metadata_input import (
+    PublicTextMetadataInput,
+)
+from rapidata.rapidata_client.datapoints.metadata._base_metadata import Metadata
+from pydantic import BaseModel
+
+
+class PublicTextMetadata(Metadata, BaseModel):
+
+    text: str
+    identifier: str = "public_text"
+
+    def to_model(self):
+        return PublicTextMetadataInput(
+            _t="PublicTextMetadataInput", identifier=self.identifier, text=self.text
+        )

rapidata/rapidata_client/{metadata → datapoints/metadata}/_select_words_metadata.py

@@ -1,15 +1,14 @@
 from rapidata.api_client.models.transcription_metadata_input import (
     TranscriptionMetadataInput,
 )
-from rapidata.rapidata_client.metadata._base_metadata import Metadata
+from rapidata.rapidata_client.datapoints.metadata._base_metadata import Metadata
+from pydantic import BaseModel
 
 
-class SelectWordsMetadata(Metadata):
+class SelectWordsMetadata(Metadata, BaseModel):
     """SelectWordsMetadata Class is used to define the Sentence that will be display to the user."""
 
-    def __init__(self, select_words: str):
-        super().__init__()
-        self.select_words = select_words
+    select_words: str
 
     def to_model(self):
         return TranscriptionMetadataInput(

rapidata/rapidata_client/demographic/demographic_manager.py

@@ -1,29 +1,54 @@
 from rapidata.service.openapi_service import OpenAPIService
-from rapidata.rapidata_client.assets import MediaAsset
-from rapidata.api_client.models.create_demographic_rapid_model import CreateDemographicRapidModel
-from rapidata.api_client.models.classify_payload import ClassifyPayload
-from rapidata.rapidata_client.logging import logger
+from rapidata.rapidata_client.config import logger
+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)
+
+    def create_demographic_rapid(
+        self, instruction: str, answer_options: list[str], datapoint: str, key: str
+    ):
+        from rapidata.api_client.models.classify_payload import ClassifyPayload
+        from rapidata.api_client.models.classify_payload_category import (
+            ClassifyPayloadCategory,
+        )
+        from rapidata.api_client.models.create_demographic_rapid_model import (
+            CreateDemographicRapidModel,
+        )
+        from rapidata.api_client.models.i_asset_input import IAssetInput
+        from rapidata.api_client.models.i_asset_input_existing_asset_input import (
+            IAssetInputExistingAssetInput,
+        )
+
         model = CreateDemographicRapidModel(
             key=key,
             payload=ClassifyPayload(
-                _t="ClassifyPayload",
-                possibleCategories=answer_options,
-                title=instruction
-            )
+                categories=[
+                    ClassifyPayloadCategory(label=option, value=option)
+                    for option in answer_options
+                ],
+                title=instruction,
+            ),
+            asset=IAssetInput(
+                actual_instance=IAssetInputExistingAssetInput(
+                    _t="ExistingAssetInput",
+                    name=self._asset_uploader.upload_asset(datapoint),
+                ),
+            ),
+        )
+
+        result = self._openapi_service.rapid_api.rapid_demographic_post(
+            create_demographic_rapid_model=model
         )
-        
-        self._openapi_service.rapid_api.rapid_demographic_post(model=model, file=[media.to_file()])
-        
+        logger.info(f"Demographic Rapid created: {result.rapid_id}")
+        return result.rapid_id
+
+    def __str__(self) -> str:
+        return "DemographicManager"
+
+    def __repr__(self) -> str:
+        return self.__str__()

rapidata/rapidata_client/exceptions/__init__.py

@@ -0,0 +1 @@
+from .failed_upload_exception import FailedUploadException