rapidata 2.36.2__py3-none-any.whl → 2.38.0__py3-none-any.whl

rapidata/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "2.36.2"
+__version__ = "2.38.0"
 
 from .rapidata_client import (
     RapidataClient,
@@ -38,9 +38,8 @@ from .rapidata_client import (
     Box,
     Datapoint,
     PromptMetadata,
-    logger,
-    configure_logger,
-    RapidataOutputManager,
     FailedUploadException,
     rapidata_config,
+    logger,
+    managed_print,
 )

rapidata/rapidata_client/__init__.py

@@ -43,9 +43,6 @@ from .filter import (
     OrFilter,
     ResponseCountFilter,
 )
-
-from .logging import configure_logger, logger, RapidataOutputManager
-
 from .validation import Box
 from .exceptions import FailedUploadException
-from .config import rapidata_config
+from .config import rapidata_config, logger, managed_print

rapidata/rapidata_client/api/{rapidata_exception.py → rapidata_api_client.py}

@@ -1,10 +1,19 @@
 from typing import Optional, Any
-from rapidata.api_client.api_client import ApiClient, rest, ApiResponse, ApiResponseT
+from rapidata.api_client.api_client import (
+    ApiClient,
+    rest,
+    ApiResponse,
+    ApiResponseT,
+)
 from rapidata.api_client.exceptions import ApiException
 import json
 import threading
 from contextlib import contextmanager
-from rapidata.rapidata_client.logging import logger
+from rapidata.rapidata_client.config import logger, tracer
+from opentelemetry import trace
+from opentelemetry.trace import format_trace_id, format_span_id, Link, SpanContext
+from opentelemetry.sdk.trace.id_generator import RandomIdGenerator
+
 
 # Thread-local storage for controlling error logging
 _thread_local = threading.local()
@@ -91,6 +100,114 @@ class RapidataError(Exception):
 class RapidataApiClient(ApiClient):
     """Custom API client that wraps errors in RapidataError."""
 
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.id_generator = RandomIdGenerator()
+
+    def call_api(
+        self,
+        method,
+        url,
+        header_params=None,
+        body=None,
+        post_params=None,
+        _request_timeout=None,
+    ) -> rest.RESTResponse:
+        # Get the current span from OpenTelemetry
+        current_span = trace.get_current_span()
+
+        # Initialize header_params if it's None
+        if header_params is None:
+            header_params = {}
+
+        # Add tracing headers if we have a valid span
+        if not current_span.is_recording():
+            return super().call_api(
+                method,
+                url,
+                header_params,
+                body,
+                post_params,
+                _request_timeout,
+            )
+
+        current_span_context = current_span.get_span_context()
+
+        # Generate a new trace ID for backend communication
+        # This separates the backend trace from the SDK trace
+        backend_trace_id = self.id_generator.generate_trace_id()
+        backend_span_id = self.id_generator.generate_span_id()
+
+        # Create a new span context for the backend trace
+        backend_span_context = SpanContext(
+            trace_id=backend_trace_id,
+            span_id=backend_span_id,
+            is_remote=True,
+            trace_flags=current_span_context.trace_flags,
+        )
+
+        # Create a link from current SDK span to the backend trace
+        link_to_backend = Link(backend_span_context)
+
+        # Create a link from backend trace back to the original SDK span
+        link_back_to_sdk = Link(current_span_context)
+
+        # Create a span in the current SDK trace that links to the backend
+        with tracer.start_span(
+            f"sdk_request_{method}_{url.replace('/', '_')}",
+            links=[link_to_backend],
+        ) as sdk_request_span:
+            # Set attributes on the SDK span
+            sdk_request_span.set_attribute("http.method", method)
+            sdk_request_span.set_attribute("http.target", url)
+            sdk_request_span.set_attribute(
+                "rapidata.backend_trace_id", format_trace_id(backend_trace_id)
+            )
+            sdk_request_span.set_attribute(
+                "rapidata.original_trace_id",
+                format_trace_id(current_span_context.trace_id),
+            )
+
+            # Now create the initial span for the backend trace that will be sent
+            # This span will be the starting point for the backend trace
+            with tracer.start_span(
+                f"backend_trace_start_{method}_{url.replace('/', '_')}",
+                context=trace.set_span_in_context(
+                    trace.NonRecordingSpan(backend_span_context)
+                ),
+                links=[link_back_to_sdk],
+            ) as backend_initial_span:
+                # Set attributes on the backend initial span
+                backend_initial_span.set_attribute("http.method", method)
+                backend_initial_span.set_attribute("http.target", url)
+                backend_initial_span.set_attribute(
+                    "rapidata.trace_type", "backend_start"
+                )
+                backend_initial_span.set_attribute(
+                    "rapidata.sdk_trace_id",
+                    format_trace_id(current_span_context.trace_id),
+                )
+
+                # Format the traceparent header with the backend trace ID
+                # The backend will receive this and continue the trace
+                header_params["traceparent"] = (
+                    "00-"
+                    + format_trace_id(backend_trace_id)
+                    + "-"
+                    + format_span_id(backend_span_id)
+                    + "-"
+                    + f"{backend_span_context.trace_flags:02x}"
+                )
+
+                return super().call_api(
+                    method,
+                    url,
+                    header_params,
+                    body,
+                    post_params,
+                    _request_timeout,
+                )
+
     def response_deserialize(
         self,
         response_data: rest.RESTResponse,

rapidata/rapidata_client/benchmark/leaderboard/rapidata_leaderboard.py

@@ -1,12 +1,13 @@
+import urllib.parse
+import webbrowser
+from colorama import Fore
 import pandas as pd
 from typing import Literal, Optional
 
-from rapidata.rapidata_client.logging import logger
+from rapidata.rapidata_client.config import logger, managed_print, tracer
 from rapidata.rapidata_client.benchmark._detail_mapper import DetailMapper
 from rapidata.service.openapi_service import OpenAPIService
-from rapidata.api_client.models.update_leaderboard_response_config_model import (
-    UpdateLeaderboardResponseConfigModel,
-)
+from rapidata.api_client.models.update_leaderboard_model import UpdateLeaderboardModel
 
 
 class RapidataLeaderboard:
@@ -32,6 +33,7 @@ class RapidataLeaderboard:
         inverse_ranking: bool,
         response_budget: int,
         min_responses_per_matchup: int,
+        benchmark_id: str,
         id: str,
         openapi_service: OpenAPIService,
     ):
@@ -43,7 +45,9 @@ class RapidataLeaderboard:
         self.__inverse_ranking = inverse_ranking
         self.__response_budget = response_budget
         self.__min_responses_per_matchup = min_responses_per_matchup
+        self.__benchmark_id = benchmark_id
         self.id = id
+        self.__leaderboard_page = f"https://app.{self.__openapi_service.environment}/mri/benchmarks/{self.__benchmark_id}/leaderboard/{self.id}"
 
     @property
     def level_of_detail(self) -> Literal["low", "medium", "high", "very high"]:
@@ -59,15 +63,10 @@ class RapidataLeaderboard:
         """
         Sets the level of detail of the leaderboard.
         """
-        logger.debug(f"Setting level of detail to {level_of_detail}")
-        self.__openapi_service.leaderboard_api.leaderboard_leaderboard_id_response_config_put(
-            leaderboard_id=self.id,
-            update_leaderboard_response_config_model=UpdateLeaderboardResponseConfigModel(
-                responseBudget=DetailMapper.get_budget(level_of_detail),
-                minResponses=self.__min_responses_per_matchup,
-            ),
-        )
-        self.__response_budget = DetailMapper.get_budget(level_of_detail)
+        with tracer.start_as_current_span("RapidataLeaderboard.level_of_detail.setter"):
+            logger.debug(f"Setting level of detail to {level_of_detail}")
+            self.__response_budget = DetailMapper.get_budget(level_of_detail)
+            self._update_config()
 
     @property
     def min_responses_per_matchup(self) -> int:
@@ -81,23 +80,20 @@ class RapidataLeaderboard:
         """
         Sets the minimum number of responses required to be considered for the leaderboard.
         """
-        if not isinstance(min_responses, int):
-            raise ValueError("Min responses per matchup must be an integer")
+        with tracer.start_as_current_span(
+            "RapidataLeaderboard.min_responses_per_matchup.setter"
+        ):
+            if not isinstance(min_responses, int):
+                raise ValueError("Min responses per matchup must be an integer")
 
-        if min_responses < 3:
-            raise ValueError("Min responses per matchup must be at least 3")
+            if min_responses < 3:
+                raise ValueError("Min responses per matchup must be at least 3")
 
-        logger.debug(
-            f"Setting min responses per matchup to {min_responses} for leaderboard {self.name}"
-        )
-        self.__openapi_service.leaderboard_api.leaderboard_leaderboard_id_response_config_put(
-            leaderboard_id=self.id,
-            update_leaderboard_response_config_model=UpdateLeaderboardResponseConfigModel(
-                responseBudget=self.__response_budget,
-                minResponses=min_responses,
-            ),
-        )
-        self.__min_responses_per_matchup = min_responses
+            logger.debug(
+                f"Setting min responses per matchup to {min_responses} for leaderboard {self.name}"
+            )
+            self.__min_responses_per_matchup = min_responses
+            self._update_config()
 
     @property
     def show_prompt_asset(self) -> bool:
@@ -134,6 +130,20 @@ class RapidataLeaderboard:
         """
         return self.__name
 
+    @name.setter
+    def name(self, name: str):
+        """
+        Sets the name of the leaderboard.
+        """
+        with tracer.start_as_current_span("RapidataLeaderboard.name.setter"):
+            if not isinstance(name, str):
+                raise ValueError("Name must be a string")
+            if len(name) < 1:
+                raise ValueError("Name must be at least 1 character long")
+
+            self.__name = name
+            self._update_config()
+
     def get_standings(self, tags: Optional[list[str]] = None) -> pd.DataFrame:
         """
         Returns the standings of the leaderboard.
@@ -146,27 +156,59 @@ class RapidataLeaderboard:
         Returns:
             A pandas DataFrame containing the standings of the leaderboard.
         """
+        with tracer.start_as_current_span("RapidataLeaderboard.get_standings"):
+            participants = self.__openapi_service.leaderboard_api.leaderboard_leaderboard_id_standings_get(
+                leaderboard_id=self.id, tags=tags
+            )
 
-        participants = self.__openapi_service.leaderboard_api.leaderboard_leaderboard_id_standings_get(
-            leaderboard_id=self.id, tags=tags
-        )
-
-        standings = []
-        for participant in participants.items:
-            standings.append(
-                {
-                    "name": participant.name,
-                    "wins": participant.wins,
-                    "total_matches": participant.total_matches,
-                    "score": (
-                        round(participant.score, 2)
-                        if participant.score is not None
-                        else None
-                    ),
-                }
+            standings = []
+            for participant in participants.items:
+                standings.append(
+                    {
+                        "name": participant.name,
+                        "wins": participant.wins,
+                        "total_matches": participant.total_matches,
+                        "score": (
+                            round(participant.score, 2)
+                            if participant.score is not None
+                            else None
+                        ),
+                    }
+                )
+
+            return pd.DataFrame(standings)
+
+    def view(self) -> None:
+        """
+        Views the leaderboard.
+        """
+        logger.info("Opening leaderboard page in browser...")
+        could_open_browser = webbrowser.open(self.__leaderboard_page)
+        if not could_open_browser:
+            encoded_url = urllib.parse.quote(
+                self.__leaderboard_page, safe="%/:=&?~#+!$,;'@()*[]"
+            )
+            managed_print(
+                Fore.RED
+                + f"Please open this URL in your browser: '{encoded_url}'"
+                + Fore.RESET
             )
 
-        return pd.DataFrame(standings)
+    def _custom_config(self, response_budget: int, min_responses_per_matchup: int):
+        self.__response_budget = response_budget
+        self.__min_responses_per_matchup = min_responses_per_matchup
+        self._update_config()
+
+    def _update_config(self):
+        with tracer.start_as_current_span("RapidataLeaderboard._update_config"):
+            self.__openapi_service.leaderboard_api.leaderboard_leaderboard_id_patch(
+                leaderboard_id=self.id,
+                update_leaderboard_model=UpdateLeaderboardModel(
+                    name=self.__name,
+                    responseBudget=self.__response_budget,
+                    minResponses=self.__min_responses_per_matchup,
+                ),
+            )
 
     def __str__(self) -> str:
         return f"RapidataLeaderboard(name={self.name}, instruction={self.instruction}, show_prompt={self.show_prompt}, leaderboard_id={self.id})"

rapidata/rapidata_client/benchmark/participant/_participant.py

@@ -3,15 +3,17 @@ import time
 from tqdm import tqdm
 
 from rapidata.rapidata_client.datapoints.assets import MediaAsset
-from rapidata.rapidata_client.logging import logger
-from rapidata.rapidata_client.logging.output_manager import RapidataOutputManager
+from rapidata.rapidata_client.config import logger
 from rapidata.api_client.models.create_sample_model import CreateSampleModel
 from rapidata.service.openapi_service import OpenAPIService
 from rapidata.rapidata_client.config.rapidata_config import rapidata_config
-from rapidata.rapidata_client.api.rapidata_exception import (
+from rapidata.rapidata_client.api.rapidata_api_client import (
     suppress_rapidata_error_logging,
 )
 
+# Add OpenTelemetry context imports for thread propagation
+from opentelemetry import context as otel_context
+
 
 class BenchmarkParticipant:
     def __init__(self, name: str, id: str, openapi_service: OpenAPIService):
@@ -42,7 +44,7 @@ class BenchmarkParticipant:
             urls = [asset.path]
 
         last_exception = None
-        for attempt in range(rapidata_config.uploadMaxRetries):
+        for attempt in range(rapidata_config.upload.maxRetries):
             try:
                 with suppress_rapidata_error_logging():
                     self.__openapi_service.participant_api.participant_participant_id_sample_post(
@@ -56,7 +58,7 @@ class BenchmarkParticipant:
 
             except Exception as e:
                 last_exception = e
-                if attempt < rapidata_config.uploadMaxRetries - 1:
+                if attempt < rapidata_config.upload.maxRetries - 1:
                     # Exponential backoff: wait 1s, then 2s, then 4s
                     retry_delay = 2**attempt
                     time.sleep(retry_delay)
@@ -64,7 +66,7 @@ class BenchmarkParticipant:
                     logger.debug(
                         "Retrying %s of %s...",
                         attempt + 1,
-                        rapidata_config.uploadMaxRetries,
+                        rapidata_config.upload.maxRetries,
                     )
 
         logger.error(f"Upload failed for {identifier}. Error: {str(last_exception)}")
@@ -85,16 +87,31 @@ class BenchmarkParticipant:
         Returns:
             tuple[list[str], list[str]]: Lists of successful and failed identifiers
         """
+
+        def upload_with_context(
+            context: otel_context.Context, asset: MediaAsset, identifier: str
+        ) -> tuple[MediaAsset | None, MediaAsset | None]:
+            """Wrapper function that runs _process_single_sample_upload with the provided context."""
+            token = otel_context.attach(context)
+            try:
+                return self._process_single_sample_upload(asset, identifier)
+            finally:
+                otel_context.detach(token)
+
         successful_uploads: list[MediaAsset] = []
         failed_uploads: list[MediaAsset] = []
         total_uploads = len(assets)
 
+        # Capture the current OpenTelemetry context before creating threads
+        current_context = otel_context.get_current()
+
         with ThreadPoolExecutor(
-            max_workers=rapidata_config.maxUploadWorkers
+            max_workers=rapidata_config.upload.maxWorkers
         ) as executor:
             futures = [
                 executor.submit(
-                    self._process_single_sample_upload,
+                    upload_with_context,
+                    current_context,
                     asset,
                     identifier,
                 )
@@ -104,7 +121,7 @@ class BenchmarkParticipant:
             with tqdm(
                 total=total_uploads,
                 desc="Uploading media",
-                disable=RapidataOutputManager.silent_mode,
+                disable=rapidata_config.logging.silent_mode,
             ) as pbar:
                 for future in as_completed(futures):
                     try: