change var name and clean code

ufo/automator/ui_control/controller.py

@@ -7,7 +7,6 @@
 from typing import Dict, List
 
 from ...config.config import Config
-from .utils import get_control_info
 from ...utils import print_with_color
 from ..basic import CommandBasic, ReceiverBasic, ReceiverFactory
 
@@ -27,7 +26,6 @@ def __init__(self, control, application):
         self.control = control
 
         if control:
-            self.control_info = get_control_info(control)
             self.control.set_focus()
             self.wait_enabled()
         self.application = application

ufo/module/basic.py

@@ -25,7 +25,8 @@
 from ..agent.agent import AgentFactory, HostAgent
 from ..config.config import Config
 from ..experience.summarizer import ExperienceSummarizer
-from .state import StatusToStateMapper
+from .state import (ErrorState, MaxStepReachedState, NoneState,
+                    SessionFinishState, StatusToStateMapper)
 
 configs = Config.get_instance().config_data
 
@@ -74,7 +75,7 @@ def __init__(self, task: str, logger: Logger, request_logger: Logger, host_agent
 
         # round_num and global step-related properties  
         self.round_num = None  
-        self.global_step = None
+        self.session_step = None
 
 
     def process_application_selection(self) -> None:
@@ -96,25 +97,25 @@ def process_action_selection(self) -> None:
 
     def get_status(self) -> str:
         """
-        Get the status of the session.
-        return: The status of the session.
+        Get the status of the round.
+        return: The status of the round.
         """
         return self._status
 
 
 
     def get_step(self) -> int:
         """
-        Get the step of the session.
-        return: The step of the session.
+        Get the local step of the round.
+        return: The step of the round.
         """
         return self._step
 
 
     def get_cost(self) -> float:
         """
-        Get the cost of the session.
-        return: The cost of the session.
+        Get the cost of the round.
+        return: The cost of the round.
         """
         return self._cost
 
@@ -132,8 +133,8 @@ def print_cost(self) -> None:
 
     def get_results(self) -> str:
         """
-        Get the results of the session.
-        return: The results of the session.
+        Get the results of the round.
+        return: The results of the round.
         """
 
         action_history = self.host_agent.get_global_action_memory().content
@@ -147,16 +148,16 @@ def get_results(self) -> str:
 
     def set_index(self, index: int) -> None:
         """
-        Set the round index of the session.
+        Set the round index of the round.
         """
         self.round_num = index
 
 
-    def set_global_step(self, global_step: int) -> None:
+    def set_session_step(self, session_step: int) -> None:
         """
-        Set the global step of the session.
+        Set the global step of the round.
         """
-        self.global_step = global_step
+        self.session_step = session_step
 
 
     def get_application_window(self) -> UIAWrapper:
@@ -177,7 +178,7 @@ def set_application_window(self, app_window: UIAWrapper) -> None:
 
     def update_cost(self, cost: float) -> None:
         """
-        Update the cost of the session.
+        Update the cost of the round.
         """
         if isinstance(cost, float) and isinstance(self._cost, float):
             self._cost += cost
@@ -291,7 +292,7 @@ def round_appagent_execution(self) -> None:
         pass
 
 
-    def get_current_round(self):
+    def get_current_round(self) -> BaseRound:
         """
         Get the current round.
         return: The current round.
@@ -393,6 +394,16 @@ def update_cost(self, cost: float) -> None:
 
 
 
+    def is_finish(self) -> bool:
+        """
+        Check if the session is ended.
+        return: True if the session is ended, otherwise False.
+        """
+
+        # Finish the session if the state is SessionFinishState, ErrorState, MaxStepReachedState, or NoneState.
+        return True if isinstance(self.get_state(), (SessionFinishState, ErrorState, MaxStepReachedState, NoneState)) else False
+
+
     def set_state(self, state) -> None:
         """
         Set the state of the session.

ufo/module/client.py

@@ -23,10 +23,10 @@ def run(self) -> None:
         Run the UFO client.
         """
 
-        while not isinstance(self.session.get_state(), (SessionFinishState, ErrorState, MaxStepReachedState, NoneState)):
-
+        while not self.session.is_finish():
             self.session.handle()
 
+        # If the session is finished normally, handle the last state with additional logic.
         if isinstance(self.session.get_state(), SessionFinishState):
             self.session.handle()
 

ufo/module/processors/basic.py

@@ -24,7 +24,7 @@ class BaseProcessor(ABC):
     """
 
     def __init__(self, round_num: int, log_path: str, request: str, request_logger: Logger, logger: Logger, 
-                 round_step: int, global_step: int, prev_status: str, app_window:UIAWrapper) -> None:
+                 round_step: int, session_step: int, prev_status: str, app_window:UIAWrapper) -> None:
         """
         Initialize the processor.
         :param round_num: The index of the processor. The round_num is the total number of rounds in the session.
@@ -33,7 +33,7 @@ def __init__(self, round_num: int, log_path: str, request: str, request_logger:
         :param request_logger: The logger for the request string.
         :param logger: The logger for the response and error.
         :param round_step: The step of the round.
-        :param global_step: The global step of the session.
+        :param session_step: The global step of the session.
         :param prev_status: The previous status of the session.
         :param app_window: The application window.
         """
@@ -46,7 +46,7 @@ def __init__(self, round_num: int, log_path: str, request: str, request_logger:
         self.logger = logger
         self._app_window = app_window
 
-        self.global_step = global_step
+        self.session_step = session_step
         self.round_step = round_step
         self.prev_status = prev_status
         self.round_num = round_num

ufo/module/processors/follower_processor.py

@@ -77,6 +77,6 @@ def get_prompt_message(self) -> None:
             self._prompt_message = self.app_agent.message_constructor(examples, tips, external_knowledge_prompt, self._image_url, request_history, action_history, 
                                                                                 self.filtered_control_info, self.prev_plan, self.request, current_state, state_diff, configs["INCLUDE_LAST_SCREENSHOT"])
 
-            log = json.dumps({"step": self.global_step, "prompt": self._prompt_message, "control_items": self._control_info, 
+            log = json.dumps({"step": self.session_step, "prompt": self._prompt_message, "control_items": self._control_info, 
                               "filted_control_items": self.filtered_control_info, "status": ""})
             self.request_logger.debug(log)

ufo/module/processors/processor.py

@@ -26,9 +26,9 @@
 class HostAgentProcessor(BaseProcessor):
 
     def __init__(self, round_num: int, log_path: str, request: str, request_logger: Logger, logger: Logger, 
-                 host_agent: HostAgent, round_step: int, global_step: int, prev_status: str, app_window=None) -> None:
+                 host_agent: HostAgent, round_step: int, session_step: int, prev_status: str, app_window=None) -> None:
 
-        super().__init__(round_num, log_path, request, request_logger, logger, round_step, global_step, prev_status, app_window)
+        super().__init__(round_num, log_path, request, request_logger, logger, round_step, session_step, prev_status, app_window)
 
         """
         Initialize the host agent processor.
@@ -39,7 +39,7 @@ def __init__(self, round_num: int, log_path: str, request: str, request_logger:
         :param logger: The logger for the response and error.
         :param host_agent: The host agent to process the session.
         :param round_step: The number of steps in the round.
-        :param global_step: The global step of the session.
+        :param session_step: The global step of the session.
         :param prev_status: The previous status of the session.
         :param app_window: The application window.
         """
@@ -63,7 +63,7 @@ def capture_screenshot(self) -> None:
         Capture the screenshot.
         """
 
-        desktop_save_path = self.log_path + f"action_step{self.global_step}.png"
+        desktop_save_path = self.log_path + f"action_step{self.session_step}.png"
 
         # Capture the desktop screenshot for all screens.
         self.photographer.capture_desktop_screen_screenshot(all_screens=True, save_path=desktop_save_path)
@@ -202,7 +202,7 @@ def update_memory(self) -> None:
         host_agent_step_memory = MemoryItem()
 
         # Log additional information for the host agent.
-        additional_memory = {"Step": self.global_step, "RoundStep": self.get_process_step(), "AgentStep": self.host_agent.get_step(), "Round": self.round_num, "ControlLabel": self._control_text, "Action": "set_focus()", 
+        additional_memory = {"Step": self.session_step, "RoundStep": self.get_process_step(), "AgentStep": self.host_agent.get_step(), "Round": self.round_num, "ControlLabel": self._control_text, "Action": "set_focus()", 
                                 "ActionType": "UIControl", "Request": self.request, "Agent": "HostAgent", "AgentName": self.host_agent.name, "Application": self.app_root, "Cost": self._cost, "Results": ""}
 
         host_agent_step_memory.set_values_from_dict(self._response_json)
@@ -297,10 +297,10 @@ def app_agent_context_provision(self, app_agent: AppAgent) -> None:
 
 class AppAgentProcessor(BaseProcessor):
 
-        def __init__(self, round_num: int, log_path: str, request: str, request_logger: Logger, logger: Logger, app_agent: AppAgent, round_step:int, global_step: int, 
+        def __init__(self, round_num: int, log_path: str, request: str, request_logger: Logger, logger: Logger, app_agent: AppAgent, round_step:int, session_step: int, 
                      process_name: str, app_window: UIAWrapper, control_reannotate: Optional[list], prev_status: str) -> None:
 
-            super().__init__(round_num, log_path, request, request_logger, logger, round_step, global_step, prev_status, app_window)
+            super().__init__(round_num, log_path, request, request_logger, logger, round_step, session_step, prev_status, app_window)
 
             """
             Initialize the app agent processor.
@@ -311,7 +311,7 @@ def __init__(self, round_num: int, log_path: str, request: str, request_logger:
             :param logger: The logger for the response and error.
             :param app_agent: The app agent to process the current step.
             :param round_step: The number of steps in the round.
-            :param global_step: The global step of the session.
+            :param session_step: The global step of the session.
             :param process_name: The process name.
             :param app_window: The application window.
             :param control_reannotate: The list of controls to reannotate.
@@ -345,9 +345,9 @@ def capture_screenshot(self) -> None:
             """
 
             # Define the paths for the screenshots saved.
-            screenshot_save_path = self.log_path + f"action_step{self.global_step}.png"
-            annotated_screenshot_save_path = self.log_path + f"action_step{self.global_step}_annotated.png"
-            concat_screenshot_save_path = self.log_path + f"action_step{self.global_step}_concat.png"
+            screenshot_save_path = self.log_path + f"action_step{self.session_step}.png"
+            annotated_screenshot_save_path = self.log_path + f"action_step{self.session_step}_annotated.png"
+            concat_screenshot_save_path = self.log_path + f"action_step{self.session_step}_concat.png"
 
             # Get the control elements in the application window if the control items are not provided for reannotation.
             if type(self._control_reannotate) == list and len(self._control_reannotate) > 0:
@@ -369,8 +369,8 @@ def capture_screenshot(self) -> None:
 
             # If the configuration is set to include the last screenshot with selected controls tagged, save the last screenshot.
             if configs["INCLUDE_LAST_SCREENSHOT"]:
-                last_screenshot_save_path = self.log_path + f"action_step{self.global_step - 1}.png"
-                last_control_screenshot_save_path = self.log_path + f"action_step{self.global_step - 1}_selected_controls.png"
+                last_screenshot_save_path = self.log_path + f"action_step{self.session_step - 1}.png"
+                last_control_screenshot_save_path = self.log_path + f"action_step{self.session_step - 1}_selected_controls.png"
                 self._image_url += [self.photographer.encode_image_from_path(last_control_screenshot_save_path if os.path.exists(last_control_screenshot_save_path) else last_screenshot_save_path)]
 
             # Whether to concatenate the screenshots of clean screenshot and annotated screenshot into one image.
@@ -431,7 +431,7 @@ def get_prompt_message(self) -> None:
 
 
             # Log the prompt message. Only save them in debug mode.
-            log = json.dumps({"step": self.global_step, "prompt": self._prompt_message, "control_items": self._control_info, 
+            log = json.dumps({"step": self.session_step, "prompt": self._prompt_message, "control_items": self._control_info, 
                               "filted_control_items": self.filtered_control_info, "status": ""})
             self.request_logger.debug(log)
 
@@ -446,7 +446,7 @@ def get_response(self) -> None:
                 self._response, self._cost = self.app_agent.get_response(self._prompt_message, "APPAGENT", use_backup_engine=True)
             except Exception as e:
                 error_trace = traceback.format_exc()
-                log = json.dumps({"step": self.global_step, "prompt": self._prompt_message, "status": str(error_trace)})
+                log = json.dumps({"step": self.session_step, "prompt": self._prompt_message, "status": str(error_trace)})
                 utils.print_with_color("Error occurs when calling LLM: {e}".format(e=str(error_trace)), "red")
                 self.request_logger.info(log)
                 self._status = "ERROR"
@@ -492,7 +492,7 @@ def execute_action(self) -> None:
                 self.app_agent.Puppeteer.receiver_manager.create_ui_control_receiver(control_selected, self._app_window)
 
                 # Save the screenshot of the tagged selected control.
-                control_screenshot_save_path = self.log_path + f"action_step{self.global_step}_selected_controls.png"
+                control_screenshot_save_path = self.log_path + f"action_step{self.session_step}_selected_controls.png"
                 self.photographer.capture_app_window_screenshot_with_rectangle(self._app_window, sub_control_list=[control_selected], save_path=control_screenshot_save_path)
 
                 # Compose the function call and the arguments string.
@@ -528,7 +528,7 @@ def execute_action(self) -> None:
 
             except Exception:
                 error_trace = traceback.format_exc()
-                utils.print_with_color(f"Error Occurs at action execution in AppAgent at step {self.global_step}", "red")
+                utils.print_with_color(f"Error Occurs at action execution in AppAgent at step {self.session_step}", "red")
                 utils.print_with_color(str(error_trace), "red")
                 utils.print_with_color(self._response, "red")
                 self.error_log(self._response, str(error_trace))
@@ -546,7 +546,7 @@ def update_memory(self) -> None:
             host_agent = self.app_agent.get_host()
 
             # Log additional information for the app agent.
-            additional_memory = {"Step": self.global_step, "RoundStep": self.get_process_step(), "AgentStep": self.app_agent.get_step(), "Round": self.round_num, "Action": self._action, 
+            additional_memory = {"Step": self.session_step, "RoundStep": self.get_process_step(), "AgentStep": self.app_agent.get_step(), "Round": self.round_num, "Action": self._action, 
                                  "ActionType": self.app_agent.Puppeteer.get_command_types(self._operation), "Request": self.request, "Agent": "ActAgent", "AgentName": self.app_agent.name, "Application": app_root, "Cost": self._cost, "Results": self._results}
             app_agent_step_memory.set_values_from_dict(self._response_json)
             app_agent_step_memory.set_values_from_dict(additional_memory)

ufo/module/round.py

@@ -40,7 +40,7 @@ def process_application_selection(self) -> None:
         Select an application to interact with.
         """
 
-        host_agent_processor = processor.HostAgentProcessor(round_num=self.round_num, log_path=self.log_path, request=self.request, round_step=self.get_step(), global_step=self.global_step,
+        host_agent_processor = processor.HostAgentProcessor(round_num=self.round_num, log_path=self.log_path, request=self.request, round_step=self.get_step(), session_step=self.session_step,
                                                             request_logger=self.request_logger, logger=self.logger, host_agent=self.host_agent, prev_status=self.get_status(), app_window=self.app_window)
 
         host_agent_processor.process()
@@ -60,7 +60,7 @@ def process_action_selection(self) -> None:
         Select an action with the application.
         """
 
-        app_agent_processor = processor.AppAgentProcessor(round_num=self.round_num, log_path=self.log_path, request=self.request, round_step=self.get_step(), global_step=self.global_step, 
+        app_agent_processor = processor.AppAgentProcessor(round_num=self.round_num, log_path=self.log_path, request=self.request, round_step=self.get_step(), session_step=self.session_step, 
                                                           process_name=self.application, request_logger=self.request_logger, logger=self.logger, app_agent=self.app_agent, app_window=self.app_window, 
                                                             control_reannotate=self.control_reannotate, prev_status=self.get_status())
 
@@ -105,7 +105,7 @@ def process_application_selection(self) -> None:
         Select an application to interact with.
         """
 
-        host_agent_processor = follower_processor.FollowerHostAgentProcessor(round_num=self.round_num, log_path=self.log_path, request=self.request, round_step=self.get_step(), global_step=self.global_step,
+        host_agent_processor = follower_processor.FollowerHostAgentProcessor(round_num=self.round_num, log_path=self.log_path, request=self.request, round_step=self.get_step(), session_step=self.session_step,
                                                             request_logger=self.request_logger, logger=self.logger, host_agent=self.host_agent, prev_status=self.get_status(), app_window=self.app_window)
 
         host_agent_processor.process()
@@ -125,7 +125,7 @@ def process_action_selection(self) -> None:
         Select an action with the application.
         """
 
-        app_agent_processor = follower_processor.FollowerAppAgentProcessor(round_num=self.round_num, log_path=self.log_path, request=self.request, round_step=self.get_step(), global_step=self.global_step, 
+        app_agent_processor = follower_processor.FollowerAppAgentProcessor(round_num=self.round_num, log_path=self.log_path, request=self.request, round_step=self.get_step(), session_step=self.session_step, 
                                                           process_name=self.application, request_logger=self.request_logger, logger=self.logger, app_agent=self.app_agent, app_window=self.app_window, 
                                                             control_reannotate=self.control_reannotate, prev_status=self.get_status())