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

sutro/cli.py

@@ -275,34 +275,34 @@ def cancel(job_id):
 
 
 @cli.group()
-def stages():
-    """Manage stages."""
+def datasets():
+    """Manage datasets."""
     pass
 
 
-@stages.command()
+@datasets.command()
 def create():
-    """Create a new stage."""
+    """Create a new dataset."""
     sdk = get_sdk()
-    stage_id = sdk.create_stage()
-    if not stage_id:
+    dataset_id = sdk.create_dataset()
+    if not dataset_id:
         return
     click.echo(
         Fore.GREEN
-        + f"Stage created successfully. Stage ID: {stage_id}"
+        + f"Dataset created successfully. Dataset ID: {dataset_id}"
         + Style.RESET_ALL
     )
 
 
-@stages.command()
+@datasets.command()
 def list():
-    """List all stages."""
+    """List all datasets."""
     sdk = get_sdk()
-    stages = sdk.list_stages()
-    if stages is None or len(stages) == 0:
-        click.echo(Fore.YELLOW + "No stages found." + Style.RESET_ALL)
+    datasets = sdk.list_datasets()
+    if datasets is None or len(datasets) == 0:
+        click.echo(Fore.YELLOW + "No datasets found." + Style.RESET_ALL)
         return
-    df = pl.DataFrame(stages)
+    df = pl.DataFrame(datasets)
 
     df = df.with_columns(
         pl.col("schema")
@@ -319,37 +319,37 @@ def list():
         print(df.select(pl.all()))
 
 
-@stages.command()
-@click.argument("stage_id")
-def files(stage_id):
-    """List all files in a stage."""
+@datasets.command()
+@click.argument("dataset_id")
+def files(dataset_id):
+    """List all files in a dataset."""
     sdk = get_sdk()
-    files = sdk.list_stage_files(stage_id)
+    files = sdk.list_dataset_files(dataset_id)
     if not files:
         return
 
-    print(Fore.YELLOW + "Files in stage " + stage_id + ":" + Style.RESET_ALL)
+    print(Fore.YELLOW + "Files in dataset " + dataset_id + ":" + Style.RESET_ALL)
     for file in files:
         print(f"\t{file}")
 
 
-@stages.command()
-@click.argument("stage_id", required=False)
+@datasets.command()
+@click.argument("dataset_id", required=False)
 @click.argument("file_path")
-def upload(file_path, stage_id):
-    """Upload files to a stage. You can provide a single file path or a directory path to upload all files in the directory."""
+def upload(file_path, dataset_id):
+    """Upload files to a dataset. You can provide a single file path or a directory path to upload all files in the directory."""
     sdk = get_sdk()
-    sdk.upload_to_stage(file_path, stage_id)
+    sdk.upload_to_dataset(file_path, dataset_id)
 
 
-@stages.command()
-@click.argument("stage_id")
+@datasets.command()
+@click.argument("dataset_id")
 @click.argument("file_name", required=False)
 @click.argument("output_path", required=False)
-def download(stage_id, file_name=None, output_path=None):
-    """Download a file/files from a stage. If no files are provided, all files in the stage will be downloaded. If no output path is provided, the file will be saved to the current working directory."""
+def download(dataset_id, file_name=None, output_path=None):
+    """Download a file/files from a dataset. If no files are provided, all files in the dataset will be downloaded. If no output path is provided, the file will be saved to the current working directory."""
     sdk = get_sdk()
-    files = sdk.download_from_stage(stage_id, [file_name], output_path)
+    files = sdk.download_from_dataset(dataset_id, [file_name], output_path)
     if not files:
         return
     for file in files:

sutro/sdk.py

@@ -114,7 +114,7 @@ class Sutro:
                 raise ValueError("Column name must be specified for DataFrame input")
             input_data = data[column].to_list()
         elif isinstance(data, str):
-            if data.startswith("stage-"):
+            if data.startswith("dataset-"):
                 input_data = data + ":" + column
             else:
                 file_ext = os.path.splitext(data)[1].lower()
@@ -172,12 +172,12 @@ class Sutro:
         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.
+        It supports various data types such as lists, pandas DataFrames, polars DataFrames, file paths and datasets.
 
         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.
+            column (str, optional): The column name to use for inference. Required if data is a DataFrame, file path, or dataset.
             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.
@@ -437,8 +437,10 @@ class Sutro:
                 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)
+            for _ in range(self.HEARTBEAT_INTERVAL_SECONDS):
+                if stop_event.is_set():
+                    break
+                time.sleep(1)
 
     @contextmanager
     def stream_heartbeat_session(self, job_id: str, session_token: str) -> Generator[requests.Session, None, None]:
@@ -449,7 +451,7 @@ class Sutro:
         # 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(
+            executor.submit(
                 self.start_heartbeat,
                 job_id,
                 session_token,
@@ -462,7 +464,6 @@ class Sutro:
             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()
 
@@ -792,22 +793,22 @@ class Sutro:
                 return
         return response.json()
 
-    def create_stage(self):
+    def create_dataset(self):
         """
-        Create a new stage.
+        Create a new dataset.
 
-        This method creates a new stage and returns its ID.
+        This method creates a new empty dataset and returns its ID.
 
         Returns:
-            str: The ID of the new stage.
+            str: The ID of the new dataset.
         """
-        endpoint = f"{self.base_url}/create-stage"
+        endpoint = f"{self.base_url}/create-dataset"
         headers = {
             "Authorization": f"Key {self.api_key}",
             "Content-Type": "application/json",
         }
         with yaspin(
-            SPINNER, text=to_colored_text("Creating stage"), color=YASPIN_COLOR
+            SPINNER, text=to_colored_text("Creating dataset"), color=YASPIN_COLOR
         ) as spinner:
             response = requests.get(endpoint, headers=headers)
             if response.status_code != 200:
@@ -819,25 +820,25 @@ class Sutro:
                 spinner.stop()
                 print(to_colored_text(response.json(), state="fail"))
                 return
-            stage_id = response.json()["stage_id"]
+            dataset_id = response.json()["dataset_id"]
             spinner.write(
-                to_colored_text(f"✔ Stage created with ID: {stage_id}", state="success")
+                to_colored_text(f"✔ Dataset created with ID: {dataset_id}", state="success")
             )
-        return stage_id
+        return dataset_id
 
-    def upload_to_stage(
+    def upload_to_dataset(
         self,
-        stage_id: Union[List[str], str] = None,
+        dataset_id: Union[List[str], str] = None,
         file_paths: Union[List[str], str] = None,
         verify_ssl: bool = True,
     ):
         """
-        Upload data to a stage.
+        Upload data to a dataset.
 
-        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.
+        This method uploads files to a dataset. Accepts a dataset 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.
+            dataset_id (str): The ID of the dataset to upload to. If not provided, a new dataset 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.
 
@@ -845,17 +846,17 @@ class Sutro:
             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 and dataset_id is not None:
+            file_paths = dataset_id
+            dataset_id = None
 
         if file_paths is None:
             raise ValueError("File paths must be provided")
 
-        if stage_id is None:
-            stage_id = self.create_stage()
+        if dataset_id is None:
+            dataset_id = self.create_dataset()
 
-        endpoint = f"{self.base_url}/upload-to-stage"
+        endpoint = f"{self.base_url}/upload-to-dataset"
 
         if isinstance(file_paths, str):
             # check if the file path is a directory
@@ -870,7 +871,7 @@ class Sutro:
 
         with yaspin(
             SPINNER,
-            text=to_colored_text(f"Uploading files to stage: {stage_id}"),
+            text=to_colored_text(f"Uploading files to dataset: {dataset_id}"),
             color=YASPIN_COLOR,
         ) as spinner:
             count = 0
@@ -886,7 +887,7 @@ class Sutro:
                 }
 
                 payload = {
-                    "stage_id": stage_id,
+                    "dataset_id": dataset_id,
                 }
 
                 headers = {
@@ -895,7 +896,7 @@ class Sutro:
                 count += 1
                 spinner.write(
                     to_colored_text(
-                        f"Uploading file {count}/{len(file_paths)} to stage: {stage_id}"
+                        f"Uploading file {count}/{len(file_paths)} to dataset: {dataset_id}"
                     )
                 )
 
@@ -922,19 +923,19 @@ class Sutro:
 
             spinner.write(
                 to_colored_text(
-                    f"✔ {count} files successfully uploaded to stage", state="success"
+                    f"✔ {count} files successfully uploaded to dataset", state="success"
                 )
             )
-        return stage_id
+        return dataset_id
 
-    def list_stages(self):
-        endpoint = f"{self.base_url}/list-stages"
+    def list_datasets(self):
+        endpoint = f"{self.base_url}/list-datasets"
         headers = {
             "Authorization": f"Key {self.api_key}",
             "Content-Type": "application/json",
         }
         with yaspin(
-            SPINNER, text=to_colored_text("Retrieving stages"), color=YASPIN_COLOR
+            SPINNER, text=to_colored_text("Retrieving datasets"), color=YASPIN_COLOR
         ) as spinner:
             response = requests.post(endpoint, headers=headers)
             if response.status_code != 200:
@@ -945,21 +946,21 @@ class Sutro:
                 )
                 print(to_colored_text(f"Error: {response.json()}", state="fail"))
                 return
-            spinner.write(to_colored_text("✔ Stages retrieved", state="success"))
-        return response.json()["stages"]
+            spinner.write(to_colored_text("✔ Datasets retrieved", state="success"))
+        return response.json()["datasets"]
 
-    def list_stage_files(self, stage_id: str):
-        endpoint = f"{self.base_url}/list-stage-files"
+    def list_dataset_files(self, dataset_id: str):
+        endpoint = f"{self.base_url}/list-dataset-files"
         headers = {
             "Authorization": f"Key {self.api_key}",
             "Content-Type": "application/json",
         }
         payload = {
-            "stage_id": stage_id,
+            "dataset_id": dataset_id,
         }
         with yaspin(
             SPINNER,
-            text=to_colored_text(f"Listing files in stage: {stage_id}"),
+            text=to_colored_text(f"Listing files in dataset: {dataset_id}"),
             color=YASPIN_COLOR,
         ) as spinner:
             response = requests.post(
@@ -974,27 +975,27 @@ class Sutro:
                 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")
+                to_colored_text(f"✔ Files listed in dataset: {dataset_id}", state="success")
             )
         return response.json()["files"]
 
-    def download_from_stage(
+    def download_from_dataset(
         self,
-        stage_id: str,
+        dataset_id: str,
         files: Union[List[str], str] = None,
         output_path: str = None,
     ):
-        endpoint = f"{self.base_url}/download-from-stage"
+        endpoint = f"{self.base_url}/download-from-dataset"
 
         if files is None:
-            files = self.list_stage_files(stage_id)
+            files = self.list_dataset_files(dataset_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"
+                    f"Couldn't find files for dataset ID: {dataset_id}", state="fail"
                 )
             )
             return
@@ -1005,7 +1006,7 @@ class Sutro:
 
         with yaspin(
             SPINNER,
-            text=to_colored_text(f"Downloading files from stage: {stage_id}"),
+            text=to_colored_text(f"Downloading files from dataset: {dataset_id}"),
             color=YASPIN_COLOR,
         ) as spinner:
             count = 0
@@ -1015,11 +1016,11 @@ class Sutro:
                     "Content-Type": "application/json",
                 }
                 payload = {
-                    "stage_id": stage_id,
+                    "dataset_id": dataset_id,
                     "file_name": file,
                 }
                 spinner.text = to_colored_text(
-                    f"Downloading file {count + 1}/{len(files)} from stage: {stage_id}"
+                    f"Downloading file {count + 1}/{len(files)} from dataset: {dataset_id}"
                 )
                 response = requests.post(
                     endpoint, headers=headers, data=json.dumps(payload)
@@ -1038,7 +1039,7 @@ class Sutro:
                 count += 1
             spinner.write(
                 to_colored_text(
-                    f"✔ {count} files successfully downloaded from stage: {stage_id}",
+                    f"✔ {count} files successfully downloaded from dataset: {dataset_id}",
                     state="success",
                 )
             )

sutro-0.1.14.dist-info/METADATA

@@ -0,0 +1,23 @@
+Metadata-Version: 2.4
+Name: sutro
+Version: 0.1.14
+Summary: Sutro Python SDK
+Project-URL: Homepage, https://sutro.sh
+Project-URL: Documentation, https://docs.sutro.sh
+License-Expression: Apache-2.0
+License-File: LICENSE
+Requires-Python: >=3.10
+Requires-Dist: click==8.1.7
+Requires-Dist: colorama==0.4.4
+Requires-Dist: numpy==2.1.1
+Requires-Dist: pandas==2.2.3
+Requires-Dist: polars==1.8.2
+Requires-Dist: pydantic==2.11.4
+Requires-Dist: requests==2.32.3
+Requires-Dist: tqdm==4.67.1
+Requires-Dist: yaspin==3.1.0
+Description-Content-Type: text/markdown
+
+# sutro-client
+
+The official Python client for Sutro. See [docs.sutro.sh](https://docs.sutro.sh/) for more information.

sutro-0.1.14.dist-info/RECORD

@@ -0,0 +1,8 @@
+sutro/__init__.py,sha256=yUiVwcZ8QamSqDdRHgzoANyTZ-x3cPzlt2Fs5OllR_w,402
+sutro/cli.py,sha256=6Qy9Vwaaho92HeO8YA_z1De4zp1dEFkSX3bEnLvdbkE,13203
+sutro/sdk.py,sha256=K86B8TNkNTSc1k0I6SVJldt6QE-ctnZmc8dMiOjK-PA,42439
+sutro-0.1.14.dist-info/METADATA,sha256=ZpPrBjBQz3CCiXtvvBnJgFyz4cekr9gzi_ETtWtWc-8,669
+sutro-0.1.14.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+sutro-0.1.14.dist-info/entry_points.txt,sha256=eXvr4dvMV4UmZgR0zmrY8KOmNpo64cJkhNDywiadRFM,40
+sutro-0.1.14.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+sutro-0.1.14.dist-info/RECORD,,

sutro-0.1.12.dist-info/METADATA

@@ -1,41 +0,0 @@
-Metadata-Version: 2.4
-Name: sutro
-Version: 0.1.12
-Summary: Sutro Python SDK
-Project-URL: Homepage, https://sutro.sh
-Project-URL: Documentation, https://docs.sutro.sh
-License-Expression: Apache-2.0
-License-File: LICENSE
-Requires-Python: >=3.10
-Requires-Dist: click==8.1.7
-Requires-Dist: colorama==0.4.4
-Requires-Dist: numpy==2.1.1
-Requires-Dist: pandas==2.2.3
-Requires-Dist: polars==1.8.2
-Requires-Dist: pydantic==2.11.4
-Requires-Dist: requests==2.32.3
-Requires-Dist: tqdm==4.67.1
-Requires-Dist: yaspin==3.1.0
-Description-Content-Type: text/markdown
-
-# sutro-client
-
-The official Python client for Sutro. See [docs.sutro.sh](https://docs.sutro.sh/) for more information.
-
-## Installing Locally (to test changes during development)
-
-Run `make install` from the root directory. This should remove the old builds and reinstall the package in your environment with the latest. You can run `uv pip list` to ensure the package is pointing at the local files instead of the PyPI package.
-
-## Creating releases
-
-Make sure you increment the version appropriately in `pyproject.toml`. Generally speaking we'll do patch versions for small tweaks, minor versions for large additions or changes to behavior, and probably do major releases once it makes sense. Since we're still in beta and `0.x.x` releases, its probably okay to add backwards-incompatible changes to minor releases, but we want to avoid this if possible. 
-
-To create a release, run: 
-
-`make release <version>` with `<version>` formatted like `0.1.1`
-
-It'll prompt you for an API key to PyPI, which you must have for it to work. 
-
-We also have a test PyPI account which you can use to test creating releases before pushing to the actual PyPI hub. I believe you can only create **one** release per version number, so it may be worth testing if you're paranoid about getting it right.
-
-Also make sure to update the docs and increment the docs version number to match the new release. Keeping these consistent will provide a better user experience. 

sutro-0.1.12.dist-info/RECORD

@@ -1,8 +0,0 @@
-sutro/__init__.py,sha256=yUiVwcZ8QamSqDdRHgzoANyTZ-x3cPzlt2Fs5OllR_w,402
-sutro/cli.py,sha256=ZvfMYNKq3iT-4TtGWmbI0qsEPTb0tjIGIiVMod1kH-s,13125
-sutro/sdk.py,sha256=dhlrIzqePS0QAh1ShonO5RbbPgYFpXFQL_gEcvoP8_c,42353
-sutro-0.1.12.dist-info/METADATA,sha256=FNpJnLVVwovtzFW6tiug46MWYPsJ4D73OAwdaew8FnA,1998
-sutro-0.1.12.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-sutro-0.1.12.dist-info/entry_points.txt,sha256=eXvr4dvMV4UmZgR0zmrY8KOmNpo64cJkhNDywiadRFM,40
-sutro-0.1.12.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-sutro-0.1.12.dist-info/RECORD,,