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

rapidata/rapidata_client/benchmark/rapidata_benchmark_manager.py

@@ -0,0 +1,219 @@
+from typing import Optional, cast
+from rapidata.rapidata_client.benchmark.rapidata_benchmark import RapidataBenchmark
+from rapidata.api_client.models.create_benchmark_model import CreateBenchmarkModel
+from rapidata.service.openapi_service import OpenAPIService
+from rapidata.api_client.models.query_model import QueryModel
+from rapidata.api_client.models.page_info import PageInfo
+from rapidata.api_client.models.root_filter import RootFilter
+from rapidata.api_client.models.filter import Filter
+from rapidata.api_client.models.sort_criterion import SortCriterion
+from rapidata.api_client.models.sort_direction import SortDirection
+from rapidata.api_client.models.filter_operator import FilterOperator
+from rapidata.rapidata_client.config import logger, tracer
+
+
+class RapidataBenchmarkManager:
+    """
+    A manager for benchmarks.
+
+    Used to create and retrieve benchmarks.
+
+    A benchmark is a collection of leaderboards.
+
+    Args:
+        openapi_service: The OpenAPIService instance for API interaction.
+    """
+
+    def __init__(self, openapi_service: OpenAPIService):
+        self.__openapi_service = openapi_service
+
+    def create_new_benchmark(
+        self,
+        name: str,
+        identifiers: Optional[list[str]] = None,
+        prompts: Optional[list[str | None] | list[str]] = None,
+        prompt_assets: Optional[list[str | None] | list[str]] = None,
+        tags: Optional[list[list[str] | None] | list[list[str]]] = None,
+    ) -> RapidataBenchmark:
+        """
+        Creates a new benchmark with the given name, identifiers, prompts, and media assets.
+        Everything is matched up by the indexes of the lists.
+
+        prompts or identifiers must be provided, as well as prompts or prompt_assets.
+
+        Args:
+            name: The name of the benchmark.
+            identifiers: The identifiers of the prompts/assets/tags that will be used to match up the media. If not provided, it will use the prompts as the identifiers.
+            prompts: The prompts that will be registered for the benchmark.
+            prompt_assets: The prompt assets that will be registered for the benchmark.
+            tags: The tags that will be associated with the prompts to use for filtering the leaderboard results. They will NOT be shown to the users.
+
+        Example:
+            ```python
+            name = "Example Benchmark"
+            identifiers = ["id1", "id2", "id3"]
+            prompts = ["prompt 1", "prompt 2", "prompt 3"]
+            prompt_assets = ["https://assets.rapidata.ai/prompt_1.jpg", "https://assets.rapidata.ai/prompt_2.jpg", "https://assets.rapidata.ai/prompt_3.jpg"]
+            tags = [["tag1", "tag2"], ["tag2"], ["tag2", "tag3"]]
+
+            benchmark = create_new_benchmark(name=name, identifiers=identifiers, prompts=prompts, prompt_assets=prompt_assets, tags=tags)
+            ```
+        """
+        with tracer.start_as_current_span(
+            "RapidataBenchmarkManager.create_new_benchmark"
+        ):
+            if not isinstance(name, str):
+                raise ValueError("Name must be a string.")
+
+            if prompts and (
+                not isinstance(prompts, list)
+                or not all(
+                    isinstance(prompt, str) or prompt is None for prompt in prompts
+                )
+            ):
+                raise ValueError("Prompts must be a list of strings or None.")
+
+            if prompt_assets and (
+                not isinstance(prompt_assets, list)
+                or not all(
+                    isinstance(asset, str) or asset is None for asset in prompt_assets
+                )
+            ):
+                raise ValueError("Media assets must be a list of strings or None.")
+
+            if identifiers and (
+                not isinstance(identifiers, list)
+                or not all(isinstance(identifier, str) for identifier in identifiers)
+            ):
+                raise ValueError("Identifiers must be a list of strings.")
+
+            if identifiers:
+                if not len(set(identifiers)) == len(identifiers):
+                    raise ValueError("Identifiers must be unique.")
+
+            if tags is not None:
+                if not isinstance(tags, list):
+                    raise ValueError("Tags must be a list of lists of strings or None.")
+
+                for tag in tags:
+                    if tag is not None and (
+                        not isinstance(tag, list)
+                        or not all(isinstance(item, str) for item in tag)
+                    ):
+                        raise ValueError(
+                            "Tags must be a list of lists of strings or None."
+                        )
+
+            if not identifiers and not prompts:
+                raise ValueError(
+                    "At least one of identifiers or prompts must be provided."
+                )
+
+            if not prompts and not prompt_assets:
+                raise ValueError(
+                    "At least one of prompts or media assets must be provided."
+                )
+
+            if not identifiers:
+                assert prompts is not None
+                if not len(set(prompts)) == len(prompts):
+                    raise ValueError(
+                        "Prompts must be unique. Otherwise use identifiers."
+                    )
+                if any(prompt is None for prompt in prompts):
+                    raise ValueError(
+                        "Prompts must not be None. Otherwise use identifiers."
+                    )
+
+                identifiers = cast(list[str], prompts)
+
+            assert identifiers is not None
+
+            expected_length = len(identifiers)
+
+            if not prompts:
+                prompts = cast(list[str | None], [None] * expected_length)
+
+            if not prompt_assets:
+                prompt_assets = cast(list[str | None], [None] * expected_length)
+
+            if not tags:
+                tags = cast(list[list[str] | None], [None] * expected_length)
+
+            # At this point, all variables are guaranteed to be lists, not None
+            assert prompts is not None
+            assert prompt_assets is not None
+            assert tags is not None
+
+            if not (expected_length == len(prompts) == len(prompt_assets) == len(tags)):
+                raise ValueError(
+                    "Identifiers, prompts, media assets, and tags must have the same length or set to None."
+                )
+
+            logger.info("Creating new benchmark %s", name)
+
+            benchmark_result = self.__openapi_service.benchmark_api.benchmark_post(
+                create_benchmark_model=CreateBenchmarkModel(
+                    name=name,
+                )
+            )
+
+            logger.info("Benchmark created with id %s", benchmark_result.id)
+
+            benchmark = RapidataBenchmark(
+                name, benchmark_result.id, self.__openapi_service
+            )
+
+            for identifier, prompt, asset, tag in zip(
+                identifiers, prompts, prompt_assets, tags
+            ):
+                benchmark.add_prompt(identifier, prompt, asset, tag)
+
+            return benchmark
+
+    def get_benchmark_by_id(self, id: str) -> RapidataBenchmark:
+        """
+        Returns a benchmark by its ID.
+        """
+        with tracer.start_as_current_span(
+            "RapidataBenchmarkManager.get_benchmark_by_id"
+        ):
+            benchmark_result = (
+                self.__openapi_service.benchmark_api.benchmark_benchmark_id_get(
+                    benchmark_id=id
+                )
+            )
+            return RapidataBenchmark(
+                benchmark_result.name, benchmark_result.id, self.__openapi_service
+            )
+
+    def find_benchmarks(
+        self, name: str = "", amount: int = 10
+    ) -> list[RapidataBenchmark]:
+        """
+        Returns a list of benchmarks by their name.
+        """
+        with tracer.start_as_current_span("RapidataBenchmarkManager.find_benchmarks"):
+            benchmark_result = self.__openapi_service.benchmark_api.benchmarks_get(
+                QueryModel(
+                    page=PageInfo(index=1, size=amount),
+                    filter=RootFilter(
+                        filters=[
+                            Filter(
+                                field="Name",
+                                operator=FilterOperator.CONTAINS,
+                                value=name,
+                            )
+                        ]
+                    ),
+                    sortCriteria=[
+                        SortCriterion(
+                            direction=SortDirection.DESC, propertyName="CreatedAt"
+                        )
+                    ],
+                )
+            )
+            return [
+                RapidataBenchmark(benchmark.name, benchmark.id, self.__openapi_service)
+                for benchmark in benchmark_result.items
+            ]

rapidata/rapidata_client/config/__init__.py

@@ -0,0 +1,4 @@
+from .logger import logger  # has to be instantiated before rapidata_config
+from .tracer import tracer  # has to be instantiated before rapidata_config
+from .rapidata_config import rapidata_config
+from .managed_print import managed_print

rapidata/rapidata_client/config/logger.py

@@ -0,0 +1,135 @@
+import logging
+from typing import Protocol, runtime_checkable
+from opentelemetry._logs import set_logger_provider
+from opentelemetry.exporter.otlp.proto.http._log_exporter import OTLPLogExporter
+from opentelemetry.sdk._logs import LoggerProvider, LoggingHandler
+from opentelemetry.sdk._logs.export import BatchLogRecordProcessor
+from opentelemetry.sdk.resources import Resource
+from rapidata import __version__
+from .logging_config import LoggingConfig, register_config_handler
+
+
+@runtime_checkable
+class LoggerProtocol(Protocol):
+    """Protocol that defines the logger interface for type checking."""
+
+    def debug(self, msg: object, *args, **kwargs) -> None: ...
+    def info(self, msg: object, *args, **kwargs) -> None: ...
+    def warning(self, msg: object, *args, **kwargs) -> None: ...
+    def warn(self, msg: object, *args, **kwargs) -> None: ...
+    def error(self, msg: object, *args, **kwargs) -> None: ...
+    def exception(self, msg: object, *args, exc_info=True, **kwargs) -> None: ...
+    def critical(self, msg: object, *args, **kwargs) -> None: ...
+    def fatal(self, msg: object, *args, **kwargs) -> None: ...
+    def log(self, level: int, msg: object, *args, **kwargs) -> None: ...
+    def isEnabledFor(self, level: int) -> bool: ...
+    def getEffectiveLevel(self) -> int: ...
+    def setLevel(self, level: int | str) -> None: ...
+    def addHandler(self, handler: logging.Handler) -> None: ...
+    def removeHandler(self, handler: logging.Handler) -> None: ...
+    @property
+    def handlers(self) -> list[logging.Handler]: ...
+    @property
+    def level(self) -> int: ...
+    @property
+    def name(self) -> str: ...
+
+
+class RapidataLogger:
+    """Logger implementation that updates when the configuration changes."""
+
+    def __init__(self, name: str = "rapidata"):
+        self._logger = logging.getLogger(name)
+        self._otlp_initialized = False
+        self._otlp_handler = None
+
+        # Register this logger to receive configuration updates
+        register_config_handler(self._handle_config_update)
+
+    def _handle_config_update(self, config: LoggingConfig) -> None:
+        """Handle configuration updates."""
+        self._update_logger(config)
+
+    def _update_logger(self, config: LoggingConfig) -> None:
+        """Update the logger based on the new configuration."""
+        # Initialize OTLP logging only once and only if not disabled
+        if not self._otlp_initialized and config.enable_otlp:
+            try:
+                logger_provider = LoggerProvider(
+                    resource=Resource.create(
+                        {
+                            "service.name": "Rapidata.Python.SDK",
+                            "service.version": __version__,
+                        }
+                    ),
+                )
+                set_logger_provider(logger_provider)
+
+                exporter = OTLPLogExporter(
+                    endpoint="https://otlp-sdk.rapidata.ai/v1/logs",
+                    timeout=30,
+                )
+
+                processor = BatchLogRecordProcessor(
+                    exporter,
+                    max_queue_size=2048,
+                    export_timeout_millis=30000,
+                    max_export_batch_size=512,
+                )
+
+                logger_provider.add_log_record_processor(processor)
+
+                # OTLP handler - captures DEBUG and above
+                self._otlp_handler = LoggingHandler(logger_provider=logger_provider)
+                self._otlp_handler.setLevel(logging.DEBUG)  # OTLP gets everything
+
+                self._otlp_initialized = True
+
+            except Exception as e:
+                self._logger.warning(f"Failed to initialize OTLP logging: {e}")
+                import traceback
+
+                traceback.print_exc()
+
+        # Console handler with configurable level
+        console_handler = logging.StreamHandler()
+        console_level = getattr(logging, config.level.upper())
+        console_handler.setLevel(console_level)
+        console_formatter = logging.Formatter(config.format)
+        console_handler.setFormatter(console_formatter)
+
+        # Configure the logger
+        self._logger.setLevel(logging.DEBUG)  # Logger must allow DEBUG for OTLP
+
+        # Remove any existing handlers (except OTLP when appropriate)
+        for handler in self._logger.handlers[:]:
+            if handler != self._otlp_handler:
+                self._logger.removeHandler(handler)
+            elif handler == self._otlp_handler and not config.enable_otlp:
+                self._logger.removeHandler(handler)
+
+        # Add OTLP handler if initialized and not disabled
+        if (
+            self._otlp_handler
+            and self._otlp_handler not in self._logger.handlers
+            and config.enable_otlp
+        ):
+            self._logger.addHandler(self._otlp_handler)
+
+        # Add console handler
+        self._logger.addHandler(console_handler)
+
+        # Add file handler if log_file is provided
+        if config.log_file:
+            file_handler = logging.FileHandler(config.log_file)
+            file_handler.setLevel(console_level)  # Use same level as console
+            file_formatter = logging.Formatter(config.format)
+            file_handler.setFormatter(file_formatter)
+            self._logger.addHandler(file_handler)
+
+    def __getattr__(self, name: str) -> object:
+        """Delegate attribute access to the underlying logger."""
+        return getattr(self._logger, name)
+
+
+logger: LoggerProtocol = RapidataLogger()  # type: ignore[assignment]

rapidata/rapidata_client/config/logging_config.py

@@ -0,0 +1,55 @@
+from typing import Callable
+from pydantic import BaseModel, Field
+from rapidata.rapidata_client.config import logger
+
+# Type alias for config update handlers
+ConfigUpdateHandler = Callable[["LoggingConfig"], None]
+
+# Global list to store registered handlers
+_config_handlers: list[ConfigUpdateHandler] = []
+
+
+def register_config_handler(handler: ConfigUpdateHandler) -> None:
+    """Register a handler to be called when the logging configuration updates."""
+    _config_handlers.append(handler)
+
+
+def unregister_config_handler(handler: ConfigUpdateHandler) -> None:
+    """Unregister a previously registered handler."""
+    if handler in _config_handlers:
+        _config_handlers.remove(handler)
+
+
+class LoggingConfig(BaseModel):
+    """
+    Holds the configuration for the logging process.
+
+    Attributes:
+        level (str): The logging level. Defaults to "WARNING".
+        log_file (str | None): The logging file. Defaults to None.
+        format (str): The logging format. Defaults to "%(asctime)s - %(name)s - %(levelname)s - %(message)s".
+        silent_mode (bool): Whether to disable the prints and progress bars. Does NOT affect the logging. Defaults to False.
+        enable_otlp (bool): Whether to enable OpenTelemetry trace logs. Defaults to True.
+    """
+
+    level: str = Field(default="WARNING")
+    log_file: str | None = Field(default=None)
+    format: str = Field(default="%(asctime)s - %(name)s - %(levelname)s - %(message)s")
+    silent_mode: bool = Field(default=False)
+    enable_otlp: bool = Field(default=True)
+
+    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 _config_handlers:
+            try:
+                handler(self)
+            except Exception as e:
+                logger.warning(f"Warning: Config handler failed: {e}")

rapidata/rapidata_client/config/managed_print.py

@@ -0,0 +1,6 @@
+from rapidata.rapidata_client.config import rapidata_config
+
+
+def managed_print(*args, **kwargs) -> None:
+    if not rapidata_config.logging.silent_mode:
+        print(*args, **kwargs)

rapidata/rapidata_client/config/order_config.py

@@ -0,0 +1,14 @@
+from pydantic import BaseModel, Field
+
+
+class OrderConfig(BaseModel):
+    """
+    Holds the configuration for the order process.
+
+    Attributes:
+        minOrderDatapointsForValidation (int): The minimum number of datapoints required so that an automatic validationset gets created if no recommended was found. Defaults to 50.
+        autoValidationSetSize (int): The maximum size of the auto-generated validation set. Defaults to 20.
+    """
+
+    autoValidationSetCreation: bool = Field(default=True)
+    minOrderDatapointsForValidation: int = Field(default=20)

rapidata/rapidata_client/config/rapidata_config.py

@@ -0,0 +1,36 @@
+from pydantic import BaseModel, Field
+
+from rapidata.rapidata_client.config.logging_config import LoggingConfig
+from rapidata.rapidata_client.config.order_config import OrderConfig
+from rapidata.rapidata_client.config.upload_config import UploadConfig
+
+
+class RapidataConfig(BaseModel):
+    """
+    Holds the configuration for the Rapidata client.
+
+    To adjust the configurations used, you can modify the `rapidata_config` object.
+
+    Attributes:
+        enableBetaFeatures (bool): Whether to enable beta features. Defaults to False.
+        upload (UploadConfig): The configuration for the upload process.
+            Such as the maximum number of worker threads for processing media paths and the maximum number of retries for failed uploads.
+        order (OrderConfig): The configuration for the order process.
+            Such as the minimum number of datapoints required so that an automatic validationset gets created if no recommended was found.
+        logging (LoggingConfig): The configuration for the logging process.
+            Such as the logging level and the logging file.
+
+    Example:
+        ```python
+        from rapidata import rapidata_config
+        rapidata_config.upload.maxUploadWorkers = 20
+        ```
+    """
+
+    enableBetaFeatures: bool = False
+    upload: UploadConfig = Field(default_factory=UploadConfig)
+    order: OrderConfig = Field(default_factory=OrderConfig)
+    logging: LoggingConfig = Field(default_factory=LoggingConfig)
+
+
+rapidata_config = RapidataConfig()

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]