rapidata 2.23.2__py3-none-any.whl → 3.0.1__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. "
+                "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,8 @@
+from ._datapoint import Datapoint
+from .metadata import (
+    Metadata,
+    PrivateTextMetadata,
+    PublicTextMetadata,
+    SelectWordsMetadata,
+    PromptIdentifierMetadata,
+)

rapidata/rapidata_client/datapoints/_asset_uploader.py

@@ -0,0 +1,194 @@
+from __future__ import annotations
+
+import re
+import os
+from typing import TYPE_CHECKING
+
+from rapidata.rapidata_client.config.upload_config import (
+    register_upload_config_handler,
+    UploadConfig,
+)
+from rapidata.service.openapi_service import OpenAPIService
+from rapidata.rapidata_client.config import logger, rapidata_config, tracer
+from rapidata.rapidata_client.datapoints._single_flight_cache import SingleFlightCache
+from diskcache import FanoutCache
+
+if TYPE_CHECKING:
+    from rapidata.api_client.models.i_asset_input import IAssetInput
+
+
+class AssetUploader:
+    _file_cache: SingleFlightCache = SingleFlightCache(
+        "File cache",
+        storage=FanoutCache(
+            rapidata_config.upload.cacheLocation,
+            shards=rapidata_config.upload.maxWorkers,
+            timeout=rapidata_config.upload.cacheTimeout,
+            size_limit=rapidata_config.upload.cacheSizeLimit,
+        ),
+    )
+    _url_cache: SingleFlightCache = SingleFlightCache("URL cache")
+
+    def __init__(self, openapi_service: OpenAPIService):
+        self.openapi_service = openapi_service
+        register_upload_config_handler(self._handle_config_update)
+
+    @classmethod
+    def _handle_config_update(cls, config: UploadConfig):
+        """Handle updates to the upload config by re-creating the file cache storage."""
+        logger.debug("Updating AssetUploader file cache with new config")
+        try:
+            cls._file_cache.set_storage(
+                FanoutCache(
+                    config.cacheLocation,
+                    shards=config.maxWorkers,
+                    timeout=config.cacheTimeout,
+                    size_limit=config.cacheSizeLimit,
+                )
+            )
+            logger.info(
+                "AssetUploader file 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 file cache: {e}")
+
+    def _get_file_cache_key(self, asset: str) -> str:
+        """Generate cache key for a file, 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 _get_url_cache_key(self, url: str) -> str:
+        """Generate cache key for a URL, including environment."""
+        env = self.openapi_service.environment
+        return f"{env}@{url}"
+
+    def _upload_url_asset(self, url: str) -> str:
+        """Upload a URL asset, with optional caching."""
+
+        def upload_url() -> str:
+            response = self.openapi_service.asset_api.asset_url_post(url=url)
+            logger.info(
+                "Asset uploaded from URL: %s, file name: %s", url, response.file_name
+            )
+            return response.file_name
+
+        if not rapidata_config.upload.cacheUploads:
+            return upload_url()
+
+        return self._url_cache.get_or_fetch(
+            self._get_url_cache_key(url),
+            upload_url,
+            should_cache=rapidata_config.upload.cacheUploads,
+        )
+
+    def _upload_file_asset(self, file_path: str) -> str:
+        """Upload a local file asset, with optional caching."""
+
+        def upload_file() -> str:
+            response = self.openapi_service.asset_api.asset_file_post(file=file_path)
+            logger.info(
+                "Asset uploaded from file: %s, file name: %s",
+                file_path,
+                response.file_name,
+            )
+            return response.file_name
+
+        if not rapidata_config.upload.cacheUploads:
+            return upload_file()
+
+        return self._file_cache.get_or_fetch(
+            self._get_file_cache_key(file_path),
+            upload_file,
+            should_cache=rapidata_config.upload.cacheUploads,
+        )
+
+    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):
+                return self._upload_url_asset(asset)
+
+            return self._upload_file_asset(asset)
+
+    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._file_cache.clear()
+        self._url_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/_single_flight_cache.py

@@ -0,0 +1,82 @@
+import threading
+from typing import Callable, cast
+
+from diskcache import FanoutCache
+from rapidata.rapidata_client.config import logger
+
+from concurrent.futures import Future
+
+
+class SingleFlightCache:
+    """Cache with single-flight pattern to prevent duplicate concurrent fetches."""
+
+    def __init__(self, name: str, storage: dict[str, str] | FanoutCache | None = None):
+        self._name = name
+        self._storage: dict[str, str] | FanoutCache = (
+            storage if storage is not None else {}
+        )
+        self._in_flight: dict[str, Future[str]] = {}
+        self._lock = threading.Lock()
+
+    def set_storage(self, storage: dict[str, str] | FanoutCache) -> None:
+        """Replace the cache storage."""
+        with self._lock:
+            old_storage = self._storage
+            self._storage = storage
+            if isinstance(old_storage, FanoutCache):
+                try:
+                    old_storage.close()
+                except Exception:
+                    pass
+
+    def get_or_fetch(
+        self,
+        key: str,
+        fetch_fn: Callable[[], str],
+        should_cache: bool = True,
+    ) -> str:
+        """Get value from cache or fetch it, preventing duplicate concurrent fetches."""
+        # Fast path - check cache without lock
+        cached = self._storage.get(key)
+        if cached is not None:
+            logger.debug("%s: cache hit", self._name)
+            return cast(str, cached)
+
+        with self._lock:
+            # Double-check cache under lock
+            cached = self._storage.get(key)
+            if cached is not None:
+                logger.debug("%s: cache hit", self._name)
+                return cast(str, cached)
+
+            # Check if there's an in-flight request
+            in_flight = self._in_flight.get(key)
+            if in_flight is None:
+                in_flight = Future()
+                self._in_flight[key] = in_flight
+                should_fetch = True
+            else:
+                should_fetch = False
+
+        if not should_fetch:
+            logger.debug("%s: waiting for in-flight request", self._name)
+            return in_flight.result()
+
+        # We need to fetch
+        try:
+            result = fetch_fn()
+            if should_cache:
+                self._storage[key] = result
+                logger.debug("%s: cached result", self._name)
+            in_flight.set_result(result)
+            return result
+        except Exception as e:
+            in_flight.set_exception(e)
+            raise
+        finally:
+            with self._lock:
+                self._in_flight.pop(key, None)
+
+    def clear(self) -> None:
+        """Clear the cache."""
+        self._storage.clear()

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

@@ -1,5 +1,5 @@
 from ._base_metadata import Metadata
 from ._private_text_metadata import PrivateTextMetadata
 from ._public_text_metadata import PublicTextMetadata
-from ._prompt_metadata import PromptMetadata
 from ._select_words_metadata import SelectWordsMetadata
+from ._prompt_identifier_metadata import PromptIdentifierMetadata

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/_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/dataset/__init__.py

@@ -0,0 +1 @@
+from ._rapidata_dataset import RapidataDataset