sutro 0.1.18__tar.gz → 0.1.19__tar.gz

{sutro-0.1.18 → sutro-0.1.19}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sutro
-Version: 0.1.18
+Version: 0.1.19
 Summary: Sutro Python SDK
 Project-URL: Homepage, https://sutro.sh
 Project-URL: Documentation, https://docs.sutro.sh

{sutro-0.1.18 → sutro-0.1.19}/pyproject.toml

@@ -9,7 +9,7 @@ installer = "uv"
 
 [project]
 name = "sutro"
-version = "0.1.18"
+version = "0.1.19"
 description = "Sutro Python SDK"
 readme = "README.md"
 requires-python = ">=3.10"

{sutro-0.1.18 → sutro-0.1.19}/sutro/sdk.py

@@ -97,6 +97,16 @@ def to_colored_text(
             # Default to blue for normal/processing states
             return f"{Fore.BLUE}{text}{Style.RESET_ALL}"
 
+# Isn't fully support in all terminals unfortunately. We should switch to Rich
+# at some point, but even Rich links aren't clickable on MacOS Terminal
+def make_clickable_link(url, text=None):
+    """
+    Create a clickable link for terminals that support OSC 8 hyperlinks.
+    Falls back to plain text for terminals that don't support it.
+    """
+    if text is None:
+        text = url
+    return f"\033]8;;{url}\033\\{text}\033]8;;\033\\"
 
 class Sutro:
     def __init__(
@@ -304,7 +314,7 @@ class Sutro:
                 else:
                     spinner.write(
                         to_colored_text(
-                            f"🛠️  Priority {job_priority} Job created with ID: {job_id}",
+                            f"🛠 Priority {job_priority} Job created with ID: {job_id}.",
                             state="success",
                         )
                     )
@@ -312,13 +322,14 @@ class Sutro:
                         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:
-            spinner.write(to_colored_text("Awaiting job start...", "info"))
+            spinner.write(to_colored_text("Awaiting job start...", ))
+            spinner.write(to_colored_text(f'Progress can also be monitored at: {make_clickable_link(f'https://app.sutro.sh/jobs/{job_id}')}'))
             started = self._await_job_start(job_id)
             if not started:
                 failure_reason = self._get_failure_reason(job_id)
@@ -599,6 +610,7 @@ class Sutro:
                     text=to_colored_text("Awaiting status updates..."),
                     color=YASPIN_COLOR,
                 )
+                spinner.write(to_colored_text(f'Progress can also be monitored at: {make_clickable_link(f'https://app.sutro.sh/jobs/{job_id}')}'))
                 spinner.start()
                 for line in streaming_response.iter_lines():
                     if line:
@@ -819,10 +831,12 @@ class Sutro:
                 return None
 
     def get_job_results(
-        self,
-        job_id: str,
-        include_inputs: bool = False,
-        include_cumulative_logprobs: bool = False,
+            self,
+            job_id: str,
+            include_inputs: bool = False,
+            include_cumulative_logprobs: bool = False,
+            with_original_df: pl.DataFrame | pd.DataFrame = None,
+            output_column: str = "inference_result",
     ):
         """
         Get the results of a job by its ID.
@@ -833,9 +847,11 @@ class Sutro:
             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.
+            with_original_df (pd.DataFrame | pl.DataFrame, optional): Original DataFrame to concatenate with results. Defaults to None.
+            output_column (str, optional): Name of the output column. Defaults to "inference_result".
 
         Returns:
-            list: The results of the job.
+            Union[pl.DataFrame, pd.DataFrame]: The results as a DataFrame. By default, returns polars.DataFrame; when with_original_df is an instance of pandas.DataFrame, returns pandas.DataFrame.
         """
         endpoint = f"{self.base_url}/job-results"
         payload = {
@@ -848,18 +864,14 @@ class Sutro:
             "Content-Type": "application/json",
         }
         with yaspin(
-            SPINNER,
-            text=to_colored_text(f"Gathering results from job: {job_id}"),
-            color=YASPIN_COLOR,
+                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:
+            if response.status_code != 200:
                 spinner.write(
                     to_colored_text(
                         f"Bad status code: {response.status_code}", state="fail"
@@ -867,8 +879,56 @@ class Sutro:
                 )
                 spinner.stop()
                 print(to_colored_text(response.json(), state="fail"))
-                return
-        return response.json()["results"]
+                return None
+
+            spinner.write(
+                to_colored_text("✔ Job results retrieved", state="success")
+            )
+
+        response_data = response.json()
+        results_df = pl.DataFrame(response_data["results"])
+
+
+        if len(results_df.columns ) == 1:
+            # Default column when API is only returning a list, and we construct the df
+            # from that
+            original_results_column = 'column_0'
+        else:
+            original_results_column = 'outputs'
+
+        results_df = results_df.rename({original_results_column: output_column})
+
+        # Ordering inputs col first seems most logical/useful
+        column_config = [
+            ('inputs', include_inputs),
+            (output_column, True),
+            ('cumulative_logprobs', include_cumulative_logprobs),
+        ]
+
+        columns_to_keep = [col for col, include in column_config
+                           if include and col in results_df.columns]
+
+        results_df = results_df.select(columns_to_keep)
+
+        # Handle concatenation with original DataFrame
+        if with_original_df is not None:
+            if isinstance(with_original_df, pd.DataFrame):
+                # Convert to polars for consistent handling
+                original_pl = pl.from_pandas(with_original_df)
+
+                combined_df = original_pl.with_columns(results_df)
+
+                # Convert back to pandas to match input type
+                return combined_df.to_pandas()
+
+            elif isinstance(with_original_df, pl.DataFrame):
+                return with_original_df.with_columns(results_df)
+
+        # Return pd.DataFrame type when appropriate
+        if with_original_df is None and isinstance(with_original_df, pd.DataFrame):
+            return results_df.to_pandas()
+
+        return results_df
 
     def cancel_job(self, job_id: str):
         """
@@ -1227,6 +1287,7 @@ class Sutro:
         with yaspin(
             SPINNER, text=to_colored_text("Awaiting job completion"), color=YASPIN_COLOR
         ) as spinner:
+            spinner.write(to_colored_text(f'Progress can also be monitored at: {make_clickable_link(f'https://app.sutro.sh/jobs/{job_id}')}'))
             while (time.time() - start_time) < timeout:
                 try:
                     status = self._fetch_job_status(job_id)