Don't wait for a confident answer if we know there won't be one

src/groundlight/client.py

@@ -25,6 +25,7 @@
     ImageQuery,
     PaginatedDetectorList,
     PaginatedImageQueryList,
+    Source,
 )
 from urllib3.exceptions import InsecureRequestWarning
 
@@ -970,6 +971,8 @@ def wait_for_confident_result(
             2. The timeout_sec is reached
             3. An error occurs
 
+            If the image query was answered on the edge or is done_processing, the result is returned immediately.
+
         **Example usage**::
 
             gl = Groundlight()
@@ -1002,9 +1005,29 @@ def wait_for_confident_result(
             :meth:`get_image_query` for checking result status without blocking
             :meth:`wait_for_ml_result` for waiting until the first ML result is available
         """
+
+        def should_stop_waiting(image_query: ImageQuery) -> bool:
+            """Checks if the image query should be returned immediately because no better answer is expected."""
+            if image_query.done_processing:
+                logger.debug(
+                    "The image query has completed escalating and will receive no new results, so we are "
+                    "returning it immediately."
+                )
+                return True
+            if image_query.result and image_query.result.source and image_query.result.source == Source.EDGE:
+                logger.debug(
+                    "The image query was answered on the edge, so we are returning it immediately and not waiting for "
+                    "a confident result."
+                )
+                return True
+            return False
+
         if isinstance(image_query, str):
             image_query = self.get_image_query(image_query)
 
+        if should_stop_waiting(image_query):
+            return image_query
+
         if confidence_threshold is None:
             confidence_threshold = self.get_detector(image_query.detector_id).confidence_threshold
 
@@ -1049,6 +1072,14 @@ def wait_for_ml_result(self, image_query: Union[ImageQuery, str], timeout_sec: f
             :meth:`get_image_query` for checking result status without blocking
             :meth:`wait_for_confident_result` for waiting until a confident result is available
         """
+        if isinstance(image_query, ImageQuery):
+            if image_query.result and image_query.result.source and image_query.result.source == Source.EDGE:
+                logger.debug(
+                    "The image query is from the edge, so we are returning it immediately and not waiting for an ML "
+                    "result."
+                )
+                return image_query
+        # TODO I think this is lying - it doesn't raise a TimeoutError if there is no ML result within timeout_sec
         return self._wait_for_result(image_query, condition=iq_is_answered, timeout_sec=timeout_sec)
 
     def _wait_for_result(

test/unit/test_wait.py

@@ -0,0 +1,20 @@
+from unittest.mock import patch
+
+from groundlight import ExperimentalApi
+from model import BinaryClassificationResult, ImageQuery, Label, Source
+
+
+def test_wait_for_confident_result_returns_immediately_when_no_better_result_expected(
+    gl_experimental: ExperimentalApi, initial_iq: ImageQuery
+):
+    with patch.object(gl_experimental, "_wait_for_result") as mock_wait_for_result:
+        # Shouldn't wait if the image query is done processing
+        initial_iq.done_processing = True
+        gl_experimental.wait_for_confident_result(initial_iq)
+        mock_wait_for_result.assert_not_called()
+
+        # Shouldn't wait if the result is from the edge
+        initial_iq.done_processing = False
+        initial_iq.result = BinaryClassificationResult(source=Source.EDGE, label=Label.YES)
+        gl_experimental.wait_for_confident_result(initial_iq)
+        mock_wait_for_result.assert_not_called()