From ca16d3de1d1af9c2933119ffff9850731114141e Mon Sep 17 00:00:00 2001
From: The android_world Authors <no-reply@google.com>
Date: Mon, 20 Jan 2025 07:58:07 -0800
Subject: [PATCH] Add metadata field to UI element that can store arbitrary
 information.

Update json extraction to try conversion with json module as well.

PiperOrigin-RevId: 717534522
---
 android_world/agents/agent_utils.py       | 19 +++++++++++++++----
 android_world/env/representation_utils.py |  1 +
 2 files changed, 16 insertions(+), 4 deletions(-)

diff --git a/android_world/agents/agent_utils.py b/android_world/agents/agent_utils.py
index 86e9603..7cf9eb8 100644
--- a/android_world/agents/agent_utils.py
+++ b/android_world/agents/agent_utils.py
@@ -15,13 +15,16 @@
 """Utilities for agents."""
 
 import ast
+import json
 import re
-from typing import Any, Optional
+from typing import Any
 
 
-def extract_json(s: str) -> Optional[dict[str, Any]]:
+def extract_json(s: str) -> dict[str, Any] | None:
   """Extracts JSON from string.
 
+  Tries conversion with ast and json modules.
+
   Args:
     s: A string with a JSON in it. E.g., "{'hello': 'world'}" or from CoT:
       "let's think step-by-step, ..., {'hello': 'world'}".
@@ -35,7 +38,15 @@ def extract_json(s: str) -> Optional[dict[str, Any]]:
     try:
       return ast.literal_eval(match.group())
     except (SyntaxError, ValueError) as error:
-      print('Cannot extract JSON, skipping due to error %s', error)
-      return None
+      try:
+        # Try conversion with json module.
+        return json.loads(match.group())
+      except (SyntaxError, ValueError) as error2:
+        print(
+            'Cannot extract JSON, skipping due to errors %s and %s',
+            error,
+            error2,
+        )
+        return None
   else:
     return None
diff --git a/android_world/env/representation_utils.py b/android_world/env/representation_utils.py
index 3e556a2..a9c5912 100644
--- a/android_world/env/representation_utils.py
+++ b/android_world/env/representation_utils.py
@@ -74,6 +74,7 @@ class UIElement:
   resource_name: Optional[str] = None
   tooltip: Optional[str] = None
   resource_id: Optional[str] = None
+  metadata: Optional[dict[str, Any]] = None
 
 
 def accessibility_node_to_ui_element(