feat: add_centroid in WindowEvent.to_prompt_dict (#840)

* add centroid in window_dict output * black --------- Co-authored-by: Richard Abrich <[email protected]> Co-authored-by: Richard Abrich <[email protected]>
OpenAdaptAI · Jul 24, 2024 · 5118b39 · 5118b39
1 parent b288c07
commit 5118b39
Showing 1 changed file with 35 additions and 8 deletions.
diff --git a/openadapt/models.py b/openadapt/models.py
@@ -533,7 +533,12 @@ def scrub(self, scrubber: ScrubbingProvider | TextScrubbingMixin) -> None:
         if self.state is not None:
             self.state = scrubber.scrub_dict(self.state)
 
-    def to_prompt_dict(self, include_data: bool = True) -> dict[str, Any]:
+    def to_prompt_dict(
+        self,
+        include_data: bool = True,
+        add_centroid: bool = True,
+        remove_bbox: bool = False,
+    ) -> dict[str, Any]:
         """Convert into a dict, excluding properties not necessary for prompting.
 
         Args:
@@ -553,6 +558,29 @@ def to_prompt_dict(self, include_data: bool = True) -> dict[str, Any]:
                 # and not isinstance(getattr(models.WindowEvent, key), property)
             }
         )
+
+        if add_centroid:
+            left = window_dict["left"]
+            top = window_dict["top"]
+            width = window_dict["width"]
+            height = window_dict["height"]
+
+            # Compute the centroid of the bounding box
+            centroid_x = left + width / 2
+            centroid_y = top + height / 2
+
+            # Add centroid in the prompt dict { "centroid": }
+            window_dict["centroid"] = {
+                "x": centroid_x,
+                "y": centroid_y,
+            }
+
+        if remove_bbox:
+            window_dict.pop("left")
+            window_dict.pop("top")
+            window_dict.pop("width")
+            window_dict.pop("height")
+
         if "state" in window_dict:
             if include_data:
                 key_suffixes = [
@@ -574,14 +602,13 @@ def to_prompt_dict(self, include_data: bool = True) -> dict[str, Any]:
                     # from pprint import pformat
                     # logger.info(f"window_dict=\n{pformat(window_dict)}")
                     # import ipdb; ipdb.set_trace()
-                if "state" in window_dict:
-                    window_state = window_dict["state"]
-                    window_state["data"] = utils.clean_dict(
-                        utils.filter_keys(
-                            window_state["data"],
-                            key_suffixes,
-                        )
+                window_state = window_dict["state"]
+                window_state["data"] = utils.clean_dict(
+                    utils.filter_keys(
+                        window_state["data"],
+                        key_suffixes,
                     )
+                )
             else:
                 window_dict["state"].pop("data")
             window_dict["state"].pop("meta")