sutro 0.1.18__tar.gz → 0.1.20__tar.gz

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

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

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

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

{sutro-0.1.18 → sutro-0.1.20}/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__(
@@ -104,7 +114,6 @@ class Sutro:
     ):
         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):
         """
@@ -276,68 +285,67 @@ class Sutro:
         job_id = None
         t = f"Creating {'[dry run] ' if dry_run else ''}priority {job_priority} job"
         spinner_text = to_colored_text(t)
-        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")
+        try:
+            with yaspin(SPINNER, text=spinner_text, color=YASPIN_COLOR) as spinner:
+                response = requests.post(
+                    endpoint, data=json.dumps(payload), headers=headers
                 )
-                spinner.stop()
-                print(to_colored_text(response.json(), state="fail"))
-                return None
-            else:
-                job_id = response_data["results"]
-                if dry_run:
+                response_data = response.json()
+                if response.status_code != 200:
                     spinner.write(
-                        to_colored_text(f"Awaiting cost estimates with job ID: {job_id}. You can safely detach and retrieve the cost estimates later.", state="info")
+                        to_colored_text(f"Error: {response.status_code}", state="fail")
                     )
                     spinner.stop()
-                    self.await_job_completion(job_id, obtain_results=False)
-                    cost_estimate = self._get_job_cost_estimate(job_id)
-                    spinner.write(
-                        to_colored_text(f"✔ Cost estimates retrieved for job {job_id}: ${cost_estimate}", state="success")
-                    )
-                    return job_id
+                    print(to_colored_text(response.json(), state="fail"))
+                    return None
                 else:
-                    spinner.write(
-                        to_colored_text(
-                            f"🛠️  Priority {job_priority} Job created with ID: {job_id}",
-                            state="success",
+                    job_id = response_data["results"]
+                    if dry_run:
+                        spinner.write(
+                            to_colored_text(f"Awaiting cost estimates with job ID: {job_id}. You can safely detach and retrieve the cost estimates later.", state="info")
                         )
-                    )
-                    if not stay_attached:
+                        spinner.stop()
+                        self.await_job_completion(job_id, obtain_results=False)
+                        cost_estimate = self._get_job_cost_estimate(job_id)
+                        spinner.write(
+                            to_colored_text(f"✔ Cost estimates retrieved for job {job_id}: ${cost_estimate}", state="success")
+                        )
+                        return job_id
+                    else:
                         spinner.write(
                             to_colored_text(
-                                f"Use `so.get_job_status('{job_id}')` to check the status of the job."
+                                f"🛠 Priority {job_priority} Job created with ID: {job_id}.",
+                                state="success",
                             )
                         )
-                        return job_id
+                        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
+        except KeyboardInterrupt:
+            pass
+        finally:
+            if spinner:
+                spinner.stop()
 
         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)
                 spinner.write(to_colored_text(f"Failure reason: {failure_reason['message']}", "fail"))
                 return 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}",
+            try:
+                with requests.get(
+                        f"{self.base_url}/stream-job-progress/{job_id}",
                         headers=headers,
                         stream=True,
                 ) as streaming_response:
@@ -348,6 +356,13 @@ class Sutro:
                         color=YASPIN_COLOR,
                     )
                     spinner.start()
+
+                    token_state = {
+                        'input_tokens': 0,
+                        'output_tokens': 0,
+                            'total_tokens_processed_per_second': 0
+                    }
+
                     for line in streaming_response.iter_lines():
                         if line:
                             try:
@@ -370,12 +385,30 @@ class Sutro:
                                     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":
+                                # Update only the values that are present in this update
+                                # Currently, the way the progress stream endpoint is defined,
+                                # its possible to have updates come in that only have 1 or 2 fields
+                                new = {
+                                    k: v for k, v in json_obj.get('result', {}).items()
+                                    if k in token_state and v >= token_state[k]
+                                }
+                                token_state.update(new)
+
                                 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.postfix = f"Input tokens processed: {token_state['input_tokens']}, Output tokens generated: {token_state['output_tokens']}, Total tokens/s: {token_state['total_tokens_processed_per_second']}"
                                     pbar.refresh()
+
+            except KeyboardInterrupt:
+                pass
+            finally:
+                # Need to clean these up on keyboard exit otherwise it causes
+                # an error
+                if pbar is not None:
+                    pbar.close()
+                if spinner is not None:
+                    spinner.stop()
             if success:
                 spinner.text = to_colored_text(
                     "✔ Job succeeded. Obtaining results...", state="success"
@@ -440,87 +473,6 @@ class Sutro:
             return None
         return None
 
-    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}")
-
-            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]:
-        """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:
-            executor.submit(
-                self.start_heartbeat,
-                job_id,
-                session_token,
-                session,
-                stop_heartbeat
-            )
-
-            try:
-                yield session
-            finally:
-                # Signal stop and cleanup
-                stop_heartbeat.set()
-                self.unregister_stream_listener(job_id, session_token)
-                session.close()
 
     def attach(self, job_id):
         """
@@ -585,11 +537,9 @@ class Sutro:
         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:
+        try:
             with s.get(
-                    f"{self.base_url}/stream-job-progress/{job_id}?request_session_token={session_token}",
+                    f"{self.base_url}/stream-job-progress/{job_id}",
                     headers=headers,
                     stream=True,
             ) as streaming_response:
@@ -599,6 +549,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:
@@ -637,6 +588,13 @@ class Sutro:
                         )
                     )
                     spinner.stop()
+        except KeyboardInterrupt:
+            pass
+        finally:
+            if pbar:
+                pbar.close()
+            if spinner:
+                spinner.stop()
 
 
 
@@ -819,10 +777,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 +793,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 +810,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 +825,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 +1233,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)