sutro 0.0.0__py3-none-any.whl → 0.1.12__py3-none-any.whl

sutro/sdk.py

@@ -0,0 +1,1097 @@
+import threading
+from concurrent.futures import ThreadPoolExecutor
+from contextlib import contextmanager
+
+import requests
+import pandas as pd
+import polars as pl
+import json
+from typing import Union, List, Optional, Literal, Generator, Dict, Any
+import os
+import sys
+from yaspin import yaspin
+from yaspin.spinners import Spinners
+from colorama import init, Fore, Back, Style
+from tqdm import tqdm
+import time
+from pydantic import BaseModel
+import json
+
+# Initialize colorama (required for Windows)
+init()
+
+
+# This is how yaspin defines is_jupyter logic
+def is_jupyter() -> bool:
+    return not sys.stdout.isatty()
+
+
+# `color` param not supported in Jupyter notebooks
+YASPIN_COLOR = None if is_jupyter() else "blue"
+SPINNER = Spinners.dots14
+
+
+def to_colored_text(
+    text: str, state: Optional[Literal["success", "fail"]] = None
+) -> str:
+    """
+    Apply color to text based on state.
+
+    Args:
+        text (str): The text to color
+        state (Optional[Literal['success', 'fail']]): The state that determines the color.
+            Options: 'success', 'fail', or None (default blue)
+
+    Returns:
+        str: Text with appropriate color applied
+    """
+    match state:
+        case "success":
+            return f"{Fore.GREEN}{text}{Style.RESET_ALL}"
+        case "fail":
+            return f"{Fore.RED}{text}{Style.RESET_ALL}"
+        case _:
+            # Default to blue for normal/processing states
+            return f"{Fore.BLUE}{text}{Style.RESET_ALL}"
+
+
+class Sutro:
+    def __init__(
+        self, api_key: str = None, base_url: str = "https://api.sutro.sh/"
+    ):
+        self.api_key = api_key or self.check_for_api_key()
+        self.base_url = base_url
+        self.HEARTBEAT_INTERVAL_SECONDS = 15  # Keep in sync w what the backend expects
+
+    def check_for_api_key(self):
+        """
+        Check for an API key in the user's home directory.
+
+        This method looks for a configuration file named 'config.json' in the
+        '.sutro' directory within the user's home directory.
+        If the file exists, it attempts to read the API key from it.
+
+        Returns:
+            str or None: The API key if found in the configuration file, or None if not found.
+
+        Note:
+            The expected structure of the config.json file is:
+            {
+                "api_key": "your_api_key_here"
+            }
+        """
+        CONFIG_DIR = os.path.expanduser("~/.sutro")
+        CONFIG_FILE = os.path.join(CONFIG_DIR, "config.json")
+        if os.path.exists(CONFIG_FILE):
+            with open(CONFIG_FILE, "r") as f:
+                config = json.load(f)
+            return config.get("api_key")
+        else:
+            return None
+
+    def set_api_key(self, api_key: str):
+        """
+        Set the API key for the Sutro API.
+
+        This method allows you to set the API key for the Sutro API.
+        The API key is used to authenticate requests to the API.
+
+        Args:
+            api_key (str): The API key to set.
+
+        Returns:
+            None
+        """
+        self.api_key = api_key
+
+    def handle_data_helper(
+        self, data: Union[List, pd.DataFrame, pl.DataFrame, str], column: str = None
+    ):
+        if isinstance(data, list):
+            input_data = data
+        elif isinstance(data, (pd.DataFrame, pl.DataFrame)):
+            if column is None:
+                raise ValueError("Column name must be specified for DataFrame input")
+            input_data = data[column].to_list()
+        elif isinstance(data, str):
+            if data.startswith("stage-"):
+                input_data = data + ":" + column
+            else:
+                file_ext = os.path.splitext(data)[1].lower()
+                if file_ext == ".csv":
+                    df = pl.read_csv(data)
+                elif file_ext == ".parquet":
+                    df = pl.read_parquet(data)
+                elif file_ext in [".txt", ""]:
+                    with open(data, "r") as file:
+                        input_data = [line.strip() for line in file]
+                else:
+                    raise ValueError(f"Unsupported file type: {file_ext}")
+
+                if file_ext in [".csv", ".parquet"]:
+                    if column is None:
+                        raise ValueError(
+                            "Column name must be specified for CSV/Parquet input"
+                        )
+                    input_data = df[column].to_list()
+        else:
+            raise ValueError(
+                "Unsupported data type. Please provide a list, DataFrame, or file path."
+            )
+
+        return input_data
+
+    def set_base_url(self, base_url: str):
+        """
+        Set the base URL for the Sutro API.
+
+        This method allows you to set the base URL for the Sutro API.
+        The base URL is used to authenticate requests to the API.
+
+        Args:
+            base_url (str): The base URL to set.
+        """
+        self.base_url = base_url
+
+    def infer(
+        self,
+        data: Union[List, pd.DataFrame, pl.DataFrame, str],
+        model: str = "llama-3.1-8b",
+        column: str = None,
+        output_column: str = "inference_result",
+        job_priority: int = 0,
+        output_schema: Union[Dict[str, Any], BaseModel] = None,
+        sampling_params: dict = None,
+        system_prompt: str = None,
+        dry_run: bool = False,
+        stay_attached: bool = False,
+        random_seed_per_input: bool = False,
+        truncate_rows: bool = False
+    ):
+        """
+        Run inference on the provided data.
+
+        This method allows you to run inference on the provided data using the Sutro API.
+        It supports various data types such as lists, pandas DataFrames, polars DataFrames, file paths and stages.
+
+        Args:
+            data (Union[List, pd.DataFrame, pl.DataFrame, str]): The data to run inference on.
+            model (str, optional): The model to use for inference. Defaults to "llama-3.1-8b".
+            column (str, optional): The column name to use for inference. Required if data is a DataFrame, file path, or stage.
+            output_column (str, optional): The column name to store the inference results in if the input is a DataFrame. Defaults to "inference_result".
+            job_priority (int, optional): The priority of the job. Defaults to 0.
+            output_schema (Union[Dict[str, Any], BaseModel], optional): A structured schema for the output.
+                Can be either a dictionary representing a JSON schema or a pydantic BaseModel. Defaults to None.
+            sampling_params: (dict, optional): The sampling parameters to use at generation time, ie temperature, top_p etc.
+            system_prompt (str, optional): A system prompt to add to all inputs. This allows you to define the behavior of the model. Defaults to None.
+            dry_run (bool, optional): If True, the method will return cost estimates instead of running inference. Defaults to False.
+            stay_attached (bool, optional): If True, the method will stay attached to the job until it is complete. Defaults to True for prototyping jobs, False otherwise.
+            random_seed_per_input (bool, optional): If True, the method will use a different random seed for each input. Defaults to False.
+            truncate_rows (bool, optional): If True, any rows that have a token count exceeding the context window length of the selected model will be truncated to the max length that will fit within the context window. Defaults to False.
+
+        Returns:
+            Union[List, pd.DataFrame, pl.DataFrame, str]: The results of the inference.
+
+        """
+        input_data = self.handle_data_helper(data, column)
+        stay_attached = stay_attached or job_priority == 0
+
+        # Convert BaseModel to dict if needed
+        if output_schema is not None:
+            if hasattr(output_schema, 'model_json_schema'):  # Check for pydantic Model interface
+                json_schema = output_schema.model_json_schema()
+            elif isinstance(output_schema, dict):
+                json_schema = output_schema
+            else:
+                raise ValueError("Invalid output schema type. Must be a dictionary or a pydantic Model.")
+        else:
+            json_schema = None
+
+        endpoint = f"{self.base_url}/batch-inference"
+        headers = {
+            "Authorization": f"Key {self.api_key}",
+            "Content-Type": "application/json",
+        }
+        payload = {
+            "model": model,
+            "inputs": input_data,
+            "job_priority": job_priority,
+            "json_schema": json_schema,
+            "system_prompt": system_prompt,
+            "dry_run": dry_run,
+            "sampling_params": sampling_params,
+            "random_seed_per_input": random_seed_per_input,
+            "truncate_rows": truncate_rows
+        }
+        if dry_run:
+            spinner_text = to_colored_text("Retrieving cost estimates...")
+        else:
+            t = f"Creating priority {job_priority} job"
+            spinner_text = to_colored_text(t)
+
+        # There are two gotchas with yaspin:
+        # 1. Can't use print while in spinner is running
+        # 2. When writing to stdout via spinner.fail, spinner.write etc, there is a pretty strict
+        # limit for content length in jupyter notebooks, where it wisll give an error about:
+        # Terminal size {self._terminal_width} is too small to display spinner with the given settings.
+        # https://github.com/pavdmyt/yaspin/blob/9c7430b499ab4611888ece39783a870e4a05fa45/yaspin/core.py#L568-L571
+        job_id = None
+        with yaspin(SPINNER, text=spinner_text, color=YASPIN_COLOR) as spinner:
+            response = requests.post(
+                endpoint, data=json.dumps(payload), headers=headers
+            )
+            response_data = response.json()
+            if response.status_code != 200:
+                spinner.write(
+                    to_colored_text(f"Error: {response.status_code}", state="fail")
+                )
+                spinner.stop()
+                print(to_colored_text(response.json(), state="fail"))
+                return
+            else:
+                if dry_run:
+                    spinner.write(
+                        to_colored_text("✔ Cost estimates retrieved", state="success")
+                    )
+                    return response_data["results"]
+                else:
+                    job_id = response_data["results"]
+                    spinner.write(
+                        to_colored_text(
+                            f"🛠️  Priority {job_priority} Job created with ID: {job_id}",
+                            state="success",
+                        )
+                    )
+                    if not stay_attached:
+                        spinner.write(
+                            to_colored_text(
+                                f"Use `so.get_job_status('{job_id}')` to check the status of the job."
+                            )
+                        )
+                        return job_id
+
+        success = False
+        if stay_attached and job_id is not None:
+            s = requests.Session()
+            payload = {
+                "job_id": job_id,
+            }
+            pbar = None
+
+            # Register for stream and get session token
+            session_token = self.register_stream_listener(job_id)
+
+            # Use the heartbeat session context manager
+            with self.stream_heartbeat_session(job_id, session_token) as s:
+                with s.get(
+                        f"{self.base_url}/stream-job-progress/{job_id}?request_session_token={session_token}",
+                        headers=headers,
+                        stream=True,
+                ) as streaming_response:
+                    streaming_response.raise_for_status()
+                    spinner = yaspin(
+                        SPINNER,
+                        text=to_colored_text("Awaiting status updates..."),
+                        color=YASPIN_COLOR,
+                    )
+                    spinner.start()
+                    for line in streaming_response.iter_lines():
+                        if line:
+                            try:
+                                json_obj = json.loads(line)
+                            except json.JSONDecodeError:
+                                print("Error: ", line, flush=True)
+                                continue
+
+                            if json_obj["update_type"] == "progress":
+                                if pbar is None:
+                                    spinner.stop()
+                                    postfix = f"Input tokens processed: 0"
+                                    pbar = self.fancy_tqdm(
+                                        total=len(input_data),
+                                        desc="Progress",
+                                        style=1,
+                                        postfix=postfix,
+                                    )
+                                if json_obj["result"] > pbar.n:
+                                    pbar.update(json_obj["result"] - pbar.n)
+                                    pbar.refresh()
+                                if json_obj["result"] == len(input_data):
+                                    pbar.close()
+                                    success = True
+                            elif json_obj["update_type"] == "tokens":
+                                if pbar is not None:
+                                    pbar.postfix = f"Input tokens processed: {json_obj['result']['input_tokens']}, Tokens generated: {json_obj['result']['output_tokens']}, Total tokens/s: {json_obj['result'].get('total_tokens_processed_per_second')}"
+                                    pbar.refresh()
+            if success:
+                spinner.text = to_colored_text(
+                    "✔ Job succeeded. Obtaining results...", state="success"
+                )
+                spinner.start()
+
+                payload = {
+                    "job_id": job_id,
+                }
+
+                # TODO: we implment retries in cases where the job hasn't written results yet
+                # it would be better if we could receive a fully succeeded status from the job
+                # and not have such a race condition
+                max_retries = 20 # winds up being 100 seconds cumulative delay
+                retry_delay = 5 # initial delay in seconds
+
+                for _ in range(max_retries):
+                    time.sleep(retry_delay)
+
+                    job_results_response = s.post(
+                        f"{self.base_url}/job-results",
+                        headers=headers,
+                        data=json.dumps(payload),
+                    )
+                    if job_results_response.status_code == 200:
+                        break
+
+                if job_results_response.status_code != 200:
+                    spinner.write(
+                        to_colored_text(
+                            "Job succeeded, but results are not yet available. Use `so.get_job_results('{job_id}')` to obtain results.",
+                            state="fail",
+                        )
+                    )
+                    spinner.stop()
+                    return
+
+                results = job_results_response.json()["results"]
+
+                spinner.write(
+                    to_colored_text(
+                        f"✔ Job results received. You can re-obtain the results with `so.get_job_results('{job_id}')`",
+                        state="success",
+                    )
+                )
+                spinner.stop()
+
+                if isinstance(data, (pd.DataFrame, pl.DataFrame)):
+                    sample_n = 1 if sampling_params is None else sampling_params["n"]
+                    if sample_n > 1:
+                        results = [
+                            results[i : i + sample_n]
+                            for i in range(0, len(results), sample_n)
+                        ]
+                    if isinstance(data, pd.DataFrame):
+                        data[output_column] = results
+                    elif isinstance(data, pl.DataFrame):
+                        data = data.with_columns(pl.Series(output_column, results))
+                    return data
+
+                return results
+
+    def register_stream_listener(self, job_id: str) -> str:
+        """Register a new stream listener and get a session token."""
+        headers = {
+            "Authorization": f"Key {self.api_key}",
+            "Content-Type": "application/json",
+        }
+        with requests.post(
+                f"{self.base_url}/register-stream-listener/{job_id}",
+                headers=headers,
+        ) as response:
+            response.raise_for_status()
+            data = response.json()
+            return data["request_session_token"]
+
+    # This is a best effort action and is ok if it sometimes doesn't complete etc
+    def unregister_stream_listener(self, job_id: str, session_token: str):
+        """Explicitly unregister a stream listener."""
+        headers = {
+            "Authorization": f"Key {self.api_key}",
+            "Content-Type": "application/json",
+        }
+        with requests.post(
+                f"{self.base_url}/unregister-stream-listener/{job_id}",
+                headers=headers,
+                json={"request_session_token": session_token},
+        ) as response:
+            response.raise_for_status()
+
+    def start_heartbeat(
+            self,
+            job_id: str,
+            session_token: str,
+            session: requests.Session,
+            stop_event: threading.Event
+    ):
+        """Send heartbeats until stopped."""
+        while not stop_event.is_set():
+            try:
+                headers = {
+                    "Authorization": f"Key {self.api_key}",
+                    "Content-Type": "application/json",
+                }
+                response = session.post(
+                    f"{self.base_url}/stream-heartbeat/{job_id}",
+                    headers=headers,
+                    params={"request_session_token": session_token},
+                )
+                response.raise_for_status()
+            except Exception as e:
+                if not stop_event.is_set():  # Only log if we weren't stopping anyway
+                    print(f"Heartbeat failed for job {job_id}: {e}")
+
+            # Use time.sleep instead of asyncio.sleep since this is synchronous
+            time.sleep(self.HEARTBEAT_INTERVAL_SECONDS)
+
+    @contextmanager
+    def stream_heartbeat_session(self, job_id: str, session_token: str) -> Generator[requests.Session, None, None]:
+        """Context manager that handles session registration and heartbeat."""
+        session = requests.Session()
+        stop_heartbeat = threading.Event()
+
+        # Run this concurrently in a thread so we can not block main SDK path/behavior
+        # but still run heartbeat requests
+        with ThreadPoolExecutor(max_workers=1) as executor:
+            future = executor.submit(
+                self.start_heartbeat,
+                job_id,
+                session_token,
+                session,
+                stop_heartbeat
+            )
+
+            try:
+                yield session
+            finally:
+                # Signal stop and cleanup
+                stop_heartbeat.set()
+                future.result()  # Wait for heartbeat to finish
+                self.unregister_stream_listener(job_id, session_token)
+                session.close()
+
+    def attach(self, job_id):
+        """
+        Attach to an existing job and stream its progress.
+
+        Args:
+            job_id (str): The ID of the job to attach to
+        """
+
+        s = requests.Session()
+        payload = {
+            "job_id": job_id,
+        }
+        pbar = None
+
+        headers = {
+            "Authorization": f"Key {self.api_key}",
+            "Content-Type": "application/json",
+        }
+
+        with yaspin(
+                SPINNER,
+                text=to_colored_text("Looking for job..."),
+                color=YASPIN_COLOR,
+        ) as spinner:
+            # Get job information from list-jobs endpoint
+            # TODO(cooper) we should add a get jobs endpoint:
+            # GET /jobs/{job_id}
+            jobs_response = s.get(
+                f"{self.base_url}/list-jobs",
+                headers=headers
+            )
+            jobs_response.raise_for_status()
+
+            # Find the specific job we want to attach to
+            job = next(
+                (job for job in jobs_response.json()["jobs"] if job["job_id"] == job_id),
+                None
+            )
+
+            if not job:
+                spinner.write(to_colored_text(f"Job {job_id} not found", state="fail"))
+                return
+
+            match job.get("status"):
+                case "SUCCEEDED":
+                    spinner.write(
+                        to_colored_text(
+                            f"Job already completed. You can obtain the results with `sutro jobs results {job_id}`"
+                        )
+                    )
+                    return
+                case "FAILED":
+                    spinner.write(to_colored_text("❌ Job is in failed state.", state="fail"))
+                    return
+                case "CANCELLED":
+                    spinner.write(to_colored_text("❌ Job was cancelled.", state="fail"))
+                    return
+                case _:
+                    spinner.write(to_colored_text("✔ Job found!", state="success"))
+
+        total_rows = job["num_rows"]
+        success = False
+
+        session_token = self.register_stream_listener(job_id)
+
+        with self.stream_heartbeat_session(job_id, session_token) as s:
+            with s.get(
+                    f"{self.base_url}/stream-job-progress/{job_id}?request_session_token={session_token}",
+                    headers=headers,
+                    stream=True,
+            ) as streaming_response:
+                streaming_response.raise_for_status()
+                spinner = yaspin(
+                    SPINNER,
+                    text=to_colored_text("Awaiting status updates..."),
+                    color=YASPIN_COLOR,
+                )
+                spinner.start()
+                for line in streaming_response.iter_lines():
+                    if line:
+                        try:
+                            json_obj = json.loads(line)
+                        except json.JSONDecodeError:
+                            print("Error: ", line, flush=True)
+                            continue
+
+                        if json_obj["update_type"] == "progress":
+                            if pbar is None:
+                                spinner.stop()
+                                postfix = f"Input tokens processed: 0"
+                                pbar = self.fancy_tqdm(
+                                    total=total_rows,
+                                    desc="Progress",
+                                    style=1,
+                                    postfix=postfix,
+                                )
+                            if json_obj["result"] > pbar.n:
+                                pbar.update(json_obj["result"] - pbar.n)
+                                pbar.refresh()
+                            if json_obj["result"] == total_rows:
+                                pbar.close()
+                                success = True
+                        elif json_obj["update_type"] == "tokens":
+                            if pbar is not None:
+                                pbar.postfix = f"Input tokens processed: {json_obj['result']['input_tokens']}, Tokens generated: {json_obj['result']['output_tokens']}, Total tokens/s: {json_obj['result'].get('total_tokens_processed_per_second')}"
+                                pbar.refresh()
+
+                if success:
+                    spinner.write(
+                        to_colored_text(
+                            f"✔ Job succeeded. Use `sutro jobs results {job_id}` to obtain results.",
+                            state="success",
+                        )
+                    )
+                    spinner.stop()
+
+
+
+    def fancy_tqdm(
+        self,
+        total: int,
+        desc: str = "Progress",
+        color: str = "blue",
+        style=1,
+        postfix: str = None,
+    ):
+        """
+        Creates a customized tqdm progress bar with different styling options.
+
+        Args:
+            total (int): Total iterations
+            desc (str): Description for the progress bar
+            color (str): Color of the progress bar (green, blue, red, yellow, magenta)
+            style (int): Style preset (1-4)
+            postfix (str): Postfix for the progress bar
+        """
+
+        # Style presets
+        style_presets = {
+            1: {
+                "bar_format": "{l_bar}{bar:30}| {n_fmt}/{total_fmt} | {percentage:3.0f}% {postfix}",
+                "ascii": "░▒█",
+            },
+            2: {
+                "bar_format": "╢{l_bar}{bar:30}╟ {percentage:3.0f}%",
+                "ascii": "▁▂▃▄▅▆▇█",
+            },
+            3: {
+                "bar_format": "{desc}: |{bar}| {percentage:3.0f}% [{elapsed}<{remaining}]",
+                "ascii": "◯◔◑◕●",
+            },
+            4: {
+                "bar_format": "⏳ {desc} {percentage:3.0f}% |{bar}| {n_fmt}/{total_fmt}",
+                "ascii": "⬜⬛",
+            },
+            5: {
+                "bar_format": "⏳ {desc} {percentage:3.0f}% |{bar}| {n_fmt}/{total_fmt}",
+                "ascii": "▏▎▍▌▋▊▉█",
+            },
+        }
+
+        # Get style configuration
+        style_config = style_presets.get(style, style_presets[1])
+
+        return tqdm(
+            total=total,
+            desc=desc,
+            colour=color,
+            bar_format=style_config["bar_format"],
+            ascii=style_config["ascii"],
+            ncols=80,
+            dynamic_ncols=True,
+            smoothing=0.3,
+            leave=True,
+            postfix=postfix,
+        )
+
+    def list_jobs(self):
+        """
+        List all jobs.
+
+        This method retrieves a list of all jobs associated with the API key.
+
+        Returns:
+            list: A list of job details.
+        """
+        endpoint = f"{self.base_url}/list-jobs"
+        headers = {
+            "Authorization": f"Key {self.api_key}",
+            "Content-Type": "application/json",
+        }
+
+        with yaspin(
+            SPINNER, text=to_colored_text("Fetching jobs"), color=YASPIN_COLOR
+        ) as spinner:
+            response = requests.get(endpoint, headers=headers)
+            if response.status_code != 200:
+                spinner.write(
+                    to_colored_text(
+                        f"Bad status code: {response.status_code}", state="fail"
+                    )
+                )
+                spinner.stop()
+                print(to_colored_text(response.json(), state="fail"))
+                return
+        return response.json()["jobs"]
+
+    def get_job_status(self, job_id: str):
+        """
+        Get the status of a job by its ID.
+
+        This method retrieves the status of a job using its unique identifier.
+
+        Args:
+            job_id (str): The ID of the job to retrieve the status for.
+
+        Returns:
+            str: The status of the job.
+        """
+        endpoint = f"{self.base_url}/job-status/{job_id}"
+        headers = {
+            "Authorization": f"Key {self.api_key}",
+            "Content-Type": "application/json",
+        }
+        with yaspin(
+            SPINNER,
+            text=to_colored_text(f"Checking job status with ID: {job_id}"),
+            color=YASPIN_COLOR,
+        ) as spinner:
+            response = requests.get(endpoint, headers=headers)
+            if response.status_code != 200:
+                spinner.write(
+                    to_colored_text(
+                        f"Bad status code: {response.status_code}", state="fail"
+                    )
+                )
+                spinner.stop()
+                print(to_colored_text(response.json(), state="fail"))
+                return
+            spinner.write(to_colored_text("✔ Job status retrieved!", state="success"))
+        return response.json()["job_status"][job_id]
+
+    def get_job_results(
+        self,
+        job_id: str,
+        include_inputs: bool = False,
+        include_cumulative_logprobs: bool = False,
+    ):
+        """
+        Get the results of a job by its ID.
+
+        This method retrieves the results of a job using its unique identifier.
+
+        Args:
+            job_id (str): The ID of the job to retrieve the results for.
+            include_inputs (bool, optional): Whether to include the inputs in the results. Defaults to False.
+            include_cumulative_logprobs (bool, optional): Whether to include the cumulative logprobs in the results. Defaults to False.
+
+        Returns:
+            list: The results of the job.
+        """
+        endpoint = f"{self.base_url}/job-results"
+        payload = {
+            "job_id": job_id,
+            "include_inputs": include_inputs,
+            "include_cumulative_logprobs": include_cumulative_logprobs,
+        }
+        headers = {
+            "Authorization": f"Key {self.api_key}",
+            "Content-Type": "application/json",
+        }
+        with yaspin(
+            SPINNER,
+            text=to_colored_text(f"Gathering results from job: {job_id}"),
+            color=YASPIN_COLOR,
+        ) as spinner:
+            response = requests.post(
+                endpoint, data=json.dumps(payload), headers=headers
+            )
+            if response.status_code == 200:
+                spinner.write(
+                    to_colored_text("✔ Job results retrieved", state="success")
+                )
+            else:
+                spinner.write(
+                    to_colored_text(
+                        f"Bad status code: {response.status_code}", state="fail"
+                    )
+                )
+                spinner.stop()
+                print(to_colored_text(response.json(), state="fail"))
+                return
+        return response.json()["results"]
+
+    def cancel_job(self, job_id: str):
+        """
+        Cancel a job by its ID.
+
+        This method allows you to cancel a job using its unique identifier.
+
+        Args:
+            job_id (str): The ID of the job to cancel.
+
+        Returns:
+            dict: The status of the job.
+        """
+        endpoint = f"{self.base_url}/job-cancel/{job_id}"
+        headers = {
+            "Authorization": f"Key {self.api_key}",
+            "Content-Type": "application/json",
+        }
+        with yaspin(
+            SPINNER,
+            text=to_colored_text(f"Cancelling job: {job_id}"),
+            color=YASPIN_COLOR,
+        ) as spinner:
+            response = requests.get(endpoint, headers=headers)
+            if response.status_code == 200:
+                spinner.write(to_colored_text("✔ Job cancelled", state="success"))
+            else:
+                spinner.write(to_colored_text("Failed to cancel job", state="fail"))
+                spinner.stop()
+                print(to_colored_text(response.json(), state="fail"))
+                return
+        return response.json()
+
+    def create_stage(self):
+        """
+        Create a new stage.
+
+        This method creates a new stage and returns its ID.
+
+        Returns:
+            str: The ID of the new stage.
+        """
+        endpoint = f"{self.base_url}/create-stage"
+        headers = {
+            "Authorization": f"Key {self.api_key}",
+            "Content-Type": "application/json",
+        }
+        with yaspin(
+            SPINNER, text=to_colored_text("Creating stage"), color=YASPIN_COLOR
+        ) as spinner:
+            response = requests.get(endpoint, headers=headers)
+            if response.status_code != 200:
+                spinner.write(
+                    to_colored_text(
+                        f"Bad status code: {response.status_code}", state="fail"
+                    )
+                )
+                spinner.stop()
+                print(to_colored_text(response.json(), state="fail"))
+                return
+            stage_id = response.json()["stage_id"]
+            spinner.write(
+                to_colored_text(f"✔ Stage created with ID: {stage_id}", state="success")
+            )
+        return stage_id
+
+    def upload_to_stage(
+        self,
+        stage_id: Union[List[str], str] = None,
+        file_paths: Union[List[str], str] = None,
+        verify_ssl: bool = True,
+    ):
+        """
+        Upload data to a stage.
+
+        This method uploads files to a stage. Accepts a stage ID and file paths. If only a single parameter is provided, it will be interpreted as the file paths.
+
+        Args:
+            stage_id (str): The ID of the stage to upload to. If not provided, a new stage will be created.
+            file_paths (Union[List[str], str]): A list of paths to the files to upload, or a single path to a collection of files.
+            verify_ssl (bool): Whether to verify SSL certificates. Set to False to bypass SSL verification for troubleshooting.
+
+        Returns:
+            dict: The response from the API.
+        """
+        # when only a single parameter is provided, it is interpreted as the file paths
+        if file_paths is None and stage_id is not None:
+            file_paths = stage_id
+            stage_id = None
+
+        if file_paths is None:
+            raise ValueError("File paths must be provided")
+
+        if stage_id is None:
+            stage_id = self.create_stage()
+
+        endpoint = f"{self.base_url}/upload-to-stage"
+
+        if isinstance(file_paths, str):
+            # check if the file path is a directory
+            if os.path.isdir(file_paths):
+                file_paths = [
+                    os.path.join(file_paths, f) for f in os.listdir(file_paths)
+                ]
+                if len(file_paths) == 0:
+                    raise ValueError("No files found in the directory")
+            else:
+                file_paths = [file_paths]
+
+        with yaspin(
+            SPINNER,
+            text=to_colored_text(f"Uploading files to stage: {stage_id}"),
+            color=YASPIN_COLOR,
+        ) as spinner:
+            count = 0
+            for file_path in file_paths:
+                file_name = os.path.basename(file_path)
+
+                files = {
+                    "file": (
+                        file_name,
+                        open(file_path, "rb"),
+                        "application/octet-stream",
+                    )
+                }
+
+                payload = {
+                    "stage_id": stage_id,
+                }
+
+                headers = {
+                    "Authorization": f"Key {self.api_key}"}
+
+                count += 1
+                spinner.write(
+                    to_colored_text(
+                        f"Uploading file {count}/{len(file_paths)} to stage: {stage_id}"
+                    )
+                )
+
+                try:
+                    response = requests.post(
+                        endpoint, headers=headers, data=payload, files=files
+                    )
+                    if response.status_code != 200:
+                        # Stop spinner before showing error to avoid terminal width error
+                        spinner.stop()
+                        print(
+                            to_colored_text(
+                                f"Error: HTTP {response.status_code}", state="fail"
+                            )
+                        )
+                        print(to_colored_text(response.json(), state="fail"))
+                        return
+
+                except requests.exceptions.RequestException as e:
+                    # Stop spinner before showing error to avoid terminal width error
+                    spinner.stop()
+                    print(to_colored_text(f"Upload failed: {str(e)}", state="fail"))
+                    return
+
+            spinner.write(
+                to_colored_text(
+                    f"✔ {count} files successfully uploaded to stage", state="success"
+                )
+            )
+        return stage_id
+
+    def list_stages(self):
+        endpoint = f"{self.base_url}/list-stages"
+        headers = {
+            "Authorization": f"Key {self.api_key}",
+            "Content-Type": "application/json",
+        }
+        with yaspin(
+            SPINNER, text=to_colored_text("Retrieving stages"), color=YASPIN_COLOR
+        ) as spinner:
+            response = requests.post(endpoint, headers=headers)
+            if response.status_code != 200:
+                spinner.fail(
+                    to_colored_text(
+                        f"Bad status code: {response.status_code}", state="fail"
+                    )
+                )
+                print(to_colored_text(f"Error: {response.json()}", state="fail"))
+                return
+            spinner.write(to_colored_text("✔ Stages retrieved", state="success"))
+        return response.json()["stages"]
+
+    def list_stage_files(self, stage_id: str):
+        endpoint = f"{self.base_url}/list-stage-files"
+        headers = {
+            "Authorization": f"Key {self.api_key}",
+            "Content-Type": "application/json",
+        }
+        payload = {
+            "stage_id": stage_id,
+        }
+        with yaspin(
+            SPINNER,
+            text=to_colored_text(f"Listing files in stage: {stage_id}"),
+            color=YASPIN_COLOR,
+        ) as spinner:
+            response = requests.post(
+                endpoint, headers=headers, data=json.dumps(payload)
+            )
+            if response.status_code != 200:
+                spinner.fail(
+                    to_colored_text(
+                        f"Bad status code: {response.status_code}", state="fail"
+                    )
+                )
+                print(to_colored_text(f"Error: {response.json()}", state="fail"))
+                return
+            spinner.write(
+                to_colored_text(f"✔ Files listed in stage: {stage_id}", state="success")
+            )
+        return response.json()["files"]
+
+    def download_from_stage(
+        self,
+        stage_id: str,
+        files: Union[List[str], str] = None,
+        output_path: str = None,
+    ):
+        endpoint = f"{self.base_url}/download-from-stage"
+
+        if files is None:
+            files = self.list_stage_files(stage_id)
+        elif isinstance(files, str):
+            files = [files]
+
+        if not files:
+            print(
+                to_colored_text(
+                    f"Couldn't find files for stage ID: {stage_id}", state="fail"
+                )
+            )
+            return
+
+        # if no output path is provided, save the files to the current working directory
+        if output_path is None:
+            output_path = os.getcwd()
+
+        with yaspin(
+            SPINNER,
+            text=to_colored_text(f"Downloading files from stage: {stage_id}"),
+            color=YASPIN_COLOR,
+        ) as spinner:
+            count = 0
+            for file in files:
+                headers = {
+                    "Authorization": f"Key {self.api_key}",
+                    "Content-Type": "application/json",
+                }
+                payload = {
+                    "stage_id": stage_id,
+                    "file_name": file,
+                }
+                spinner.text = to_colored_text(
+                    f"Downloading file {count + 1}/{len(files)} from stage: {stage_id}"
+                )
+                response = requests.post(
+                    endpoint, headers=headers, data=json.dumps(payload)
+                )
+                if response.status_code != 200:
+                    spinner.fail(
+                        to_colored_text(
+                            f"Bad status code: {response.status_code}", state="fail"
+                        )
+                    )
+                    print(to_colored_text(f"Error: {response.json()}", state="fail"))
+                    return
+                file_content = response.content
+                with open(os.path.join(output_path, file), "wb") as f:
+                    f.write(file_content)
+                count += 1
+            spinner.write(
+                to_colored_text(
+                    f"✔ {count} files successfully downloaded from stage: {stage_id}",
+                    state="success",
+                )
+            )
+
+    def try_authentication(self, api_key: str):
+        """
+        Try to authenticate with the API key.
+
+        This method allows you to authenticate with the API key.
+
+        Args:
+            api_key (str): The API key to authenticate with.
+
+        Returns:
+            dict: The status of the authentication.
+        """
+        endpoint = f"{self.base_url}/try-authentication"
+        headers = {
+            "Authorization": f"Key {api_key}",
+            "Content-Type": "application/json",
+        }
+        with yaspin(
+            SPINNER, text=to_colored_text("Checking API key"), color=YASPIN_COLOR
+        ) as spinner:
+            response = requests.get(endpoint, headers=headers)
+            if response.status_code == 200:
+                spinner.write(to_colored_text("✔"))
+            else:
+                spinner.write(
+                    to_colored_text(
+                        f"API key failed to authenticate: {response.status_code}",
+                        state="fail",
+                    )
+                )
+                return
+        return response.json()
+
+    def get_quotas(self):
+        endpoint = f"{self.base_url}/get-quotas"
+        headers = {
+            "Authorization": f"Key {self.api_key}",
+            "Content-Type": "application/json",
+        }
+        with yaspin(
+            SPINNER, text=to_colored_text("Fetching quotas"), color=YASPIN_COLOR
+        ) as spinner:
+            response = requests.get(endpoint, headers=headers)
+            if response.status_code != 200:
+                spinner.fail(
+                    to_colored_text(
+                        f"Bad status code: {response.status_code}", state="fail"
+                    )
+                )
+                print(to_colored_text(f"Error: {response.json()}", state="fail"))
+                return
+        return response.json()["quotas"]