rapidata 2.37.0__py3-none-any.whl → 3.0.0a8__py3-none-any.whl

rapidata/rapidata_client/config/tracer.py

@@ -0,0 +1,184 @@
+from typing import Protocol, runtime_checkable, Any
+import platform
+import sys
+import os
+from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.sdk.trace.export import BatchSpanProcessor
+from opentelemetry.sdk.resources import Resource
+from rapidata import __version__
+from .logging_config import LoggingConfig, register_config_handler
+from rapidata.rapidata_client.config import logger
+
+
+def get_system_attributes() -> dict[str, str | int | None]:
+    """Gather system telemetry for traces."""
+    try:
+        attrs = {
+            "system.os": platform.system(),
+            "system.os.version": platform.release(),
+            "system.arch": platform.machine(),
+            "python.version": f"{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}",
+            "process.cpu_count": os.cpu_count(),
+        }
+        logger.debug(f"System attributes: {attrs}")
+        return attrs
+    except Exception:
+        logger.debug("Failed to get system attributes, returning empty dict")
+        return {}
+
+
+@runtime_checkable
+class TracerProtocol(Protocol):
+    """Protocol that defines the tracer interface for type checking."""
+
+    def start_span(self, name: str, *args, **kwargs) -> Any: ...
+    def start_as_current_span(self, name: str, *args, **kwargs) -> Any: ...
+    def set_session_id(self, session_id: str) -> None: ...
+
+
+class NoOpSpan:
+    """A no-op span that does nothing when tracing is disabled."""
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args):
+        pass
+
+    def set_attribute(self, *args, **kwargs):
+        pass
+
+    def set_status(self, *args, **kwargs):
+        pass
+
+    def add_event(self, *args, **kwargs):
+        pass
+
+    def end(self, *args, **kwargs):
+        pass
+
+    def __getattr__(self, name: str) -> Any:
+        """Return self for any method call to maintain chainability."""
+        return lambda *args, **kwargs: self
+
+
+class NoOpTracer:
+    """A no-op tracer that returns no-op spans when tracing is disabled."""
+
+    def start_span(self, name: str, *args, **kwargs) -> NoOpSpan:
+        return NoOpSpan()
+
+    def start_as_current_span(self, name: str, *args, **kwargs) -> NoOpSpan:
+        return NoOpSpan()
+
+    def set_session_id(self, session_id: str) -> None:
+        pass
+
+    def __getattr__(self, name: str) -> Any:
+        """Delegate to no-op behavior."""
+        return lambda *args, **kwargs: NoOpSpan()
+
+
+class SpanContextManagerWrapper:
+    """Wrapper for span context managers to add session_id on enter."""
+
+    def __init__(self, context_manager: Any, session_id: str | None):
+        self._context_manager = context_manager
+        self.session_id = session_id
+
+    def __enter__(self):
+        span = self._context_manager.__enter__()
+        if self.session_id and hasattr(span, "set_attribute"):
+            span.set_attribute("SDK.session.id", self.session_id)
+        return span
+
+    def __exit__(self, *args):
+        return self._context_manager.__exit__(*args)
+
+
+class RapidataTracer:
+    """Tracer implementation that updates when the configuration changes."""
+
+    def __init__(self, name: str = __name__):
+        self._name = name
+        self._otlp_initialized = False
+        self._tracer_provider = None
+        self._real_tracer = None
+        self._no_op_tracer = NoOpTracer()
+        self._enabled = True  # Default to enabled
+        self.session_id: str | None = None
+
+        # Register this tracer to receive configuration updates
+        register_config_handler(self._handle_config_update)
+
+    def _handle_config_update(self, config: LoggingConfig) -> None:
+        """Handle configuration updates."""
+        self._update_tracer(config)
+
+    def _update_tracer(self, config: LoggingConfig) -> None:
+        """Update the tracer based on the new configuration."""
+        self._enabled = config.enable_otlp
+
+        # Initialize OTLP tracing only once and only if not disabled
+        if not self._otlp_initialized and config.enable_otlp:
+            try:
+                resource_attributes = {
+                    "service.name": "Rapidata.Python.SDK",
+                    "service.version": __version__,
+                    **get_system_attributes(),
+                }
+
+                resource = Resource.create(resource_attributes)
+
+                self._tracer_provider = TracerProvider(resource=resource)
+
+                exporter = OTLPSpanExporter(
+                    endpoint="https://otlp-sdk.rapidata.ai/v1/traces",
+                    timeout=30,
+                )
+
+                span_processor = BatchSpanProcessor(exporter)
+                self._tracer_provider.add_span_processor(span_processor)
+
+                self._real_tracer = self._tracer_provider.get_tracer(self._name)
+                self._otlp_initialized = True
+
+            except Exception as e:
+                logger.warning(f"Failed to initialize tracing: {e}")
+                self._enabled = False
+
+    def _add_session_id_to_span(self, span: Any) -> Any:
+        """Add session_id attribute to a span if session_id is set."""
+        if self.session_id and hasattr(span, "set_attribute"):
+            span.set_attribute("SDK.session.id", self.session_id)
+        return span
+
+    def start_span(self, name: str, *args, **kwargs) -> Any:
+        """Start a span, or return a no-op span if tracing is disabled."""
+        if self._enabled and self._real_tracer:
+            span = self._real_tracer.start_span(name, *args, **kwargs)
+            return self._add_session_id_to_span(span)
+        return self._no_op_tracer.start_span(name, *args, **kwargs)
+
+    def start_as_current_span(self, name: str, *args, **kwargs) -> Any:
+        """Start a span as current, or return a no-op span if tracing is disabled."""
+        if self._enabled and self._real_tracer:
+            context_manager = self._real_tracer.start_as_current_span(
+                name, *args, **kwargs
+            )
+            return SpanContextManagerWrapper(context_manager, self.session_id)
+        return self._no_op_tracer.start_as_current_span(name, *args, **kwargs)
+
+    def set_session_id(self, session_id: str) -> None:
+        self.session_id = session_id
+
+    def __getattr__(self, name: str) -> Any:
+        """Delegate attribute access to the appropriate tracer."""
+        if self._enabled and self._real_tracer:
+            return getattr(self._real_tracer, name)
+        return getattr(self._no_op_tracer, name)
+
+
+# Create the main tracer instance - type checkers will see it as TracerProtocol
+tracer: TracerProtocol = RapidataTracer()  # type: ignore[assignment]

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

@@ -1,11 +1,8 @@
 from ._datapoint import Datapoint
-from .assets import MediaAsset, MultiAsset, TextAsset
 from .metadata import (
     Metadata,
-    PromptMetadata,
     PrivateTextMetadata,
     PublicTextMetadata,
     SelectWordsMetadata,
-    MediaAssetMetadata,
     PromptIdentifierMetadata,
 )

rapidata/rapidata_client/datapoints/_asset_uploader.py

@@ -0,0 +1,184 @@
+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,
+    )
+    _url_memory_cache: dict[str, str] = {}
+
+    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 _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_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):
+                url_key = self._get_url_cache_key(asset)
+                cached_value = self._url_memory_cache.get(url_key)
+                if cached_value is not None:
+                    logger.debug("URL found in memory cache")
+                    return cached_value
+
+                response = self.openapi_service.asset_api.asset_url_post(
+                    url=asset,
+                )
+                self._url_memory_cache[url_key] = response.file_name
+                logger.debug("URL added to memory cache")
+                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()
+        self._url_memory_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})"