Automated Code Change

PiperOrigin-RevId: 944185836

Author: h-joo

android_world/agents/infer.py

@@ -116,7 +116,7 @@ def __init__(
       raise RuntimeError('GCP API key not set.')
     genai.configure(api_key=os.environ['GCP_API_KEY'])
     self.llm = genai.GenerativeModel(
-        model_name,
+        model_name,  # pyrefly: ignore[bad-argument-type]
         safety_settings=None
         if enable_safety_checks
         else SAFETY_SETTINGS_BLOCK_NONE,
@@ -305,7 +305,7 @@ def predict_mm(
     for image in images:
       payload['messages'][0]['content'].append({
           'type': 'image_url',
-          'image_url': {
+          'image_url': {  # pyrefly: ignore[bad-assignment]
               'url': f'data:image/jpeg;base64,{self.encode_image(image)}'
           },
       })

android_world/agents/m3a.py

@@ -285,9 +285,9 @@ def _action_selection_prompt(
     The text prompt for action selection that will be sent to gpt4v.
   """
   if history:
-    history = '\n'.join(history)
+    history = '\n'.join(history)  # pyrefly: ignore[bad-assignment]
   else:
-    history = 'You just started, no action has been performed yet.'
+    history = 'You just started, no action has been performed yet.'  # pyrefly: ignore[bad-assignment]
 
   extra_guidelines = ''
   if additional_guidelines:
@@ -392,7 +392,7 @@ def step(self, goal: str) -> base_agent.AgentInteractionResult:
     before_ui_elements_list = _generate_ui_elements_description_list(
         before_ui_elements, logical_screen_size
     )
-    step_data['raw_screenshot'] = state.pixels.copy()
+    step_data['raw_screenshot'] = state.pixels.copy()  # pyrefly: ignore[bad-assignment]
     before_screenshot = state.pixels.copy()
     for index, ui_element in enumerate(before_ui_elements):
       if m3a_utils.validate_ui_element(ui_element, logical_screen_size):
@@ -404,7 +404,7 @@ def step(self, goal: str) -> base_agent.AgentInteractionResult:
             physical_frame_boundary,
             orientation,
         )
-    step_data['before_screenshot_with_som'] = before_screenshot.copy()
+    step_data['before_screenshot_with_som'] = before_screenshot.copy()  # pyrefly: ignore[bad-assignment]
 
     action_prompt = _action_selection_prompt(
         goal,
@@ -415,10 +415,10 @@ def step(self, goal: str) -> base_agent.AgentInteractionResult:
         before_ui_elements_list,
         self.additional_guidelines,
     )
-    step_data['action_prompt'] = action_prompt
+    step_data['action_prompt'] = action_prompt  # pyrefly: ignore[bad-assignment]
     action_output, is_safe, raw_response = self.llm.predict_mm(
         action_prompt,
-        [
+        [  # pyrefly: ignore[bad-argument-type]
             step_data['raw_screenshot'],
             before_screenshot,
         ],
@@ -431,7 +431,7 @@ def step(self, goal: str) -> base_agent.AgentInteractionResult:
 
     if not raw_response:
       raise RuntimeError('Error calling LLM in action selection phase.')
-    step_data['action_output'] = action_output
+    step_data['action_output'] = action_output  # pyrefly: ignore[bad-assignment]
     step_data['action_raw_response'] = raw_response
 
     reason, action = m3a_utils.parse_reason_action_output(action_output)
@@ -441,6 +441,7 @@ def step(self, goal: str) -> base_agent.AgentInteractionResult:
     if (not reason) or (not action):
       logging.info('Action prompt output is not in the correct format.')
       step_data['summary'] = (
+          # pyrefly: ignore[bad-assignment]
           'Output for action selection is not in the correct format, so no'
           ' action is performed.'
       )
@@ -453,17 +454,18 @@ def step(self, goal: str) -> base_agent.AgentInteractionResult:
 
     logging.info('Action: %s', action)
     logging.info('Reason: %s', reason)
-    step_data['action_reason'] = reason
+    step_data['action_reason'] = reason  # pyrefly: ignore[bad-assignment]
 
     try:
       converted_action = json_action.JSONAction(
-          **agent_utils.extract_json(action),
+          **agent_utils.extract_json(action),  # pyrefly: ignore[bad-unpacking]
       )
-      step_data['action_output_json'] = converted_action
+      step_data['action_output_json'] = converted_action  # pyrefly: ignore[bad-assignment]
     except Exception as e:  # pylint: disable=broad-exception-caught
       logging.info('Failed to convert the output to a valid action.')
       logging.info(str(e))
       step_data['summary'] = (
+          # pyrefly: ignore[bad-assignment]
           'Can not parse the output to a valid action. Please make sure to pick'
           ' the action from the list with required parameters (if any) in the'
           ' correct JSON format!'
@@ -482,14 +484,15 @@ def step(self, goal: str) -> base_agent.AgentInteractionResult:
         in ['click', 'long_press', 'input_text', 'scroll']
         and action_index is not None
     ):
-      if action_index >= num_ui_elements:
+      if action_index >= num_ui_elements:  # pyrefly: ignore[unsupported-operation]
         logging.info(
             'Index out of range, prediction index is %s, but the'
             ' UI element list only has %d elements.',
             action_index,
             num_ui_elements,
         )
         step_data['summary'] = (
+            # pyrefly: ignore[bad-assignment]
             'The parameter index is out of range. Remember the index must be in'
             ' the UI element list!'
         )
@@ -498,8 +501,8 @@ def step(self, goal: str) -> base_agent.AgentInteractionResult:
 
       # Add mark to the target element.
       m3a_utils.add_ui_element_mark(
-          step_data['raw_screenshot'],
-          before_ui_elements[action_index],
+          step_data['raw_screenshot'],  # pyrefly: ignore[bad-argument-type]
+          before_ui_elements[action_index],  # pyrefly: ignore[bad-index]
           action_index,
           logical_screen_size,
           physical_frame_boundary,
@@ -509,7 +512,7 @@ def step(self, goal: str) -> base_agent.AgentInteractionResult:
     if converted_action.action_type == 'status':
       if converted_action.goal_status == 'infeasible':
         logging.info('Agent stopped since it thinks mission impossible.')
-      step_data['summary'] = 'Agent thinks the request has been completed.'
+      step_data['summary'] = 'Agent thinks the request has been completed.'  # pyrefly: ignore[bad-assignment]
       self.history.append(step_data)
       return base_agent.AgentInteractionResult(
           True,
@@ -525,6 +528,7 @@ def step(self, goal: str) -> base_agent.AgentInteractionResult:
       logging.info('Failed to execute action.')
       logging.info(str(e))
       step_data['summary'] = (
+          # pyrefly: ignore[bad-assignment]
           'Can not execute the action, make sure to select the action with'
           ' the required parameters (if any) in the correct JSON format!'
       )
@@ -556,10 +560,10 @@ def step(self, goal: str) -> base_agent.AgentInteractionResult:
         )
 
     m3a_utils.add_screenshot_label(
-        step_data['before_screenshot_with_som'], 'before'
+        step_data['before_screenshot_with_som'], 'before'  # pyrefly: ignore[bad-argument-type]
     )
     m3a_utils.add_screenshot_label(after_screenshot, 'after')
-    step_data['after_screenshot_with_som'] = after_screenshot.copy()
+    step_data['after_screenshot_with_som'] = after_screenshot.copy()  # pyrefly: ignore[bad-assignment]
 
     summary_prompt = _summarize_prompt(
         action,
@@ -587,7 +591,7 @@ def step(self, goal: str) -> base_agent.AgentInteractionResult:
           summary,
       )
       step_data['summary'] = (
-          'Some error occurred calling LLM during summarization phase: %s'
+          'Some error occurred calling LLM during summarization phase: %s'  # pyrefly: ignore[bad-assignment]
           % summary
       )
       self.history.append(step_data)
@@ -596,8 +600,8 @@ def step(self, goal: str) -> base_agent.AgentInteractionResult:
           step_data,
       )
 
-    step_data['summary_prompt'] = summary_prompt
-    step_data['summary'] = f'Action selected: {action}. {summary}'
+    step_data['summary_prompt'] = summary_prompt  # pyrefly: ignore[bad-assignment]
+    step_data['summary'] = f'Action selected: {action}. {summary}'  # pyrefly: ignore[bad-assignment]
     logging.info('Summary: %s', summary)
     step_data['summary_raw_response'] = raw_response
 

android_world/agents/random_agent.py

@@ -62,8 +62,8 @@ def _generate_random_action(
       json_action.SWIPE,
       json_action.INPUT_TEXT,
   ]:
-    action_details['x'] = random.randint(0, screen_size[0] - 1)
-    action_details['y'] = random.randint(0, screen_size[1] - 1)
+    action_details['x'] = random.randint(0, screen_size[0] - 1)  # pyrefly: ignore[bad-assignment]
+    action_details['y'] = random.randint(0, screen_size[1] - 1)  # pyrefly: ignore[bad-assignment]
     if action_type == json_action.INPUT_TEXT:
       action_details['text'] = ''.join(
           random.choices(text_characters, k=10)

android_world/agents/seeact.py

@@ -200,7 +200,7 @@ def step(
     try:
       action_ground_response = result["action_ground_response"]
       seeact_action = seeact_utils.extract_element_action_value(
-          action_ground_response.split("\n")
+          action_ground_response.split("\n")  # pyrefly: ignore[missing-attribute]
       )
       action = seeact_utils.convert_seeact_action_to_json_action(
           seeact_action, actionable_elements
@@ -217,7 +217,7 @@ def step(
           seeact_action, actionable_elements
       )
       action_description = seeact_utils.generate_action_description(
-          seeact_action, target_element
+          seeact_action, target_element  # pyrefly: ignore[bad-argument-type]
       )
       actuation.execute_adb_action(
           action,

android_world/agents/seeact_utils.py

@@ -430,8 +430,8 @@ def extract_element_action_value(lines: list[str]) -> SeeActAction:
     elif line.startswith("VALUE:"):
       value = line.split(":")[1].strip().strip(".")
 
-  _validate_action(element, action, value)
-  return SeeActAction(action=action, element=element, value=value)
+  _validate_action(element, action, value)  # pyrefly: ignore[bad-argument-type]
+  return SeeActAction(action=action, element=element, value=value)  # pyrefly: ignore[bad-argument-type]
 
 
 @dataclasses.dataclass
@@ -699,7 +699,7 @@ def convert_seeact_action_to_json_action(
   if action_type == json_action.INPUT_TEXT:
     text = action.value
   elif action_type == json_action.SCROLL:
-    direction = _swipe_to_scroll(action.value)
+    direction = _swipe_to_scroll(action.value)  # pyrefly: ignore[bad-argument-type]
   elif action_type == json_action.OPEN_APP:
     app_name = action.value
   elif action_type == json_action.ANSWER:

android_world/agents/t3a.py

@@ -221,9 +221,9 @@ def _action_selection_prompt(
     The text prompt for action selection that will be sent to gpt4v.
   """
   if history:
-    history = '\n'.join(history)
+    history = '\n'.join(history)  # pyrefly: ignore[bad-assignment]
   else:
-    history = 'You just started, no action has been performed yet.'
+    history = 'You just started, no action has been performed yet.'  # pyrefly: ignore[bad-assignment]
 
   extra_guidelines = ''
   if additional_guidelines:
@@ -372,7 +372,7 @@ def step(self, goal: str) -> base_agent.AgentInteractionResult:
 
     try:
       converted_action = json_action.JSONAction(
-          **agent_utils.extract_json(action),
+          **agent_utils.extract_json(action),  # pyrefly: ignore[bad-unpacking]
       )
     except Exception as e:  # pylint: disable=broad-exception-caught
       print('Failed to convert the output to a valid action.')
@@ -389,7 +389,7 @@ def step(self, goal: str) -> base_agent.AgentInteractionResult:
       )
 
     if converted_action.action_type in ['click', 'long-press', 'input-text']:
-      if converted_action.index is not None and converted_action.index >= len(
+      if converted_action.index is not None and converted_action.index >= len(  # pyrefly: ignore[unsupported-operation]
           ui_elements
       ):
         print('Index out of range.')
@@ -403,8 +403,8 @@ def step(self, goal: str) -> base_agent.AgentInteractionResult:
         # Add mark for the target ui element, just used for visualization.
         m3a_utils.add_ui_element_mark(
             step_data['before_screenshot'],
-            ui_elements[converted_action.index],
-            converted_action.index,
+            ui_elements[converted_action.index],  # pyrefly: ignore[bad-index]
+            converted_action.index,  # pyrefly: ignore[bad-argument-type]
             logical_screen_size,
             adb_utils.get_physical_frame_boundary(self.env.controller),
             adb_utils.get_orientation(self.env.controller),
@@ -421,7 +421,7 @@ def step(self, goal: str) -> base_agent.AgentInteractionResult:
       )
 
     if converted_action.action_type == 'answer':
-      print('Agent answered with: ' + converted_action.text)
+      print('Agent answered with: ' + converted_action.text)  # pyrefly: ignore[unsupported-operation]
 
     try:
       self.env.execute_action(converted_action)
@@ -432,7 +432,7 @@ def step(self, goal: str) -> base_agent.AgentInteractionResult:
       )
       print(str(e))
       step_data['summary'] = (
-          'Some error happened executing the action '
+          'Some error happened executing the action '  # pyrefly: ignore[unsupported-operation]
           + converted_action.action_type
       )
       self.history.append(step_data)

android_world/env/actuation.py

@@ -44,12 +44,12 @@ def execute_adb_action(
     x = action.x
     y = action.y
     if idx is not None:
-      if idx < 0 or idx >= len(screen_elements):
+      if idx < 0 or idx >= len(screen_elements):  # pyrefly: ignore[unsupported-operation]
         raise ValueError(
             f'Invalid element index: {idx}, must be between 0 and'
             f' {len(screen_elements)-1}.'
         )
-      element = screen_elements[idx]
+      element = screen_elements[idx]  # pyrefly: ignore[bad-index]
       if element.bbox_pixels is None:
         raise ValueError('Bbox is not present on element.')
       x, y = element.bbox_pixels.center
@@ -119,7 +119,7 @@ def execute_adb_action(
     adb_utils.press_back_button(env)
 
   elif action.action_type == 'press_keyboard':
-    adb_utils.press_keyboard_generic(action.keycode, env)
+    adb_utils.press_keyboard_generic(action.keycode, env)  # pyrefly: ignore[bad-argument-type]
   elif action.action_type == 'drag_and_drop':
     if action.touch_xy is not None and action.lift_xy is not None:
       command = adb_utils.generate_drag_and_drop_command(
@@ -140,10 +140,10 @@ def execute_adb_action(
     screen_width, screen_height = screen_size
     if action.index:
       x_min, y_min, x_max, y_max = (
-          max(screen_elements[action.index].bbox_pixels.x_min, 0),
-          max(screen_elements[action.index].bbox_pixels.y_min, 0),
-          min(screen_elements[action.index].bbox_pixels.x_max, screen_width),
-          min(screen_elements[action.index].bbox_pixels.y_max, screen_height),
+          max(screen_elements[action.index].bbox_pixels.x_min, 0),  # pyrefly: ignore[bad-index]
+          max(screen_elements[action.index].bbox_pixels.y_min, 0),  # pyrefly: ignore[bad-index]
+          min(screen_elements[action.index].bbox_pixels.x_max, screen_width),  # pyrefly: ignore[bad-index]
+          min(screen_elements[action.index].bbox_pixels.y_max, screen_height),  # pyrefly: ignore[bad-index]
       )
     else:
       x_min, y_min, x_max, y_max = (0, 0, screen_width, screen_height)
@@ -218,7 +218,7 @@ def execute_adb_action(
       )
       adb_utils.issue_generic_request(request, env)
   elif action.action_type == 'change_orientation':
-    adb_utils.change_orientation(action.orientation, env)
+    adb_utils.change_orientation(action.orientation, env)  # pyrefly: ignore[bad-argument-type]
   elif action.action_type == json_action.UNKNOWN:
     print('Unknown action type; no action will be executed. Try again...')
   else:

android_world/env/adb_utils.py

@@ -1730,7 +1730,7 @@ def retry(n: int) -> Callable[[Any], Any]:
   """Decorator to retry ADB commands."""
 
   def decorator(func: Callable[..., T]) -> Callable[..., T]:
-    def wrapper(*args: Any, **kwargs: Any) -> T:
+    def wrapper(*args: Any, **kwargs: Any) -> T:  # pyrefly: ignore[bad-return]
       attempts = 0
       while attempts < n:
         try:

android_world/env/interface.py

@@ -296,7 +296,7 @@ def get_state(self, wait_to_stabilize: bool = False) -> State:
 
   def execute_action(self, action: json_action.JSONAction) -> None:
     if action.action_type == json_action.ANSWER:
-      self.interaction_cache = action.text
+      self.interaction_cache = action.text  # pyrefly: ignore[bad-assignment]
       if action.text:
         self.display_message(action.text, header='Agent answered:')
       return