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

rapidata/rapidata_client/rapidata_client.py

@@ -1,32 +1,48 @@
+import json
+from typing import Any
 import requests
 from packaging import version
 from rapidata import __version__
-
+import uuid
+import random
 from rapidata.service.openapi_service import OpenAPIService
 
+from rapidata.rapidata_client.benchmark.rapidata_benchmark_manager import (
+    RapidataBenchmarkManager,
+)
+from rapidata.rapidata_client.audience.rapidata_audience_manager import (
+    RapidataAudienceManager,
+)
 from rapidata.rapidata_client.order.rapidata_order_manager import RapidataOrderManager
-from rapidata.rapidata_client.benchmark.rapidata_benchmark_manager import RapidataBenchmarkManager
-
 from rapidata.rapidata_client.validation.validation_set_manager import (
     ValidationSetManager,
 )
 
 from rapidata.rapidata_client.demographic.demographic_manager import DemographicManager
 
-from rapidata.rapidata_client.logging import logger, managed_print
+from rapidata.rapidata_client.config import (
+    logger,
+    tracer,
+    managed_print,
+    rapidata_config,
+)
+
+from rapidata.rapidata_client.datapoints._asset_uploader import AssetUploader
+from rapidata.rapidata_client.job.job_manager import JobManager
+
 
 class RapidataClient:
     """The Rapidata client is the main entry point for interacting with the Rapidata API. It allows you to create orders and validation sets."""
 
     def __init__(
-            self,
-            client_id: str | None = None,
-            client_secret: str | None = None,
-            environment: str = "rapidata.ai",
-            oauth_scope: str = "openid",
-            cert_path: str | None = None,
-            token: dict | None = None,
-            leeway: int = 60
+        self,
+        client_id: str | None = None,
+        client_secret: str | None = None,
+        environment: str = "rapidata.ai",
+        oauth_scope: str = "openid roles",
+        cert_path: str | None = None,
+        token: dict | None = None,
+        leeway: int = 60,
     ):
         """Initialize the RapidataClient. If both the client_id and client_secret are None, it will try using your credentials under "~/.config/rapidata/credentials.json".
         If this is not successful, it will open a browser window and ask you to log in, then save your new credentials in said json file.
@@ -43,48 +59,110 @@ class RapidataClient:
         Attributes:
             order (RapidataOrderManager): The RapidataOrderManager instance.
             validation (ValidationSetManager): The ValidationSetManager instance.
+            audience (RapidataAudienceManager): The RapidataAudienceManager instance.
+            job (JobManager): The JobManager instance.
+            mri (RapidataBenchmarkManager): The RapidataBenchmarkManager instance.
         """
-        logger.debug("Checking version")
-        self._check_version()
-
-        logger.debug("Initializing OpenAPIService")
-        self._openapi_service = OpenAPIService(
-            client_id=client_id,
-            client_secret=client_secret,
-            environment=environment,
-            oauth_scope=oauth_scope,
-            cert_path=cert_path,
-            token=token,
-            leeway=leeway,
+        tracer.set_session_id(
+            uuid.UUID(int=random.Random().getrandbits(128), version=4).hex
         )
 
-        logger.debug("Initializing RapidataOrderManager")
-        self.order = RapidataOrderManager(openapi_service=self._openapi_service)
+        with tracer.start_as_current_span("RapidataClient.__init__"):
+            logger.debug("Checking version")
+            self._check_version()
+            if environment != "rapidata.ai":
+                rapidata_config.logging.enable_otlp = False
+
+            logger.debug("Initializing OpenAPIService")
+            self._openapi_service = OpenAPIService(
+                client_id=client_id,
+                client_secret=client_secret,
+                environment=environment,
+                oauth_scope=oauth_scope,
+                cert_path=cert_path,
+                token=token,
+                leeway=leeway,
+            )
+
+            self._asset_uploader = AssetUploader(openapi_service=self._openapi_service)
+
+            logger.debug("Initializing RapidataOrderManager")
+            self.order = RapidataOrderManager(openapi_service=self._openapi_service)
+
+            logger.debug("Initializing ValidationSetManager")
+            self.validation = ValidationSetManager(openapi_service=self._openapi_service)
 
-        logger.debug("Initializing ValidationSetManager")
-        self.validation = ValidationSetManager(openapi_service=self._openapi_service)
+            logger.debug("Initializing JobManager")
+            self.job = JobManager(openapi_service=self._openapi_service)
 
-        logger.debug("Initializing DemographicManager")
-        self._demographic = DemographicManager(openapi_service=self._openapi_service)
+            logger.debug("Initializing RapidataBenchmarkManager")
+            self.mri = RapidataBenchmarkManager(openapi_service=self._openapi_service)
+
+            logger.debug("Initializing RapidataAudienceManager")
+            self.audience = RapidataAudienceManager(
+                openapi_service=self._openapi_service
+            )
+
+            logger.debug("Initializing RapidataDemographicManager")
+            self._demographic = DemographicManager(
+                openapi_service=self._openapi_service
+            )
+
+        self._check_beta_features()  # can't be in the trace for some reason
 
-        logger.debug("Initializing RapidataBenchmarkManager")
-        self.mri = RapidataBenchmarkManager(openapi_service=self._openapi_service)
-        
     def reset_credentials(self):
         """Reset the credentials saved in the configuration file for the current environment."""
+        logger.info("Resetting credentials")
         self._openapi_service.reset_credentials()
+        logger.info("Credentials reset")
+
+    def clear_all_caches(self):
+        """Clear all caches for the client."""
+        self._asset_uploader.clear_cache()
+        logger.info("All caches cleared")
+
+    def _check_beta_features(self):
+        """Enable beta features for the client."""
+        with tracer.start_as_current_span("RapidataClient.check_beta_features"):
+            result: dict[str, Any] = json.loads(
+                self._openapi_service.api_client.call_api(
+                    "GET",
+                    f"https://auth.{self._openapi_service.environment}/connect/userinfo",
+                )
+                .read()
+                .decode("utf-8")
+            )
+            logger.debug("Userinfo: %s", result)
+            if "Admin" not in result.get("role", []):
+                logger.debug("User is not an admin, not enabling beta features")
+                return
+
+            logger.debug("User is an admin, enabling beta features")
+            rapidata_config.enableBetaFeatures = True
 
     def _check_version(self):
         try:
             response = requests.get(
                 "https://api.github.com/repos/RapidataAI/rapidata-python-sdk/releases/latest",
                 headers={"Accept": "application/vnd.github.v3+json"},
-                timeout=3
+                timeout=3,
             )
             if response.status_code == 200:
                 latest_version = response.json()["tag_name"].lstrip("v")
                 if version.parse(latest_version) > version.parse(__version__):
-                    managed_print(f"""A new version of the Rapidata SDK is available: {latest_version}
-Your current version is: {__version__}""")
+                    managed_print(
+                        f"""A new version of the Rapidata SDK is available: {latest_version}
+Your current version is: {__version__}"""
+                    )
+                else:
+                    logger.debug(
+                        "Current version is up to date. Version: %s", __version__
+                    )
         except Exception as e:
-            logger.debug(f"Failed to check for updates: {e}")
+            logger.debug("Failed to check for updates: %s", e)
+
+    def __str__(self) -> str:
+        return f"RapidataClient(environment={self._openapi_service.environment})"
+
+    def __repr__(self) -> str:
+        return self.__str__()

rapidata/rapidata_client/referee/__init__.py

@@ -1,3 +1,3 @@
 from ._base_referee import Referee
-from ._naive_referee import NaiveReferee #as MaxVoteReferee
+from ._naive_referee import NaiveReferee  # as MaxVoteReferee
 from ._early_stopping_referee import EarlyStoppingReferee

rapidata/rapidata_client/referee/_base_referee.py

@@ -1,10 +1,15 @@
 from abc import ABC, abstractmethod
-from typing import Any, Mapping
+from typing import Mapping
+import json
+
+from rapidata.api_client.models.i_referee_model import IRefereeModel
+
 
 class Referee(ABC):
     """
-    The referee defines when a rapid is considered complete. 
+    The referee defines when a rapid is considered complete.
     """
+
     @abstractmethod
     def _to_dict(self) -> Mapping[str, str | int | float]:
         """
@@ -13,8 +18,14 @@ class Referee(ABC):
         pass
 
     @abstractmethod
-    def _to_model(self) -> Any:
+    def _to_model(self) -> IRefereeModel:
         """
         Convert the referee to a referee configuration model.
         """
         pass
+
+    def __str__(self) -> str:
+        return json.dumps(self._to_dict())
+
+    def __repr__(self) -> str:
+        return self.__str__()

rapidata/rapidata_client/referee/_early_stopping_referee.py

@@ -1,8 +1,10 @@
-from typing import Any
+from __future__ import annotations
+
 from rapidata.rapidata_client.referee._base_referee import Referee
-from rapidata.api_client.models.early_stopping_referee_model import (
-    EarlyStoppingRefereeModel,
-)
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from rapidata.api_client.models.i_referee_model import IRefereeModel
 
 
 class EarlyStoppingReferee(Referee):
@@ -15,39 +17,52 @@ class EarlyStoppingReferee(Referee):
     The threshold behaves logarithmically, meaning small increments (e.g., from 0.99
     to 0.999) can significantly impact the stopping criteria.
 
-    This referee is supported for the classification and compare tasks (in compare, 
+    This referee is supported for the classification and compare tasks (in compare,
     the two options are treated as the categories).
 
     Args:
         threshold (float, optional): The confidence threshold for early stopping.
             Defaults to 0.999.
-        max_vote_count (int, optional): The maximum number of votes allowed
+        max_responses (int, optional): The maximum number of votes allowed
             before stopping. Defaults to 100.
 
     Attributes:
         threshold (float): The confidence threshold for early stopping.
-        max_vote_count (int): The maximum number of votes allowed before stopping.
+        max_responses (int): The maximum number of votes allowed before stopping.
     """
 
-    def __init__(self, threshold: float = 0.999, max_vote_count: int = 100):
+    def __init__(self, threshold: float = 0.999, max_responses: int = 100):
         if threshold <= 0 or threshold >= 1:
             raise ValueError("The threshold must be between 0 and 1.")
-        if max_vote_count < 1:
+        if max_responses < 1:
             raise ValueError("The number of responses must be greater than 0.")
-        
+
         self.threshold = threshold
-        self.max_vote_count = max_vote_count
+        self.max_responses = max_responses
 
     def _to_dict(self):
         return {
             "_t": "ProbabilisticAttachCategoryRefereeConfig",
             "threshold": self.threshold,
-            "maxVotes": self.max_vote_count,
+            "maxVotes": self.max_responses,
         }
 
-    def _to_model(self) -> Any:
-        return EarlyStoppingRefereeModel(
-            _t="EarlyStoppingReferee",
-            threshold=self.threshold,
-            maxVotes=self.max_vote_count,
+    def _to_model(self) -> IRefereeModel:
+        from rapidata.api_client.models.i_referee_model_early_stopping_referee_model import (
+            IRefereeModelEarlyStoppingRefereeModel,
         )
+        from rapidata.api_client.models.i_referee_model import IRefereeModel
+
+        return IRefereeModel(
+            actual_instance=IRefereeModelEarlyStoppingRefereeModel(
+                _t="EarlyStoppingReferee",
+                threshold=self.threshold,
+                maxVotes=self.max_responses,
+            )
+        )
+
+    def __str__(self) -> str:
+        return f"EarlyStoppingReferee(threshold={self.threshold}, max_responses={self.max_responses})"
+
+    def __repr__(self) -> str:
+        return self.__str__()

rapidata/rapidata_client/referee/_naive_referee.py

@@ -1,5 +1,10 @@
-from rapidata.api_client.models.naive_referee_model import NaiveRefereeModel
+from __future__ import annotations
+
 from rapidata.rapidata_client.referee._base_referee import Referee
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from rapidata.api_client.models.i_referee_model import IRefereeModel
 
 
 class NaiveReferee(Referee):
@@ -29,5 +34,21 @@ class NaiveReferee(Referee):
             "guessesRequired": self.responses,
         }
 
-    def _to_model(self):
-        return NaiveRefereeModel(_t="NaiveReferee", totalVotes=self.responses)
+    def _to_model(self) -> IRefereeModel:
+        from rapidata.api_client.models.i_referee_model_naive_referee_model import (
+            IRefereeModelNaiveRefereeModel,
+        )
+        from rapidata.api_client.models.i_referee_model import IRefereeModel
+
+        return IRefereeModel(
+            actual_instance=IRefereeModelNaiveRefereeModel(
+                _t="NaiveReferee",
+                totalVotes=self.responses,
+            )
+        )
+
+    def __str__(self) -> str:
+        return f"NaiveReferee(responses={self.responses})"
+
+    def __repr__(self) -> str:
+        return self.__str__()

rapidata/rapidata_client/results/__init__.py

@@ -0,0 +1,3 @@
+from .rapidata_results import RapidataResults
+
+__all__ = ["RapidataResults"]

rapidata/rapidata_client/{order → results}/rapidata_results.py

@@ -1,57 +1,69 @@
-import pandas as pd
-from typing import Any
-from pandas.core.indexes.base import Index
+from typing import TYPE_CHECKING, Any
+from rapidata.rapidata_client.config import managed_print
 import json
-from rapidata.rapidata_client.logging import managed_print
+
+if TYPE_CHECKING:
+    import pandas as pd
+    from pandas.core.indexes.base import Index
+
 
 class RapidataResults(dict):
     """
     A specialized dictionary class for handling Rapidata API results.
     Extends the built-in dict class with specialized methods.
     """
-    def to_pandas(self, split_details: bool = False) -> pd.DataFrame:
+
+    def to_pandas(self, split_details: bool = False) -> "pd.DataFrame":
         """
         Warning:
             This method is currently under development. The structure of the results may change in the future.
 
         Converts the results to a pandas DataFrame.
-        
+
         For Compare results, creates standardized A/B columns for metrics.
         For regular results, flattens nested dictionaries into columns with underscore-separated names.
-        
+
         Args:
             split_details: If True, splits each datapoint by its detailed results,
                           creating a row for each response with global metrics copied.
-                          
+
         Returns:
             pd.DataFrame: A DataFrame containing the processed results
-            
+
         Raises:
             ValueError: If split_details is True but no detailed results are found
         """
+        import pandas as pd
+        from pandas.core.indexes.base import Index
+
         if "results" not in self or not self["results"]:
             return pd.DataFrame()
-        
+
         if self["info"].get("orderType") is None:
-            managed_print("Warning: Results are old and Order type is not specified. Dataframe might be wrong.")
-        
+            managed_print(
+                "Warning: Results are old and Order type is not specified. Dataframe might be wrong."
+            )
+
         # Check for detailed results if split_details is True
         if split_details:
             if not self._has_detailed_results():
                 raise ValueError("No detailed results found in the data")
             return self._to_pandas_with_detailed_results()
-            
-        if self["info"].get("orderType") == "Compare" or self["info"].get("orderType") == "Ranking":
+
+        if (
+            self["info"].get("orderType") == "Compare"
+            or self["info"].get("orderType") == "Ranking"
+        ):
             return self._compare_to_pandas()
-        
+
         # Get the structure from first item
         first_item = self["results"][0]
         columns = []
         path_map = {}  # Maps flattened column names to paths to reach the values
-        
+
         # Build the column structure once
         self._build_column_structure(first_item, columns, path_map)
-        
+
         # Extract data using the known structure
         data = []
         for item in self["results"]:
@@ -60,82 +72,86 @@ class RapidataResults(dict):
                 value = self._get_value_from_path(item, path)
                 row.append(value)
             data.append(row)
-            
+
         return pd.DataFrame(data, columns=Index(columns))
-    
+
     def _has_detailed_results(self) -> bool:
         """
         Checks if the results contain detailed results.
-        
+
         Returns:
             bool: True if detailed results exist, False otherwise
         """
         if not self.get("results"):
             return False
-        
+
         first_result = self["results"][0]
-        return "detailedResults" in first_result and isinstance(first_result["detailedResults"], list)
-    
-    def _to_pandas_with_detailed_results(self) -> pd.DataFrame:
+        return "detailedResults" in first_result and isinstance(
+            first_result["detailedResults"], list
+        )
+
+    def _to_pandas_with_detailed_results(self) -> "pd.DataFrame":
         """
         Converts results to a pandas DataFrame with detailed results split into separate rows.
-        
+
         Returns:
             pd.DataFrame: A DataFrame with one row per detailed result
         """
+        import pandas as pd
+
         rows = []
-        
+
         for result in self["results"]:
             # Get all non-detailed results fields
             base_data = {k: v for k, v in result.items() if k != "detailedResults"}
-            
+
             # Process each detailed result
             for detailed_result in result["detailedResults"]:
                 row = base_data.copy()  # Copy base data for each detailed result
-                
+
                 # Add flattened detailed result data
                 flattened = self._flatten_dict(detailed_result)
                 for key, value in flattened.items():
                     row[key] = value
-                
+
                 rows.append(row)
-        
+
         return pd.DataFrame(rows)
-    
-    def _flatten_dict(self, d: dict[str, Any], parent_key: str = '') -> dict[str, Any]:
+
+    def _flatten_dict(self, d: dict[str, Any], parent_key: str = "") -> dict[str, Any]:
         """
         Flattens a nested dictionary into a single-level dictionary with underscore-separated keys.
-        
+
         Args:
             d: The dictionary to flatten
             parent_key: The parent key for nested dictionaries
-            
+
         Returns:
             dict: A flattened dictionary
         """
         items: list[tuple[str, Any]] = []
-        
+
         for key, value in d.items():
             new_key = f"{parent_key}_{key}" if parent_key else key
-            
+
             if isinstance(value, dict):
                 items.extend(self._flatten_dict(value, new_key).items())
             else:
                 items.append((new_key, value))
-                
+
         return dict(items)
 
     def _build_column_structure(
-        self, 
-        d: dict[str, Any], 
-        columns: list[str], 
-        path_map: dict[str, list[str]], 
-        parent_key: str = '', 
-        current_path: list[str] | None = None
+        self,
+        d: dict[str, Any],
+        columns: list[str],
+        path_map: dict[str, list[str]],
+        parent_key: str = "",
+        current_path: list[str] | None = None,
     ) -> None:
         """
         Builds the column structure and paths to reach values in nested dictionaries.
-        
+
         Args:
             d: The dictionary to analyze
             columns: List to store column names
@@ -145,25 +161,27 @@ class RapidataResults(dict):
         """
         if current_path is None:
             current_path = []
-            
+
         for key, value in d.items():
             new_key = f"{parent_key}_{key}" if parent_key else key
             new_path: list[str] = current_path + [key]
-            
+
             if isinstance(value, dict):
-                self._build_column_structure(value, columns, path_map, new_key, new_path)
+                self._build_column_structure(
+                    value, columns, path_map, new_key, new_path
+                )
             else:
                 columns.append(new_key)
                 path_map[new_key] = new_path
-    
+
     def _get_value_from_path(self, d: dict[str, Any], path: list[str]) -> Any:
         """
         Retrieves a value from a nested dictionary using a path list.
-        
+
         Args:
             d: The dictionary to retrieve the value from
             path: List of keys forming the path to the desired value
-            
+
         Returns:
             The value at the specified path, or None if the path doesn't exist
         """
@@ -171,10 +189,12 @@ class RapidataResults(dict):
             d = d.get(key, {})
         return d.get(path[-1])
 
-    def _compare_to_pandas(self) -> pd.DataFrame:
+    def _compare_to_pandas(self) -> "pd.DataFrame":
         """
         Converts Compare results to a pandas DataFrame dynamically.
         """
+        import pandas as pd
+
         if not self.get("results"):
             return pd.DataFrame()
 
@@ -190,10 +210,11 @@ class RapidataResults(dict):
                 continue
 
             assets = [asset for asset in assets if asset not in ["Both", "Neither"]]
-            
+
             # Initialize row with non-comparative fields
             row = {
-                key: value for key, value in result.items() 
+                key: value
+                for key, value in result.items()
                 if not isinstance(value, dict)
             }
             row["assetA"] = assets[0]
@@ -203,26 +224,28 @@ class RapidataResults(dict):
             for key, values in result.items():
                 if isinstance(values, dict) and len(values) >= 2:
                     # Add main asset columns
-                    for i, asset in enumerate(assets[:2]):  # Limit to first 2 main assets
+                    for i, asset in enumerate(
+                        assets[:2]
+                    ):  # Limit to first 2 main assets
                         column_prefix = "A_" if i == 0 else "B_"
-                        row[f'{column_prefix}{key}'] = values.get(asset, 0)
-                    
+                        row[f"{column_prefix}{key}"] = values.get(asset, 0)
+
                     # Add special option columns if they exist
                     if "Both" in values:
-                        row[f'Both_{key}'] = values.get("Both", 0)
+                        row[f"Both_{key}"] = values.get("Both", 0)
                     if "Neither" in values:
-                        row[f'Neither_{key}'] = values.get("Neither", 0)
-                        
+                        row[f"Neither_{key}"] = values.get("Neither", 0)
+
             rows.append(row)
-            
+
         return pd.DataFrame(rows)
 
-    def to_json(self, path: str="./results.json") -> None:
+    def to_json(self, path: str = "./results.json") -> None:
         """
         Saves the results to a JSON file.
-        
+
         Args:
             path: The file path where the JSON should be saved. Defaults to "./results.json".
         """
-        with open(path, 'w') as f:
+        with open(path, "w") as f:
             json.dump(self, f)

rapidata/rapidata_client/selection/__init__.py

@@ -7,5 +7,5 @@ from .capped_selection import CappedSelection
 from .shuffling_selection import ShufflingSelection
 from .ab_test_selection import AbTestSelection
 from .static_selection import StaticSelection
-from .retrieval_modes import RetrievalMode
+from .rapidata_retrieval_modes import RapidataRetrievalMode
 from .effort_selection import EffortSelection