From 79f8e2697319b4d54670498e965687cbebc220f8 Mon Sep 17 00:00:00 2001 From: CoreyEWood Date: Wed, 15 Jan 2025 15:40:29 -0800 Subject: [PATCH] check for edge query --- src/groundlight/client.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/src/groundlight/client.py b/src/groundlight/client.py index 0a8d4626..712fce74 100644 --- a/src/groundlight/client.py +++ b/src/groundlight/client.py @@ -999,6 +999,15 @@ def wait_for_confident_result( confidence_threshold = self.get_detector(image_query.detector_id).confidence_threshold confidence_above_thresh = partial(iq_is_confident, confidence_threshold=confidence_threshold) # type: ignore + + def is_from_edge(iq: ImageQuery) -> bool: + return iq.metadata and iq.metadata.get("is_from_edge", False) + + if is_from_edge(image_query) and not confidence_above_thresh(image_query): + # If the query is from the edge and the confidence is not above the threshold, it means the client wants + # only edge answers, so we don't want to poll the cloud. + return image_query + return self._wait_for_result(image_query, condition=confidence_above_thresh, timeout_sec=timeout_sec) def wait_for_ml_result(self, image_query: Union[ImageQuery, str], timeout_sec: float = 30.0) -> ImageQuery: