rapidata 1.6.1__py3-none-any.whl → 1.6.3__py3-none-any.whl

rapidata/rapidata_client/rapidata_client.py

@@ -27,10 +27,10 @@ class RapidataClient:
 
     def __init__(
         self,
-        client_id: str,
-        client_secret: str,
+        client_id: str | None = None,
+        client_secret: str | None = None,
         endpoint: str = "https://api.rapidata.ai",
-        token_url: str = "https://auth.rapidata.ai/connect/token",
+        token_url: str = "https://auth.rapidata.ai",
         oauth_scope: str = "openid",
         cert_path: str | None = None,
     ):

rapidata/service/credential_manager.py

@@ -0,0 +1,232 @@
+import json
+import os
+import time
+import webbrowser
+from datetime import datetime, timezone
+from pathlib import Path
+from socket import gethostname
+from typing import Dict, List, Optional, Tuple
+
+import requests
+from pydantic import BaseModel
+
+
+class ClientCredential(BaseModel):
+    display_name: str
+    client_id: str
+    client_secret: str
+    endpoint: str
+    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')})"
+
+
+class BridgeToken(BaseModel):
+    read_key: str
+    write_key: str
+
+
+class CredentialManager:
+    def __init__(
+        self,
+        endpoint: str,
+        cert_path: str | None = None,
+        poll_timeout: int = 60,
+        poll_interval: int = 1,
+    ):
+        self.endpoint = endpoint
+        self.cert_path = cert_path
+        self.poll_timeout = poll_timeout
+        self.poll_interval = poll_interval
+
+        self.config_dir = Path.home() / ".config" / "rapidata"
+        self.config_path = self.config_dir / "credentials.json"
+
+        # Ensure config directory exists
+        self.config_dir.mkdir(parents=True, exist_ok=True)
+
+    def _read_credentials(self) -> Dict[str, List[ClientCredential]]:
+        """Read all stored credentials from the config file."""
+        if not self.config_path.exists():
+            return {}
+
+        try:
+            with open(self.config_path, "r") as f:
+                data = json.load(f)
+                return {
+                    env: [ClientCredential.from_dict(cred) for cred in creds]
+                    for env, creds in data.items()
+                }
+        except json.JSONDecodeError:
+            return {}
+
+    def _write_credentials(
+        self, credentials: Dict[str, List[ClientCredential]]
+    ) -> None:
+        data = {
+            env: [cred.to_dict() for cred in creds]
+            for env, creds in credentials.items()
+        }
+
+        with open(self.config_path, "w") as f:
+            json.dump(data, f, indent=2)
+
+        # Ensure file is only readable by the user
+        os.chmod(self.config_path, 0o600)
+
+    def _store_credential(self, credential: ClientCredential) -> None:
+        credentials = self._read_credentials()
+
+        if credential.endpoint not in credentials:
+            credentials[credential.endpoint] = []
+
+        credentials[credential.endpoint].append(credential)
+        self._write_credentials(credentials)
+
+    @staticmethod
+    def _select_credential(
+        credentials: List[ClientCredential],
+    ) -> Optional[ClientCredential]:
+        if not credentials:
+            return None
+
+        if len(credentials) == 1:
+            return credentials[0]
+
+        return max(credentials, key=lambda c: c.last_used)
+
+    def get_client_credentials(self) -> Optional[ClientCredential]:
+        """Gets stored client credentials or create new ones via browser auth."""
+        credentials = self._read_credentials()
+        env_credentials = credentials.get(self.endpoint, [])
+
+        if env_credentials:
+            credential = self._select_credential(env_credentials)
+            if credential:
+                credential.last_used = datetime.now(timezone.utc)
+                self._write_credentials(credentials)
+                return credential
+
+        return self._create_new_credentials()
+
+    def _get_bridge_tokens(self) -> Optional[BridgeToken]:
+        """Get bridge tokens from the identity endpoint."""
+        try:
+            bridge_endpoint = (
+                f"{self.endpoint}/Identity/CreateBridgeToken?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}")
+                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}")
+            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"
+        start_time = time.time()
+
+        while time.time() - start_time < self.poll_timeout:
+            try:
+                response = requests.get(
+                    read_endpoint, params={"readKey": read_key}, verify=self.cert_path
+                )
+
+                if response.status_code == 200:
+                    return response.json().get("accessToken")
+                elif response.status_code == 202:
+                    # Still processing
+                    time.sleep(self.poll_interval)
+                    continue
+                else:
+                    # Error occurred
+                    print(f"Error polling read key: {response.status_code}")
+                    return None
+
+            except requests.RequestException as e:
+                print(f"Error polling read key: {e}")
+                return None
+
+        print("Polling timed out")
+        return None
+
+    def _create_client(self, access_token: str) -> Optional[Tuple[str, str, str]]:
+        """Create a new client using the access token."""
+        try:
+            # set the display name to the hostname
+            display_name = f"{gethostname()} - CLI"
+            response = requests.post(
+                f"{self.endpoint}/Client",
+                headers={
+                    "Authorization": f"Bearer {access_token}",
+                    "Content-Type": "application/json",
+                    "Accept": "*/*",
+                },
+                json={"displayName": display_name},
+                verify=self.cert_path,
+            )
+            response.raise_for_status()
+            data = response.json()
+            return data.get("clientId"), data.get("clientSecret"), display_name
+        except requests.RequestException as e:
+            print(f"Failed to create client: {e}")
+            return None
+
+    def _create_new_credentials(self) -> Optional[ClientCredential]:
+        bridge_endpoint = self._get_bridge_tokens()
+        if not bridge_endpoint:
+            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)
+
+        access_token = self._poll_read_key(bridge_endpoint.read_key)
+        if not access_token:
+            return None
+
+        client_state = self._create_client(access_token)
+
+        if not client_state:
+            raise ValueError("Failed to create client")
+
+        client_id, client_secret, display_name = client_state
+
+        credential = ClientCredential(
+            client_id=client_id,
+            client_secret=client_secret,
+            display_name=display_name,
+            endpoint=self.endpoint,
+            created_at=datetime.now(timezone.utc),
+            last_used=datetime.now(timezone.utc),
+        )
+
+        self._store_credential(credential)
+        return credential

rapidata/service/openapi_service.py

@@ -1,7 +1,3 @@
-import json
-import time
-import requests
-import threading
 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
@@ -11,61 +7,44 @@ 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 rapidata.api_client.configuration import Configuration
+from rapidata.service.token_manager import TokenManager, TokenInfo
 
 
 class OpenAPIService:
-
-    _TOKEN_EXPIRATION_MINUTES = 30
-
     def __init__(
         self,
-        client_id: str,
-        client_secret: str,
+        client_id: str | None,
+        client_secret: str | None,
         endpoint: str,
         token_url: str,
         oauth_scope: str,
-        cert_path: str | None = None
+        cert_path: str | None = None,
     ):
+        token_manager = TokenManager(
+            client_id=client_id,
+            client_secret=client_secret,
+            endpoint=token_url,
+            oauth_scope=oauth_scope,
+            cert_path=cert_path
+        )
         client_configuration = Configuration(host=endpoint, ssl_ca_cert=cert_path)
         self.api_client = ApiClient(configuration=client_configuration)
 
         self._client_id = client_id
         self._client_secret = client_secret
         self._oauth_scope = oauth_scope
-        self._token_url = token_url
+        self._token_url = f"{token_url}/connect/token"
         self._cert_path = cert_path
 
-        self._api_client = ApiClient()
-        self._order_api = OrderApi(self.api_client)
-        self._dataset_api = DatasetApi(self.api_client)
-
-        api_token = self.__fetch_token(
-            self._client_id, self._client_secret, self._oauth_scope, self._token_url, self._cert_path
-        )
-        self.api_client.configuration.api_key["bearer"] = f"Bearer {api_token}"
-
-        refresh_thread = threading.Thread(
-            target=lambda: self.__refresh_token_periodically(self._TOKEN_EXPIRATION_MINUTES - 1)
-        )
-        refresh_thread.daemon = True
-        refresh_thread.start()
-
-    def __refresh_token_periodically(self, refresh_interval):
-        while True:
-            new_token = self.__fetch_token(
-                self._client_id, self._client_secret, self._oauth_scope, self._token_url, self._cert_path
-            )
-            self.api_client.configuration.api_key["bearer"] = f"Bearer {new_token}"
-
-            time.sleep(refresh_interval)
+        token_manager.start_token_refresh(token_callback=self._set_token)
 
     @property
     def order_api(self) -> OrderApi:
-        return self._order_api
+        return OrderApi(self.api_client)
 
     @property
     def dataset_api(self) -> DatasetApi:
-        return self._dataset_api
+        return DatasetApi(self.api_client)
 
     @property
     def validation_api(self) -> ValidationApi:
@@ -87,23 +66,5 @@ class OpenAPIService:
     def workflow_api(self) -> WorkflowApi:
         return WorkflowApi(self.api_client)
 
-    @staticmethod
-    def __fetch_token(client_id: str, client_secret: str, scope: str, token_url: str, cert_path: str | None = None) -> str:
-        try:
-            return requests.post(
-                token_url,
-                data={
-                    'grant_type': 'client_credentials',
-                    'client_id': client_id,
-                    'client_secret': client_secret,
-                    'scope': scope,
-                },
-                verify=cert_path
-            ).json()['access_token']
-        except requests.RequestException as e:
-            raise Exception(f"Failed to fetch token: {e}")
-        except json.JSONDecodeError as e:
-            raise Exception(f"Failed to parse token response: {e}")
-        except KeyError as e:
-            raise Exception(f"Failed to extract token from response: {e}")
-
+    def _set_token(self, token: TokenInfo):
+        self.api_client.configuration.api_key["bearer"] = f"Bearer {token.access_token}"

rapidata/service/token_manager.py

@@ -0,0 +1,175 @@
+import json
+import logging
+import threading
+from datetime import datetime, timedelta
+from typing import Optional, Callable
+
+import requests
+from pydantic import BaseModel
+
+from rapidata.service.credential_manager import CredentialManager
+
+logger = logging.getLogger(__name__)
+
+
+class TokenInfo(BaseModel):
+    access_token: str
+    expires_in: int
+    issued_at: datetime
+    token_type: str = "Bearer"
+
+    @property
+    def auth_header(self):
+        return f"{self.token_type} {self.access_token}"
+
+    @property
+    def time_remaining(self):
+        remaining = (
+            (self.issued_at + timedelta(seconds=self.expires_in)) - datetime.now()
+        ).total_seconds()
+        return max(0.0, remaining)
+
+
+class TokenManager:
+    def __init__(
+        self,
+        client_id: str | None = None,
+        client_secret: str | None = None,
+        endpoint: str = "https://auth.rapidata.ai",
+        oauth_scope: str = "openid profile email",
+        cert_path: str | None = None,
+        refresh_threshold: float = 0.8,
+        max_sleep_time: float = 30,
+    ):
+        self._client_id = client_id
+        self._client_secret = client_secret
+
+        if not client_id or not client_secret:
+            credential_manager = CredentialManager(
+                endpoint=endpoint, cert_path=cert_path
+            )
+            credentials = credential_manager.get_client_credentials()
+            if not credentials:
+                raise ValueError("Failed to fetch client credentials")
+            self._client_id = credentials.client_id
+            self._client_secret = credentials.client_secret
+
+        self._endpoint = endpoint
+        self._oauth_scope = oauth_scope
+        self._cert_path = cert_path
+        self._refresh_threshold = refresh_threshold
+        self._max_sleep_time = max_sleep_time
+
+        self._token_lock = threading.Lock()
+        self._current_token: Optional[TokenInfo] = None
+        self._refresh_thread: Optional[threading.Thread] = None
+        self._should_stop = threading.Event()
+
+    def fetch_token(self):
+        try:
+            response = requests.post(
+                f"{self._endpoint}/connect/token",
+                data={
+                    "grant_type": "client_credentials",
+                    "client_id": self._client_id,
+                    "client_secret": self._client_secret,
+                    "scope": self._oauth_scope,
+                },
+                verify=self._cert_path,
+            )
+
+            if response.ok:
+                data = response.json()
+                return TokenInfo(
+                    access_token=data["access_token"],
+                    token_type=data["token_type"],
+                    expires_in=data["expires_in"],
+                    issued_at=datetime.now(),
+                )
+
+            else:
+                data = response.text
+                error_description = "An unknown error occurred"
+                if "error_description" in data:
+                    error_description = (
+                        data.split("error_description")[1].split("\n")[0].strip()
+                    )
+                raise ValueError(f"Failed to fetch token: {error_description}")
+        except requests.RequestException as e:
+            raise ValueError(f"Failed to fetch token: {e}")
+        except json.JSONDecodeError as e:
+            raise ValueError(f"Failed to parse token response: {e}")
+        except KeyError as e:
+            raise ValueError(f"Failed to extract token from response: {e}")
+
+    def start_token_refresh(self, token_callback: Callable[[TokenInfo], None]) -> None:
+        if self._refresh_thread and self._refresh_thread.is_alive():
+            logger.error("Token refresh thread is already running")
+            return
+
+        def refresh_loop():
+            while not self._should_stop.is_set():
+                try:
+                    with self._token_lock:
+                        if self._should_refresh_token(self._current_token):
+                            logger.debug("Refreshing token")
+                            self._current_token = self.fetch_token()
+                            token_callback(self._current_token)
+
+                    if self._current_token:
+                        time_until_refresh_threshold = (
+                            self._current_token.time_remaining
+                            - (
+                                self._current_token.expires_in
+                                * (1 - self._refresh_threshold)
+                            )
+                        )
+                        logger.debug("Time until refresh threshold: %s", time_until_refresh_threshold)
+                        sleep_time = min(
+                            self._max_sleep_time, time_until_refresh_threshold
+                        )
+                        logger.debug(
+                            f"Sleeping for {sleep_time} until checking the token again"
+                        )
+                        self._should_stop.wait(timeout=max(1.0, sleep_time))
+                    else:
+                        self._should_stop.wait(timeout=self._max_sleep_time)
+                except Exception as e:
+                    logger.error("Failed to refresh token: %s", e)
+                    self._should_stop.wait(timeout=5)
+
+        self._should_stop.clear()
+        self._refresh_thread = threading.Thread(target=refresh_loop, daemon=True)
+        self._refresh_thread.start()
+
+    def stop_token_refresh(self):
+        self._should_stop.set()
+        if self._refresh_thread:
+            self._refresh_thread.join(timeout=1)
+            self._refresh_thread = None
+
+    def get_current_token(self) -> Optional[TokenInfo]:
+        with self._token_lock:
+            return self._current_token
+
+    def _should_refresh_token(self, token: TokenInfo | None) -> bool:
+        if not token:
+            return True
+
+        limit = token.expires_in * (1 - self._refresh_threshold)
+
+        logger.debug(
+            "The token was issued at %s, it expires in %s. It has %s seconds remaining and we refresh the token when it has %s seconds remaining",
+            token.issued_at,
+            token.expires_in,
+            token.time_remaining,
+            limit,
+        )
+        return token.time_remaining < limit
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.stop_token_refresh()
+        return False

{rapidata-1.6.1.dist-info → rapidata-1.6.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: rapidata
-Version: 1.6.1
+Version: 1.6.3
 Summary: Rapidata package containing the Rapidata Python Client to interact with the Rapidata Web API in an easy way.
 License: Apache-2.0
 Author: Rapidata AG

{rapidata-1.6.1.dist-info → rapidata-1.6.3.dist-info}/RECORD

@@ -345,7 +345,7 @@ rapidata/rapidata_client/metadata/transcription_metadata.py,sha256=THtDEVCON4Ulc
 rapidata/rapidata_client/order/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 rapidata/rapidata_client/order/rapidata_order.py,sha256=t5ddz_6Dk2DOpsDQ2EiZHJJB1RzU-etKwBL0RErEsSY,4567
 rapidata/rapidata_client/order/rapidata_order_builder.py,sha256=b--9byhsiAW1fL0mVSPzGJT0X4MQ-tCC0BNjIm2vu-Q,16406
-rapidata/rapidata_client/rapidata_client.py,sha256=DwWi7gtACV_OxQDEBHCILv8921tWQy3UtXzJ0ZcjjTo,8018
+rapidata/rapidata_client/rapidata_client.py,sha256=a_OZBd3sldws3ElqZt_eyqrSzdUNLxynjG5XCNLmmnY,8032
 rapidata/rapidata_client/referee/__init__.py,sha256=E1VODxTjoQRnxzdgMh3aRlDLouxe1nWuvozEHXD2gq4,150
 rapidata/rapidata_client/referee/base_referee.py,sha256=bMy7cw0a-pGNbFu6u_1_Jplu0A483Ubj4oDQzh8vu8k,493
 rapidata/rapidata_client/referee/early_stopping_referee.py,sha256=Dg2Kk7OiLBtS3kknsLxyJIlS27xmPvsikFR6g4xlbTE,1862
@@ -371,9 +371,11 @@ rapidata/rapidata_client/workflow/evaluation_workflow.py,sha256=IBQoVFxOaeCDIBfa
 rapidata/rapidata_client/workflow/free_text_workflow.py,sha256=VaypoG3yKgsbtVyqxta3W28eDwdnGebCy2xDWPCBMyo,1566
 rapidata/rapidata_client/workflow/transcription_workflow.py,sha256=_KDtGCdRhauJm3jQHpwhY-Hq79CLg5I8q2RgOz5lo1g,1404
 rapidata/service/__init__.py,sha256=s9bS1AJZaWIhLtJX_ZA40_CK39rAAkwdAmymTMbeWl4,68
+rapidata/service/credential_manager.py,sha256=CNr7jasOR3vN2XCYtCk4EI38KV2ggSVAg1w_Hhp5LII,7929
 rapidata/service/local_file_service.py,sha256=pgorvlWcx52Uh3cEG6VrdMK_t__7dacQ_5AnfY14BW8,877
-rapidata/service/openapi_service.py,sha256=l-ga9TLKWR6Gwx9nbDSQ2J1HZRPhgk_mGgnC3kud3uw,3688
-rapidata-1.6.1.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-rapidata-1.6.1.dist-info/METADATA,sha256=2Gft6rEkY_wWlOaEef9TuKZPy52w0tfsMES4zfdSxMo,1056
-rapidata-1.6.1.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
-rapidata-1.6.1.dist-info/RECORD,,
+rapidata/service/openapi_service.py,sha256=ejAIilSjRweGcR4nN4txqZNDuoI-WIJEuagdN2oSd58,2335
+rapidata/service/token_manager.py,sha256=JZ5YbR5Di8dO3H4kK11d0kzWlrXxjgCmeNkHA4AapCM,6425
+rapidata-1.6.3.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+rapidata-1.6.3.dist-info/METADATA,sha256=q-e7zziNADs-DUkyxagIGXd1wh9TRiPAs_ELku0_LEs,1056
+rapidata-1.6.3.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
+rapidata-1.6.3.dist-info/RECORD,,