typing

ufo/agent/agent.py

@@ -11,18 +11,20 @@
 from ..automator import puppeteer
 from ..automator.ui_control import openfile
 from ..automator.ui_control import utils as control
+from ..config.config import Config
 from ..prompter.agent_prompter import (AppAgentPrompter, FollowerAgentPrompter,
                                        HostAgentPrompter)
 from .basic import BasicAgent, Memory, MemoryItem
 
+configs = Config.get_instance().config_data
 
 class AgentFactory:  
     """  
     Factory class to create agents.  
     """  
 
     @staticmethod  
-    def create_agent(agent_type: str, *args, **kwargs):  
+    def create_agent(agent_type: str, *args, **kwargs) -> BasicAgent:  
         """  
         Create an agent based on the given type.  
         :param agent_type: The type of agent to create.  
@@ -361,16 +363,18 @@ def message_constructor(self, image_list: List, request_history: str, action_his
 
         return hostagent_prompt_message
 
+
     def app_file_manager(self, app_file_info: dict):
         '''
         Open the application or file for the user.
         :param app_file_info: The information of the application or file. {'APP': name of app, 'file_path': path}
         :return: The window of the application.
         '''
+
         utils.print_with_color("Opening the required application or file...", "yellow")
         file_manager = openfile.FileController()
         results = file_manager.execute_code(app_file_info)
-        time.sleep(5)
+        time.sleep(configs.get("SLEEP_TIME", 5))
         desktop_windows_dict, _ = control.get_desktop_app_info_dict()
         if not results:
             self.status = "ERROR in openning the application or file."

ufo/agent/basic.py

@@ -21,7 +21,7 @@ class MemoryItem:
 
     _memory_attributes = []
 
-    def to_dict(self) -> dict:
+    def to_dict(self) -> Dict[str, str]:
         """
         Convert the MemoryItem to a dictionary.
         :return: The dictionary.
@@ -59,7 +59,7 @@ def set_value(self, key: str, value: str) -> None:
             self._memory_attributes.append(key)
 
 
-    def set_values_from_dict(self, values: dict) -> None:
+    def set_values_from_dict(self, values: Dict[str, str]) -> None:
         """
         Add fields to the memory item.
         :param values: The values of the fields.
@@ -105,15 +105,15 @@ class Memory():
     _content: List[MemoryItem] = field(default_factory=list)
 
 
-    def load(self, content: List[MemoryItem]) -> dict:
+    def load(self, content: List[MemoryItem]) -> None:
         """
         Load the data from the memory.
         :param key: The key of the data.
         """
         self._content = content
 
 
-    def filter_memory_from_steps(self, steps: List[int]) -> List[dict]:
+    def filter_memory_from_steps(self, steps: List[int]) -> List[Dict[str, str]]:
         """
         Filter the memory from the steps.
         :param steps: The steps to filter.
@@ -122,7 +122,7 @@ def filter_memory_from_steps(self, steps: List[int]) -> List[dict]:
         return [item.to_dict() for item in self._content if item.step in steps]
 
 
-    def filter_memory_from_keys(self, keys: List[str]) -> List[dict]:
+    def filter_memory_from_keys(self, keys: List[str]) -> List[Dict[str, str]]:
         """
         Filter the memory from the keys. If an item does not have the key, the key will be ignored.
         :param keys: The keys to filter.
@@ -199,7 +199,7 @@ class BasicAgent(ABC):
     The BasicAgent class is the abstract class for the agent.
     """
 
-    def __init__(self, name: str):
+    def __init__(self, name: str) -> None:
         """
         Initialize the BasicAgent.
         :param name: The name of the agent.
@@ -281,7 +281,7 @@ def get_response(cls, message: List[dict], namescope, use_backup_engine) -> str:
 
 
     @staticmethod
-    def response_to_dict(response: str) -> dict:
+    def response_to_dict(response: str) -> Dict[str, str]:
         """
         Convert the response to a dictionary.
         :param response: The response.

ufo/module/processors/follower_processor.py

@@ -14,7 +14,7 @@ class FollowerHostAgentProcessor(HostAgentProcessor):
 
     def create_sub_agent(self) -> FollowerAgent:
         """
-        Create a sub agent for the host agent.
+        Create a follower subagent for the host agent.
         :return: The created sub agent.
         """
 
@@ -42,7 +42,7 @@ class FollowerAppAgentProcessor(AppAgentProcessor):
 
     def get_prompt_message(self) -> None:
             """
-            Get the prompt message for the AppAgent.
+            Get the prompt message for the AppAgent in the follower mode. It may accept additional prompts as input.
             """
 
             if configs["RAG_EXPERIENCE"]:

ufo/module/processors/processor.py

@@ -28,7 +28,8 @@
 class HostAgentProcessor(BaseProcessor):
 
     def __init__(self, round_num: int, log_path: str, photographer: PhotographerFacade, request: str, request_logger: Logger, logger: Logger, 
-                 host_agent: HostAgent, round_step: int, global_step: int, prev_status: str, app_window=None):
+                 host_agent: HostAgent, round_step: int, global_step: int, prev_status: str, app_window=None) -> None:
+
         super().__init__(round_num, log_path, photographer, request, request_logger, logger, round_step, global_step, prev_status, app_window)
 
         """
@@ -263,7 +264,8 @@ def app_agent_context_provision(self, app_agent: AppAgent) -> None:
 class AppAgentProcessor(BaseProcessor):
 
         def __init__(self, round_num: int, log_path: str, photographer: PhotographerFacade, request: str, request_logger: Logger, logger: Logger, app_agent: AppAgent, round_step:int, global_step: int, 
-                     process_name: str, app_window: UIAWrapper, control_reannotate: Optional[list], prev_status: str):
+                     process_name: str, app_window: UIAWrapper, control_reannotate: Optional[list], prev_status: str) -> None:
+
             super().__init__(round_num, log_path, photographer, request, request_logger, logger, round_step, global_step, prev_status, app_window)
 
             """

ufo/module/round.py

@@ -5,12 +5,13 @@
 from logging import Logger
 from typing import Optional
 
-from ..agent.agent import HostAgent, FollowerAgent
+from pywinauto.controls.uiawrapper import UIAWrapper
+
+from ..agent.agent import FollowerAgent, HostAgent
 from ..automator.ui_control.screenshot import PhotographerFacade
 from ..config.config import Config
-from .processors import processor, follower_processor   
-
 from .basic import BaseRound
+from .processors import follower_processor, processor
 
 configs = Config.get_instance().config_data
 

ufo/module/session.py

@@ -3,7 +3,7 @@
 
 import json
 import os
-from typing import List
+from typing import List, Optional
 
 from .. import utils
 from ..automator.ui_control.screenshot import PhotographerFacade
@@ -38,7 +38,7 @@ def get_task(self) -> str:
         return self.plan.get("task", "")
 
 
-    def get_steps(self) -> list:
+    def get_steps(self) -> List[str]:
         """
         Get the steps in the plan.
         :return: The steps in the plan.
@@ -83,7 +83,7 @@ def get_host_agent_request(self) -> str:
         return request
 
 
-    def next_step(self) -> dict:
+    def next_step(self) -> Optional[str]:
         """
         Get the next step in the plan.
         :return: The next step.
@@ -152,7 +152,7 @@ def is_folder(path: str) -> bool:
 
 
     @staticmethod
-    def get_plan_files(path: str) -> list:
+    def get_plan_files(path: str) -> List[str]:
         """
         Get the plan files in the folder. The plan file should have the extension ".json".
         :param path: The path of the folder.
@@ -265,7 +265,7 @@ class FollowerSession(Session):
     A session for following a list of plan for action taken.
     """
 
-    def __init__(self, task: str, plan_file: str):
+    def __init__(self, task: str, plan_file: str) -> None:
         """
         Initialize a session.
         :param task: The name of current task.
@@ -285,7 +285,7 @@ def __init__(self, task: str, plan_file: str):
 
 
 
-    def create_round(self) -> round.Round:
+    def create_round(self) -> round.FollowerRound:
         """
         Create a new round.
         """

ufo/module/state.py

@@ -11,7 +11,7 @@
 from __future__ import annotations
 
 from abc import ABC
-from typing import Dict, Type
+from typing import TYPE_CHECKING, Dict, Type
 
 from ..config.config import Config
 from ..utils import print_with_color
@@ -20,6 +20,12 @@
 configs = Config.get_instance().config_data
 
 
+# To avoid circular import
+if TYPE_CHECKING:
+    from .session import Session
+
+
+
 class StatusToStateMapper(ABC):  
     """  
     A class to map the status to the appropriate state.  
@@ -62,7 +68,7 @@ def __init__(self):
         """
         self.state_mapping = StatusToStateMapper()
 
-    def handle(self, session):
+    def handle(self, session: "Session") -> None:
         """
         Handle the session.
         :param session: The session.
@@ -83,7 +89,7 @@ class NoneState(SessionState):
     The state when the session is None.
     """
 
-    def handle(self, session):
+    def handle(self, session: "Session") -> None:
         """
         Handle the session. Do nothing.
         :param session: The session.
@@ -96,7 +102,7 @@ class RoundFinishState(SessionState):
     The state when a single round is finished.
     """
 
-    def handle(self, session):
+    def handle(self, session: "Session") -> None:
         """
         Handle the session. Either start a new round or finish the session.
         :param session: The session.
@@ -128,7 +134,7 @@ class SessionFinishState(SessionState):
     The state when the entire session is finished.
     """
 
-    def handle(self, session):
+    def handle(self, session: "Session") -> None:
         """
         Handle the session. Finish the entire session, and save the experience if needed.
         :param session: The session.
@@ -146,7 +152,7 @@ class ErrorState(SessionState):
     The state when an error occurs.
     """
 
-    def handle(self, session):
+    def handle(self, session: "Session") -> None:
         """
         Handle the session. Do nothing.
         :param session: The session.
@@ -159,7 +165,7 @@ class AppSelectionState(SessionState):
     The state when the application selection is needed by a HostAgent.
     """
 
-    def handle(self, session):
+    def handle(self, session: "Session") -> None:
         """
         Handle the session. Process the application selection.
         :param session: The session.
@@ -184,7 +190,7 @@ class ContinueState(SessionState):
     The state when the session needs to continue by the AppAgent.
     """
 
-    def handle(self, session):
+    def handle(self, session: "Session") -> None:
         """
         Handle the session. Process the action selection.
         :param session: The session.
@@ -210,7 +216,7 @@ class AnnotationState(ContinueState):
     The state when the session needs to re-nnotate the screenshot.
     """
 
-    def handle(self, session):
+    def handle(self, session: "Session") -> None:
         """
         Handle the session. Process the action selection with the re-annotation screenshot. Same as ContinueState.
         :param session: The session.
@@ -225,7 +231,7 @@ class MaxStepReachedState(SessionState):
     The state when the maximum step is reached.
     """
 
-    def handle(self, session):
+    def handle(self, session: "Session") -> None:
         """
         Handle the session. Finish the session when the maximum step is reached.
         :param session: The session.