rapidata 1.8.3__py3-none-any.whl → 2.44.4__py3-none-any.whl

rapidata/rapidata_client/workflow/_select_words_workflow.py

@@ -0,0 +1,74 @@
+from rapidata.api_client.models.i_order_workflow_model import IOrderWorkflowModel
+from rapidata.api_client.models.i_order_workflow_model_simple_workflow_model import (
+    IOrderWorkflowModelSimpleWorkflowModel,
+)
+from rapidata.api_client.models.i_rapid_blueprint import IRapidBlueprint
+from rapidata.api_client.models.i_rapid_blueprint_transcription_rapid_blueprint import (
+    IRapidBlueprintTranscriptionRapidBlueprint,
+)
+from rapidata.rapidata_client.workflow._base_workflow import Workflow
+from rapidata.api_client.models.i_rapid_payload_transcription_payload import (
+    IRapidPayloadTranscriptionPayload,
+)
+from rapidata.api_client.models.transcription_word import TranscriptionWord
+from rapidata.rapidata_client.datapoints._datapoint import Datapoint
+from rapidata.api_client.models.rapid_modality import RapidModality
+from rapidata.api_client.models.i_rapid_payload import IRapidPayload
+
+
+class SelectWordsWorkflow(Workflow):
+    """
+    A workflow for select words tasks.
+
+    This class represents a select words workflow
+    where datapoints have a sentence attached to them where words can be selected.
+
+    Attributes:
+        _instruction (str): The instruction for the select words task.
+
+    Args:
+        instruction (str): The instruction to be provided for the select words task.
+    """
+
+    modality = RapidModality.TRANSCRIPTION
+
+    def __init__(self, instruction: str):
+        super().__init__(type="SimpleWorkflowConfig")
+        self._instruction = instruction
+
+    def _get_instruction(self) -> str:
+        return self._instruction
+
+    def _to_model(self) -> IOrderWorkflowModel:
+        blueprint = IRapidBlueprintTranscriptionRapidBlueprint(
+            _t="TranscriptionBlueprint", title=self._instruction
+        )
+
+        return IOrderWorkflowModel(
+            actual_instance=IOrderWorkflowModelSimpleWorkflowModel(
+                _t="SimpleWorkflow",
+                blueprint=IRapidBlueprint(actual_instance=blueprint),
+            )
+        )
+
+    def _to_payload(self, datapoint: Datapoint) -> IRapidPayload:
+        assert (
+            datapoint.sentence is not None
+        ), "SelectWordsWorkflow requires a sentence datapoint"
+
+        return IRapidPayload(
+            actual_instance=IRapidPayloadTranscriptionPayload(
+                _t="TranscriptionPayload",
+                title=self._instruction,
+                transcription=[
+                    TranscriptionWord(word=word, wordIndex=i)
+                    for i, word in enumerate(datapoint.sentence.split())
+                ],
+            )
+        )
+
+    def __str__(self) -> str:
+        return f"SelectWordsWorkflow(instruction='{self._instruction}')"
+
+    def __repr__(self) -> str:
+        return f"SelectWordsWorkflow(instruction={self._instruction!r})"

rapidata/rapidata_client/workflow/_timestamp_workflow.py

@@ -0,0 +1,65 @@
+from rapidata.api_client.models.i_order_workflow_model import IOrderWorkflowModel
+from rapidata.api_client.models.i_order_workflow_model_simple_workflow_model import (
+    IOrderWorkflowModelSimpleWorkflowModel,
+)
+from rapidata.api_client.models.i_rapid_blueprint import IRapidBlueprint
+from rapidata.api_client.models.i_rapid_blueprint_scrub_rapid_blueprint import (
+    IRapidBlueprintScrubRapidBlueprint,
+)
+from rapidata.rapidata_client.workflow._base_workflow import Workflow
+from rapidata.api_client.models.i_rapid_payload_scrub_payload import (
+    IRapidPayloadScrubPayload,
+)
+from rapidata.rapidata_client.datapoints._datapoint import Datapoint
+from rapidata.api_client.models.rapid_modality import RapidModality
+from rapidata.api_client.models.i_rapid_payload import IRapidPayload
+
+
+class TimestampWorkflow(Workflow):
+    """
+    A workflow for timestamp tasks.
+
+    This class represents a timestamp workflow
+    where audio or video content receives timestamps based on the instruction.
+
+    Attributes:
+        _instruction (str): The instruction for the timestamp task.
+
+    Args:
+        instruction (str): The instruction to be provided for the timestamp task.
+    """
+
+    modality = RapidModality.SCRUB
+
+    def __init__(self, instruction: str):
+        super().__init__(type="SimpleWorkflowConfig")
+        self._instruction = instruction
+
+    def _get_instruction(self) -> str:
+        return self._instruction
+
+    def _to_model(self) -> IOrderWorkflowModel:
+        blueprint = IRapidBlueprintScrubRapidBlueprint(
+            _t="ScrubBlueprint", target=self._instruction
+        )
+
+        return IOrderWorkflowModel(
+            actual_instance=IOrderWorkflowModelSimpleWorkflowModel(
+                _t="SimpleWorkflow",
+                blueprint=IRapidBlueprint(actual_instance=blueprint),
+            )
+        )
+
+    def _to_payload(self, datapoint: Datapoint) -> IRapidPayload:
+        return IRapidPayload(
+            actual_instance=IRapidPayloadScrubPayload(
+                _t="ScrubPayload",
+                target=self._instruction,
+            )
+        )
+
+    def __str__(self) -> str:
+        return f"TimestampWorkflow(instruction='{self._instruction}')"
+
+    def __repr__(self) -> str:
+        return f"TimestampWorkflow(instruction={self._instruction!r})"

rapidata/service/__init__.py

@@ -1 +0,0 @@
-from .local_file_service import LocalFileService as LocalFileService

rapidata/service/credential_manager.py

@@ -1,6 +1,7 @@
 import json
 import os
 import time
+import urllib.parse
 import webbrowser
 from datetime import datetime, timezone
 from pathlib import Path
@@ -8,7 +9,9 @@ from socket import gethostname
 from typing import Dict, List, Optional, Tuple
 
 import requests
+from colorama import Fore
 from pydantic import BaseModel
+from rapidata.rapidata_client.config import logger, managed_print
 
 
 class ClientCredential(BaseModel):
@@ -19,27 +22,6 @@ class ClientCredential(BaseModel):
     created_at: datetime
     last_used: datetime
 
-    @classmethod
-    def from_dict(cls, data: Dict):
-        return cls(
-            display_name=data["display_name"],
-            client_id=data["client_id"],
-            client_secret=data["client_secret"],
-            endpoint=data["endpoint"],
-            created_at=datetime.fromisoformat(data["created_at"]),
-            last_used=datetime.fromisoformat(data["last_used"]),
-        )
-
-    def to_dict(self) -> Dict:
-        return {
-            "display_name": self.display_name,
-            "client_id": self.client_id,
-            "client_secret": self.client_secret,
-            "endpoint": self.endpoint,
-            "created_at": self.created_at.isoformat(),
-            "last_used": self.last_used.isoformat(),
-        }
-
     def get_display_string(self):
         return f"{self.display_name} - Client ID: {self.client_id} (Created: {self.created_at.strftime('%Y-%m-%d %H:%M:%S')})"
 
@@ -70,6 +52,7 @@ class CredentialManager:
 
     def _read_credentials(self) -> Dict[str, List[ClientCredential]]:
         """Read all stored credentials from the config file."""
+        logger.debug("Reading credentials from %s", self.config_path)
         if not self.config_path.exists():
             return {}
 
@@ -77,7 +60,7 @@ class CredentialManager:
             with open(self.config_path, "r") as f:
                 data = json.load(f)
                 return {
-                    env: [ClientCredential.from_dict(cred) for cred in creds]
+                    env: [ClientCredential.model_validate(cred) for cred in creds]
                     for env, creds in data.items()
                 }
         except json.JSONDecodeError:
@@ -87,15 +70,20 @@ class CredentialManager:
         self, credentials: Dict[str, List[ClientCredential]]
     ) -> None:
         data = {
-            env: [cred.to_dict() for cred in creds]
+            env: [cred.model_dump(mode="json") for cred in creds]
             for env, creds in credentials.items()
         }
 
         with open(self.config_path, "w") as f:
             json.dump(data, f, indent=2)
 
+        logger.debug("Credentials written to %s with data: %s", self.config_path, data)
+
         # Ensure file is only readable by the user
         os.chmod(self.config_path, 0o600)
+        logger.debug(
+            "Set permissions for %s to read/write for user only.", self.config_path
+        )
 
     def _store_credential(self, credential: ClientCredential) -> None:
         credentials = self._read_credentials()
@@ -121,37 +109,50 @@ class CredentialManager:
     def get_client_credentials(self) -> Optional[ClientCredential]:
         """Gets stored client credentials or create new ones via browser auth."""
         credentials = self._read_credentials()
+        logger.debug("Stored credentials: %s", credentials)
         env_credentials = credentials.get(self.endpoint, [])
 
         if env_credentials:
+            logger.debug("Found credentials for %s: %s", self.endpoint, env_credentials)
             credential = self._select_credential(env_credentials)
+            logger.debug("Selected credential: %s", credential)
             if credential:
                 credential.last_used = datetime.now(timezone.utc)
                 self._write_credentials(credentials)
                 return credential
 
+        logger.debug("No credentials found for %s. Creating new ones.", self.endpoint)
         return self._create_new_credentials()
 
+    def reset_credentials(self) -> None:
+        """Reset the stored credentials for current environment."""
+        credentials = self._read_credentials()
+        if self.endpoint in credentials:
+            del credentials[self.endpoint]
+            self._write_credentials(credentials)
+            logger.info("Credentials for %s have been reset.", self.endpoint)
+
     def _get_bridge_tokens(self) -> Optional[BridgeToken]:
         """Get bridge tokens from the identity endpoint."""
+        logger.debug("Getting bridge tokens")
         try:
             bridge_endpoint = (
-                f"{self.endpoint}/Identity/CreateBridgeToken?clientId=rapidata-cli"
+                f"{self.endpoint}/identity/bridge-token?clientId=rapidata-cli"
             )
             response = requests.post(bridge_endpoint, verify=self.cert_path)
             if not response.ok:
-                print(f"Failed to get bridge tokens: {response.status_code}")
+                logger.error("Failed to get bridge tokens: %s", response.status_code)
                 return None
 
             data = response.json()
             return BridgeToken(read_key=data["readKey"], write_key=data["writeKey"])
         except requests.RequestException as e:
-            print(f"Failed to get bridge tokens: {e}")
+            logger.error("Failed to get bridge tokens: %s", e)
             return None
 
     def _poll_read_key(self, read_key: str) -> Optional[str]:
         """Poll the read key endpoint until we get an access token."""
-        read_endpoint = f"{self.endpoint}/Identity/ReadBridgeToken"
+        read_endpoint = f"{self.endpoint}/identity/bridge-token"
         start_time = time.time()
 
         while time.time() - start_time < self.poll_timeout:
@@ -168,14 +169,14 @@ class CredentialManager:
                     continue
                 else:
                     # Error occurred
-                    print(f"Error polling read key: {response.status_code}")
+                    logger.error("Error polling read key: %s", response.status_code)
                     return None
 
             except requests.RequestException as e:
-                print(f"Error polling read key: {e}")
+                logger.error("Error polling read key: %s", e)
                 return None
 
-        print("Polling timed out")
+        logger.error("Polling timed out")
         return None
 
     def _create_client(self, access_token: str) -> Optional[Tuple[str, str, str]]:
@@ -197,7 +198,7 @@ class CredentialManager:
             data = response.json()
             return data.get("clientId"), data.get("clientSecret"), display_name
         except requests.RequestException as e:
-            print(f"Failed to create client: {e}")
+            logger.error("Failed to create client: %s", e)
             return None
 
     def _create_new_credentials(self) -> Optional[ClientCredential]:
@@ -206,7 +207,15 @@ class CredentialManager:
             return None
 
         auth_url = f"{self.endpoint}/connect/authorize/external?clientId=rapidata-cli&scope=openid profile email roles&writeKey={bridge_endpoint.write_key}"
-        webbrowser.open(auth_url)
+        could_open_browser = webbrowser.open(auth_url)
+
+        if not could_open_browser:
+            encoded_url = urllib.parse.quote(auth_url, safe="%/:=&?~#+!$,;'@()*[]")
+            managed_print(
+                Fore.RED
+                + f'Please open the following URL in your browser to log in: "{encoded_url}"'
+                + Fore.RESET
+            )
 
         access_token = self._poll_read_key(bridge_endpoint.read_key)
         if not access_token:

rapidata/service/openapi_service.py

@@ -1,13 +1,26 @@
-from rapidata.api_client.api.campaign_api import CampaignApi
-from rapidata.api_client.api.dataset_api import DatasetApi
-from rapidata.api_client.api.order_api import OrderApi
-from rapidata.api_client.api.pipeline_api import PipelineApi
-from rapidata.api_client.api.rapid_api import RapidApi
-from rapidata.api_client.api.validation_api import ValidationApi
-from rapidata.api_client.api.workflow_api import WorkflowApi
-from rapidata.api_client.api_client import ApiClient
+from __future__ import annotations
+
+import subprocess
+from importlib.metadata import version, PackageNotFoundError
+from typing import TYPE_CHECKING
+
 from rapidata.api_client.configuration import Configuration
-from rapidata.service.token_manager import TokenManager, TokenInfo
+from rapidata.service.credential_manager import CredentialManager
+from rapidata.rapidata_client.api.rapidata_api_client import RapidataApiClient
+from rapidata.rapidata_client.config import logger
+
+if TYPE_CHECKING:
+    from rapidata.api_client import CustomerRapidApi
+    from rapidata.api_client.api.campaign_api import CampaignApi
+    from rapidata.api_client.api.asset_api import AssetApi
+    from rapidata.api_client.api.dataset_api import DatasetApi
+    from rapidata.api_client.api.benchmark_api import BenchmarkApi
+    from rapidata.api_client.api.order_api import OrderApi
+    from rapidata.api_client.api.pipeline_api import PipelineApi
+    from rapidata.api_client.api.leaderboard_api import LeaderboardApi
+    from rapidata.api_client.api.validation_set_api import ValidationSetApi
+    from rapidata.api_client.api.workflow_api import WorkflowApi
+    from rapidata.api_client.api.participant_api import ParticipantApi
 
 
 class OpenAPIService:
@@ -15,62 +28,167 @@ class OpenAPIService:
         self,
         client_id: str | None,
         client_secret: str | None,
-        enviroment: str,
+        environment: str,
         oauth_scope: str,
         cert_path: str | None = None,
+        token: dict | None = None,
+        leeway: int = 60,
     ):
-        self.enviroment = enviroment  
-        endpoint=f"https://api.{enviroment}"
-        token_url=f"https://auth.{enviroment}"
-        token_manager = TokenManager(
-            client_id=client_id,
-            client_secret=client_secret,
-            endpoint=token_url,
-            oauth_scope=oauth_scope,
-            cert_path=cert_path,
+        self.environment = environment
+        endpoint = f"https://api.{environment}"
+        auth_endpoint = f"https://auth.{environment}"
+
+        if environment == "rapidata.dev" and not cert_path:
+            cert_path = _get_local_certificate()
+
+        logger.debug(
+            "Using cert_path: %s environment: %s client_id: %s",
+            cert_path,
+            environment,
+            client_id,
         )
-        client_configuration = Configuration(host=endpoint, ssl_ca_cert=cert_path)
-        self.api_client = ApiClient(configuration=client_configuration)
+        logger.debug("Initializing OpenAPIService")
+        self.credential_manager = CredentialManager(
+            endpoint=auth_endpoint, cert_path=cert_path
+        )
+        logger.debug("CredentialManager initialized")
 
-        self.api_client.configuration.api_key["bearer"] = (
-            f"Bearer {token_manager.fetch_token().access_token}"
+        logger.debug("Initializing RapidataApiClient")
+        client_configuration = Configuration(host=endpoint, ssl_ca_cert=cert_path)
+        logger.debug("Client configuration: %s", client_configuration)
+        self.api_client = RapidataApiClient(
+            configuration=client_configuration,
+            header_name="X-Client",
+            header_value=f"RapidataPythonSDK/{self._get_rapidata_package_version()}",
         )
+        logger.debug("RapidataApiClient initialized")
 
-        self._client_id = client_id
-        self._client_secret = client_secret
-        self._oauth_scope = oauth_scope
-        self._token_url = f"{token_url}/connect/token"
-        self._cert_path = cert_path
+        if token:
+            logger.debug("Using token for authentication")
+            self.api_client.rest_client.setup_oauth_with_token(
+                token=token,
+                token_endpoint=f"{auth_endpoint}/connect/token",
+                client_id=client_id,
+                client_secret=client_secret,
+                leeway=leeway,
+            )
+            logger.debug("Token authentication setup complete")
+            return
 
-        token_manager.start_token_refresh(token_callback=self._set_token)
+        if not client_id or not client_secret:
+            logger.debug(
+                "Client ID and secret not provided, fetching from credential manager"
+            )
+            credentials = self.credential_manager.get_client_credentials()
+            if not credentials:
+                raise ValueError("Failed to fetch client credentials")
+            client_id = credentials.client_id
+            client_secret = credentials.client_secret
+
+        self.api_client.rest_client.setup_oauth_client_credentials(
+            client_id=client_id,
+            client_secret=client_secret,
+            token_endpoint=f"{auth_endpoint}/connect/token",
+            scope=oauth_scope,
+        )
+        logger.debug("Client credentials authentication setup complete")
+
+    def reset_credentials(self):
+        self.credential_manager.reset_credentials()
 
     @property
     def order_api(self) -> OrderApi:
+        from rapidata.api_client.api.order_api import OrderApi
+
         return OrderApi(self.api_client)
 
+    @property
+    def asset_api(self) -> AssetApi:
+        from rapidata.api_client.api.asset_api import AssetApi
+
+        return AssetApi(self.api_client)
+
     @property
     def dataset_api(self) -> DatasetApi:
+        from rapidata.api_client.api.dataset_api import DatasetApi
+
         return DatasetApi(self.api_client)
 
     @property
-    def validation_api(self) -> ValidationApi:
-        return ValidationApi(self.api_client)
+    def validation_api(self) -> ValidationSetApi:
+        from rapidata.api_client.api.validation_set_api import ValidationSetApi
+
+        return ValidationSetApi(self.api_client)
 
     @property
-    def rapid_api(self) -> RapidApi:
-        return RapidApi(self.api_client)
+    def rapid_api(self) -> CustomerRapidApi:
+        from rapidata.api_client.api.customer_rapid_api import CustomerRapidApi
+
+        return CustomerRapidApi(self.api_client)
 
     @property
     def campaign_api(self) -> CampaignApi:
+        from rapidata.api_client.api.campaign_api import CampaignApi
+
         return CampaignApi(self.api_client)
 
     @property
     def pipeline_api(self) -> PipelineApi:
+        from rapidata.api_client.api.pipeline_api import PipelineApi
+
         return PipelineApi(self.api_client)
 
     @property
     def workflow_api(self) -> WorkflowApi:
+        from rapidata.api_client.api.workflow_api import WorkflowApi
+
         return WorkflowApi(self.api_client)
 
-    def _set_token(self, token: TokenInfo):
-        self.api_client.configuration.api_key["bearer"] = f"Bearer {token.access_token}"
+    @property
+    def leaderboard_api(self) -> LeaderboardApi:
+        from rapidata.api_client.api.leaderboard_api import LeaderboardApi
+
+        return LeaderboardApi(self.api_client)
+
+    @property
+    def benchmark_api(self) -> BenchmarkApi:
+        from rapidata.api_client.api.benchmark_api import BenchmarkApi
+
+        return BenchmarkApi(self.api_client)
+
+    @property
+    def participant_api(self) -> ParticipantApi:
+        from rapidata.api_client.api.participant_api import ParticipantApi
+
+        return ParticipantApi(self.api_client)
+
+    def _get_rapidata_package_version(self):
+        """
+        Returns the version of the currently installed rapidata package.
+
+        Returns:
+            Optional[str]: The version of the installed rapidata package,
+                        or None if the package is not found
+        """
+        try:
+            return version("rapidata")
+        except PackageNotFoundError:
+            return None
+
+    def __repr__(self) -> str:
+        return self.__str__()
+
+    def __str__(self) -> str:
+        return f"OpenAPIService(environment={self.environment})"
+
+
+def _get_local_certificate() -> str | None:
+    result = subprocess.run(["mkcert", "-CAROOT"], capture_output=True)
+    if result.returncode != 0:
+        return None
+
+    output = result.stdout.decode("utf-8").strip()
+    if not output:
+        return None
+
+    return f"{output}/rootCA.pem"