diff --git a/src/groundlight/client.py b/src/groundlight/client.py index 9a55a906..fbfa4154 100644 --- a/src/groundlight/client.py +++ b/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( diff --git a/test/unit/test_wait.py b/test/unit/test_wait.py new file mode 100644 index 00000000..6862cc36 --- /dev/null +++ b/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()