lybic-guiagents 0.1.0__py3-none-any.whl

gui_agents/agents/agent_s.py

@@ -0,0 +1,832 @@
+import json
+import logging
+from math import log
+import os
+import platform
+import textwrap
+from typing import Dict, List, Optional, Tuple
+
+from gui_agents.agents.grounding import ACI
+from gui_agents.agents.worker import Worker
+from gui_agents.agents.manager import Manager
+from gui_agents.agents.grounding import Grounding, FastGrounding
+from gui_agents.utils.common_utils import Node
+from gui_agents.agents.global_state import GlobalState
+from gui_agents.store.registry import Registry
+from gui_agents.utils.common_utils import (
+    # call_llm_safe,
+    parse_single_code_from_string,
+    sanitize_code,
+    extract_first_agent_function,
+    agent_log_to_string,
+)
+from gui_agents.tools.tools import Tools
+
+logger = logging.getLogger("desktopenv.agent")
+
+class UIAgent:
+    """Base class for UI automation agents"""
+
+    def __init__(
+        self,
+        platform: str = platform.system().lower(),
+    ):
+        """Initialize UIAgent
+
+        Args:
+            platform: Operating system platform (macos, linux, windows)
+        """
+        self.platform = platform
+
+    def reset(self) -> None:
+        """Reset agent state"""
+        pass
+
+    def predict(self, instruction: str, observation: Dict) -> Tuple[Dict, List[str]]|None:
+        """Generate next action prediction
+
+        Args:
+            instruction: Natural language instruction
+            observation: Current UI state observation
+
+        Returns:
+            Tuple containing agent info dictionary and list of actions
+        """
+        pass
+
+    def update_narrative_memory(self, trajectory: str) -> None:
+        """Update narrative memory with task trajectory
+
+        Args:
+            trajectory: String containing task execution trajectory
+        """
+        pass
+
+    def update_episodic_memory(self, meta_data: Dict, subtask_trajectory: str) -> str|None:
+        """Update episodic memory with subtask trajectory
+
+        Args:
+            meta_data: Metadata about current subtask execution
+            subtask_trajectory: String containing subtask execution trajectory
+
+        Returns:
+            Updated subtask trajectory
+        """
+        pass
+
+class AgentS2(UIAgent):
+    """Agent that uses hierarchical planning and directed acyclic graph modeling for UI automation"""
+
+    def __init__(
+        self,
+        platform: str = platform.system().lower(),
+        screen_size: List[int] = [1920, 1080],
+        memory_root_path: str = os.getcwd(),
+        memory_folder_name: str = "kb_s2",
+        kb_release_tag: str = "v0.2.2",
+        enable_takeover: bool = False,
+        enable_search: bool = True,
+    ):
+        """Initialize AgentS2
+
+        Args:
+            platform: Operating system platform (darwin, linux, windows)
+            memory_root_path: Path to memory directory. Defaults to current working directory.
+            memory_folder_name: Name of memory folder. Defaults to "kb_s2".
+            kb_release_tag: Release tag for knowledge base. Defaults to "v0.2.2".
+            enable_takeover: Whether to enable user takeover functionality. Defaults to False.
+            enable_search: Whether to enable web search functionality. Defaults to True.
+        """
+        super().__init__(
+            platform,
+        )
+
+        self.memory_root_path = memory_root_path
+        self.memory_folder_name = memory_folder_name
+        self.kb_release_tag = kb_release_tag
+        self.screen_size = screen_size
+        self.enable_takeover = enable_takeover
+        self.enable_search = enable_search
+
+        # Load tools configuration from tools_config.json
+        tools_config_path = os.path.join(os.path.dirname(os.path.dirname(__file__)), "tools", "tools_config.json")
+        with open(tools_config_path, "r") as f:
+            self.tools_config = json.load(f)
+            print(f"Loaded tools configuration from: {tools_config_path}")
+            self.Tools_dict = {}
+            for tool in self.tools_config["tools"]:
+                tool_name = tool["tool_name"]
+                self.Tools_dict[tool_name] = {
+                    "provider": tool["provider"],
+                    "model": tool["model_name"]
+                }
+            print(f"Tools configuration: {self.Tools_dict}")
+
+        # Initialize agent's knowledge base path
+        self.local_kb_path = os.path.join(
+            self.memory_root_path, self.memory_folder_name
+        )
+
+        # Check if knowledge base exists
+        kb_platform_path = os.path.join(self.local_kb_path, self.platform)
+        if not os.path.exists(kb_platform_path):
+            print(f"Warning: Knowledge base for {self.platform} platform not found in {self.local_kb_path}")
+            os.makedirs(kb_platform_path, exist_ok=True)
+            print(f"Created directory: {kb_platform_path}")
+            # raise FileNotFoundError(f"Knowledge base path does not exist: {kb_platform_path}")
+        else:
+            print(f"Found local knowledge base path: {kb_platform_path}")
+
+        self.reset()
+
+    def reset(self) -> None:
+        """Reset agent state and initialize components"""
+        # Initialize core components
+        
+        self.manager = Manager(
+            Tools_dict=self.Tools_dict,
+            local_kb_path=self.local_kb_path,
+            platform=self.platform,
+            enable_search=self.enable_search,  # Pass global switch to Manager
+        )
+        
+        self.worker = Worker(
+            Tools_dict=self.Tools_dict,
+            local_kb_path=self.local_kb_path,
+            platform=self.platform,
+            enable_takeover=self.enable_takeover,
+            enable_search=self.enable_search,  # Pass global switch to Worker
+            tools_config=self.tools_config,    # Pass complete tools configuration
+        )
+
+        self.grounding = Grounding(
+            Tools_dict=self.Tools_dict,
+            platform=self.platform,
+            width=self.screen_size[0],
+            height=self.screen_size[1]
+        )
+
+        # Reset state variables
+        self.requires_replan: bool = True
+        self.needs_next_subtask: bool = True
+        self.step_count: int = 0
+        self.turn_count: int = 0
+        self.failure_subtask: Optional[Node] = None
+        self.should_send_action: bool = False
+        self.completed_tasks: List[Node] = []
+        self.current_subtask: Optional[Node] = None
+        self.subtasks: List[Node] = []
+        self.search_query: str = ""
+        self.subtask_status: str = "Start"
+        self.global_state: GlobalState = Registry.get("GlobalStateStore") # type: ignore
+
+    def reset_executor_state(self) -> None:
+        """Reset executor and step counter"""
+        self.worker.reset()
+        self.step_count = 0
+
+    def predict(self, instruction: str, observation: Dict) -> Tuple[Dict, List[str]]:
+        # Initialize the three info dictionaries
+        planner_info = {}
+        executor_info = {}
+        evaluator_info = {
+            "obs_evaluator_response": "",
+            "num_input_tokens_evaluator": 0,
+            "num_output_tokens_evaluator": 0,
+            "evaluator_cost": 0.0,
+        }
+        actions = []
+
+        # 记录预测开始时间
+        import time
+        predict_start_time = time.time()
+
+        # If the DONE response by the executor is for a subtask, then the agent should continue with the next subtask without sending the action to the environment
+        while not self.should_send_action:
+            time.sleep(5.0)
+            self.subtask_status = "In"
+            # Always time get_action_queue, even if not called
+            import time
+            manager_start = time.time()
+            # If replan is true, generate a new plan. True at start, after a failed plan, or after subtask completion
+            if self.requires_replan:
+                logger.info("(RE)PLANNING...")
+                Manager_info, self.subtasks = self.manager.get_action_queue(
+                    Tu=self.global_state.get_Tu(),
+                    observation=self.global_state.get_obs_for_manager(),
+                    running_state=self.global_state.get_running_state(),
+                    failed_subtask=self.failure_subtask,
+                    completed_subtasks_list=self.global_state.get_completed_subtasks(),
+                    remaining_subtasks_list=self.global_state.get_remaining_subtasks(),
+                )
+                self.global_state.set_remaining_subtasks(self.subtasks) # type: ignore
+
+                self.requires_replan = False
+                if "search_query" in Manager_info:
+                    self.search_query = Manager_info["search_query"]
+                else:
+                    self.search_query = ""
+            get_action_queue_time = time.time() - manager_start
+            logger.info(f"[Timing] manager.get_action_queue execution time: {get_action_queue_time:.2f} seconds")
+            self.global_state.log_operation(
+                module="manager",
+                operation="manager.get_action_queue",
+                data={"duration": get_action_queue_time}
+            )
+
+            # use the exectuor to complete the topmost subtask
+            if self.needs_next_subtask:
+                logger.info("GETTING NEXT SUBTASK...")
+
+                # this can be empty if the DAG planner deems that all subtasks are completed
+                if len(self.subtasks) <= 0:
+                    self.requires_replan = True
+                    self.needs_next_subtask = True
+                    self.failure_subtask = None
+                    if self.current_subtask is not None:
+                        self.global_state.add_completed_subtask(self.current_subtask)
+                    # reset executor state
+                    self.reset_executor_state()
+                    self.should_send_action = True
+                    self.subtask_status = "Done"
+                    executor_info = {
+                        "executor_plan": "agent.done()",
+                        "plan_code": "agent.done()",
+                        "reflection": "agent.done()",
+                    }
+                    actions = [{"type": "DONE"}]
+                    
+                    # 记录任务完成
+                    self.global_state.log_operation(
+                        module="agent",
+                        operation="task_complete",
+                        data={
+                            "content": "All subtasks completed, task finished",
+                            "status": "done"
+                        }
+                    )
+                    break
+
+                self.current_subtask = self.subtasks.pop(0)
+                self.global_state.set_remaining_subtasks(self.subtasks)
+                logger.info(f"NEXT SUBTASK: {self.current_subtask}")
+                logger.info(f"REMAINING SUBTASKS: {self.subtasks}")
+                logger.info(f"REMAINING SUBTASKS FROM GLOBAL STATE: {self.global_state.get_remaining_subtasks()}")
+                self.needs_next_subtask = False
+                self.subtask_status = "Start"
+                
+                self.global_state.log_operation(
+                    module="agent",
+                    operation="current_subtask",
+                    data={
+                        "content": str(self.current_subtask),
+                        "status": "start"
+                    }
+                )
+
+            worker_start_time = time.time()
+            
+            # get the next action from the worker
+            executor_info = self.worker.generate_next_action(
+                Tu=instruction,
+                search_query=self.search_query,
+                subtask=self.current_subtask.name, # type: ignore
+                subtask_info=self.current_subtask.info, # type: ignore
+                future_tasks=self.global_state.get_remaining_subtasks(),
+                done_task=self.global_state.get_completed_subtasks(),
+                obs=self.global_state.get_obs_for_manager(),
+            )
+            
+            worker_execution_time = time.time() - worker_start_time
+            
+            self.global_state.log_operation(
+                module="agent",
+                operation="worker_execution",
+                data={
+                    "duration": worker_execution_time,
+                    "subtask": self.current_subtask.name # type: ignore
+                }
+            )
+
+            try:
+                grounding_start_time = time.time()
+                current_width, current_height = self.global_state.get_screen_size()
+                self.grounding.reset_screen_size(current_width, current_height)
+                self.grounding.assign_coordinates(executor_info["executor_plan"], observation)
+                plan_code = parse_single_code_from_string(executor_info["executor_plan"].split("Grounded Action")[-1])
+                plan_code = sanitize_code(plan_code)
+                plan_code = extract_first_agent_function(plan_code)
+                agent: Grounding = self.grounding # type: ignore
+                exec_code = eval(plan_code) # type: ignore
+                grounding_execution_time = time.time() - grounding_start_time
+                
+                # 记录grounding执行时间
+                self.global_state.log_operation(
+                    module="agent",
+                    operation="grounding_execution",
+                    data={
+                        "duration": grounding_execution_time,
+                        "content": plan_code
+                    }
+                )
+            except Exception as e:
+                logger.error("Error in parsing plan code: %s", e)
+                plan_code = "agent.wait(1.0)"
+                agent: Grounding = self.grounding # this agent will be used in next code
+                exec_code = eval(plan_code) # type: ignore
+                
+                # 记录grounding错误
+                self.global_state.log_operation(
+                    module="agent",
+                    operation="grounding_error",
+                    data={
+                        "content": str(e),
+                        "fallback_action": plan_code
+                    }
+                )
+
+            actions = [exec_code]
+
+            self.step_count += 1
+
+            # set the should_send_action flag to True if the executor returns an action
+            self.should_send_action = True
+
+            # replan on failure
+            if "fail" in actions[0]["type"].lower():
+                self.requires_replan = True
+                self.needs_next_subtask = True
+
+                # assign the failed subtask
+                self.global_state.add_failed_subtask(self.current_subtask) # type: ignore
+                self.failure_subtask = self.global_state.get_latest_failed_subtask()
+                
+                # 记录失败的子任务
+                self.global_state.log_operation(
+                    module="agent",
+                    operation="subtask_failed",
+                    data={
+                        "content": str(self.current_subtask),
+                        "status": "failed"
+                    }
+                )
+
+                # reset the step count, executor, and evaluator
+                self.reset_executor_state()
+
+                # if more subtasks are remaining, we don't want to send DONE to the environment but move on to the next subtask
+                if self.subtasks:
+                    self.should_send_action = False
+
+            # replan on subtask completion
+            elif "done" in actions[0]["type"].lower():
+                self.requires_replan = True
+                self.needs_next_subtask = True
+                self.failure_subtask = None
+                self.global_state.add_completed_subtask(self.current_subtask) # type: ignore
+                
+                # 记录完成的子任务
+                self.global_state.log_operation(
+                    module="agent",
+                    operation="subtask_completed",
+                    data={
+                        "content": str(self.current_subtask),
+                        "status": "completed"
+                    }
+                )
+
+                # reset the step count, executor, and evaluator
+                self.reset_executor_state()
+
+                # if more subtasks are remaining, we don't want to send DONE to the environment but move on to the next subtask
+                if self.subtasks:
+                    self.should_send_action = False
+                self.subtask_status = "Done"
+
+            self.turn_count += 1
+
+        # reset the should_send_action flag for next iteration
+        self.should_send_action = False
+
+        # concatenate the three info dictionaries
+        info = {
+            **{
+                k: v
+                for d in [planner_info or {}, executor_info or {}, evaluator_info or {}]
+                for k, v in d.items()
+            }
+        }
+        info.update(
+            {
+                "subtask": self.current_subtask.name, # type: ignore
+                "subtask_info": self.current_subtask.info, # type: ignore
+                "subtask_status": self.subtask_status,
+            }
+        )
+        
+        # 记录predict函数总执行时间
+        predict_total_time = time.time() - predict_start_time
+        self.global_state.log_operation(
+            module="agent",
+            operation="predict_execution",
+            data={
+                "duration": predict_total_time,
+                "step_count": self.step_count,
+                "turn_count": self.turn_count,
+                "subtask_status": self.subtask_status
+            }
+        )
+
+        return info, actions # type: ignore
+
+    def update_narrative_memory(self, trajectory: str) -> None:
+        """Update narrative memory from task trajectory
+
+        Args:
+            trajectory: String containing task execution trajectory
+        """
+        try:
+            reflection_path = os.path.join(
+                self.local_kb_path, self.platform, "narrative_memory.json"
+            )
+            try:
+                reflections = json.load(open(reflection_path))
+            except:
+                reflections = {}
+
+            if self.search_query not in reflections:
+                reflection = self.manager.summarize_narrative(trajectory)
+                reflections[self.search_query] = reflection
+
+            with open(reflection_path, "w") as f:
+                json.dump(reflections, f, indent=2)
+
+        except Exception as e:
+            logger.error(f"Failed to update narrative memory: {e}")
+
+    def update_episodic_memory(self, meta_data: Dict, subtask_trajectory: str) -> str:
+        """Update episodic memory from subtask trajectory
+
+        Args:
+            meta_data: Metadata about current subtask execution
+            subtask_trajectory: String containing subtask execution trajectory
+
+        Returns:
+            Updated subtask trajectory
+        """
+        subtask = meta_data["subtask"]
+        subtask_info = meta_data["subtask_info"]
+        subtask_status = meta_data["subtask_status"]
+        # Handle subtask trajectory
+        if subtask_status == "Start" or subtask_status == "Done":
+            # If it's a new subtask start, finalize the previous subtask trajectory if it exists
+            if subtask_trajectory:
+                subtask_trajectory += "\nSubtask Completed.\n"
+                subtask_key = subtask_trajectory.split(
+                    "\n----------------------\n\nPlan:\n"
+                )[0]
+                try:
+                    subtask_path = os.path.join(
+                        self.local_kb_path, self.platform, "episodic_memory.json"
+                    )
+                    kb = json.load(open(subtask_path))
+                except:
+                    kb = {}
+                if subtask_key not in kb.keys():
+                    subtask_summarization = self.manager.summarize_episode(
+                        subtask_trajectory
+                    )
+                    kb[subtask_key] = subtask_summarization
+                else:
+                    subtask_summarization = kb[subtask_key]
+                logger.info("subtask_key: %s", subtask_key)
+                logger.info("subtask_summarization: %s", subtask_summarization)
+                with open(subtask_path, "w") as fout:
+                    json.dump(kb, fout, indent=2)
+                # Reset for the next subtask
+                subtask_trajectory = ""
+            # Start a new subtask trajectory
+            subtask_trajectory = (
+                "Task:\n"
+                + self.search_query
+                + "\n\nSubtask: "
+                + subtask
+                + "\nSubtask Instruction: "
+                + subtask_info
+                + "\n----------------------\n\nPlan:\n"
+                + meta_data["executor_plan"]
+                + "\n"
+            )
+        elif subtask_status == "In":
+            # Continue appending to the current subtask trajectory if it's still ongoing
+            subtask_trajectory += (
+                "\n----------------------\n\nPlan:\n"
+                + meta_data["executor_plan"]
+                + "\n"
+            )
+
+        return subtask_trajectory
+
+class AgentSFast(UIAgent):
+    """Fast version of AgentS2 that directly generates actions using the fast_action_generator tool"""
+
+    def __init__(
+        self,
+        platform: str = platform.system().lower(),
+        screen_size: List[int] = [1920, 1080],
+        memory_root_path: str = os.getcwd(),
+        memory_folder_name: str = "kb_s2",
+        kb_release_tag: str = "v0.2.2",
+        enable_takeover: bool = False,
+        enable_search: bool = True,
+        enable_reflection: bool = True,
+        # enable_reflection: bool = False,
+    ):
+        """Initialize AgentSFast
+
+        Args:
+            platform: Operating system platform (darwin, linux, windows)
+            memory_root_path: Path to memory directory. Defaults to current working directory.
+            memory_folder_name: Name of memory folder. Defaults to "kb_s2".
+            kb_release_tag: Release tag for knowledge base. Defaults to "v0.2.2".
+            enable_takeover: Whether to enable user takeover functionality. Defaults to False.
+            enable_search: Whether to enable web search functionality. Defaults to True.
+            enable_reflection: Whether to enable reflection functionality. Defaults to True.
+        """
+        super().__init__(
+            platform,
+        )
+
+        self.memory_root_path = memory_root_path
+        self.memory_folder_name = memory_folder_name
+        self.kb_release_tag = kb_release_tag
+        self.screen_size = screen_size
+        self.enable_takeover = enable_takeover
+        self.enable_search = enable_search
+        self.enable_reflection = enable_reflection
+
+        # Load tools configuration from tools_config.json
+        tools_config_path = os.path.join(os.path.dirname(os.path.dirname(__file__)), "tools", "tools_config.json")
+        with open(tools_config_path, "r") as f:
+            self.tools_config = json.load(f)
+            print(f"Loaded tools configuration from: {tools_config_path}")
+            self.Tools_dict = {}
+            for tool in self.tools_config["tools"]:
+                tool_name = tool["tool_name"]
+                self.Tools_dict[tool_name] = {
+                    "provider": tool["provider"],
+                    "model": tool["model_name"]
+                }
+            print(f"Tools configuration: {self.Tools_dict}")
+
+        # Initialize agent's knowledge base path
+        self.local_kb_path = os.path.join(
+            self.memory_root_path, self.memory_folder_name
+        )
+
+        # Check if knowledge base exists
+        kb_platform_path = os.path.join(self.local_kb_path, self.platform)
+        if not os.path.exists(kb_platform_path):
+            print(f"Warning: Knowledge base for {self.platform} platform not found in {self.local_kb_path}")
+            os.makedirs(kb_platform_path, exist_ok=True)
+            print(f"Created directory: {kb_platform_path}")
+        else:
+            print(f"Found local knowledge base path: {kb_platform_path}")
+
+        self.reset()
+
+    def reset(self) -> None:
+        """Reset agent state and initialize components"""
+        # Initialize the fast action generator tool
+        self.fast_action_generator = Tools()
+        self.fast_action_generator_tool = "fast_action_generator_with_takeover" if self.enable_takeover else "fast_action_generator"
+        
+        # Get tool configuration from tools_config
+        tool_config = None
+        for tool in self.tools_config["tools"]:
+            if tool["tool_name"] == self.fast_action_generator_tool:
+                tool_config = tool
+                break
+        
+        # Prepare tool parameters
+        tool_params = {}
+        
+        # First check global search switch
+        if not self.enable_search:
+            # If global search is disabled, force disable search for this tool
+            tool_params["enable_search"] = False
+            logger.info(f"Configuring {self.fast_action_generator_tool} with search DISABLED (global switch off)")
+        else:
+            # If global search is enabled, check tool-specific config
+            if tool_config and "enable_search" in tool_config:
+                # Use enable_search from config file
+                enable_search = tool_config.get("enable_search", False)
+                tool_params["enable_search"] = enable_search
+                tool_params["search_provider"] = tool_config.get("search_provider", "bocha")
+                tool_params["search_model"] = tool_config.get("search_model", "")
+                
+                logger.info(f"Configuring {self.fast_action_generator_tool} with search enabled: {enable_search} (from config)")
+        
+        # Register the tool with parameters
+        self.fast_action_generator.register_tool(
+            self.fast_action_generator_tool, 
+            self.Tools_dict[self.fast_action_generator_tool]["provider"], 
+            self.Tools_dict[self.fast_action_generator_tool]["model"],
+            **tool_params
+        )
+
+        if self.enable_reflection:
+            self.reflection_agent = Tools()
+            self.reflection_agent.register_tool(
+                "traj_reflector", self.Tools_dict["traj_reflector"]["provider"],
+                self.Tools_dict["traj_reflector"]["model"])
+            self.reflections = []
+            self.planner_history = []
+
+        self.grounding_width, self.grounding_height = self.fast_action_generator.tools[self.fast_action_generator_tool].get_grounding_wh()
+        if self.grounding_width is None or self.grounding_height is None:
+            self.grounding_width = self.screen_size[0]
+            self.grounding_height = self.screen_size[1]
+        self.grounding = FastGrounding(
+            Tools_dict=self.Tools_dict,
+            platform=self.platform,
+            width=self.screen_size[0],
+            height=self.screen_size[1],
+            grounding_width=self.grounding_width,
+            grounding_height=self.grounding_height
+        )
+
+        # Reset state variables
+        self.step_count: int = 0
+        self.turn_count: int = 0
+        self.global_state: GlobalState = Registry.get("GlobalStateStore") # type: ignore
+        self.latest_action = None
+
+    def predict(self, instruction: str, observation: Dict) -> Tuple[Dict, List[str]]:
+        """Generate next action prediction using only the fast_action_generator tool
+
+        Args:
+            instruction: Natural language instruction
+            observation: Current UI state observation
+
+        Returns:
+            Tuple containing agent info dictionary and list of actions
+        """
+        import time
+        predict_start_time = time.time()
+        
+        fast_action_start_time = time.time()
+
+        reflection = None
+        if self.enable_reflection:
+            if self.turn_count == 0:
+                text_content = textwrap.dedent(f"""
+                    Task Description: {instruction}
+                    """)
+                self.reflection_agent.tools["traj_reflector"].llm_agent.add_message(
+                    text_content + "\n\nThe initial screen is provided. No action has been taken yet.",
+                    image_content=observation["screenshot"],
+                    role="user")
+                self.global_state.add_agent_log({
+                    "type": "passive",
+                    "content": "Reflection: " + text_content + "\n\nThe initial screen is provided. No action has been taken yet."
+                })
+            else:
+                agent_log = agent_log_to_string(self.global_state.get_agent_log())
+                text_content = f"Please refer to the agent log to understand the progress and context of the task so far.\n{agent_log}"
+
+                reflection_start = time.time()
+                reflection, total_tokens, cost_string = self.reflection_agent.execute_tool(
+                    "traj_reflector", {
+                        "str_input": text_content,
+                        "img_input": observation["screenshot"]
+                    })
+                reflection = str(reflection)
+                self.reflection_agent.reset("traj_reflector")
+                self.global_state.add_agent_log({
+                    "type": "passive",
+                    "content": "Reflection: " + reflection
+                })
+                logger.info(f"Trajectory reflector tokens: {total_tokens}, cost: {cost_string}")
+                reflection_time = time.time() - reflection_start
+                logger.info(f"[Timing] AgentSFast.traj_reflector execution time: {reflection_time:.2f} seconds")
+                self.reflections.append(reflection)
+                logger.info("REFLECTION: %s", reflection)
+                self.global_state.log_operation(
+                    module="agent",
+                    operation="reflection",
+                    data={
+                        "tokens": total_tokens,
+                        "cost": cost_string,
+                        "content": reflection,
+                        "duration": reflection_time
+                    })
+
+        agent_log = agent_log_to_string(self.global_state.get_agent_log())
+        
+        generator_message = textwrap.dedent(f"""
+            Task Description: {instruction}
+        """)
+        
+        generator_message += f"\n\nPlease refer to the agent log to understand the progress and context of the task so far.\n{agent_log}"
+
+        fast_action_start_time = time.time()
+        
+        plan, total_tokens, cost_string = self.fast_action_generator.execute_tool(
+            self.fast_action_generator_tool,
+            {
+                "str_input": generator_message,
+                "img_input": observation["screenshot"]
+            }
+        )
+        self.fast_action_generator.reset(self.fast_action_generator_tool)
+        
+        fast_action_execution_time = time.time() - fast_action_start_time
+        
+        self.global_state.log_operation(
+            module="agent",
+            operation="fast_action_execution",
+            data={
+                "duration": fast_action_execution_time,
+                "tokens": total_tokens,
+                "cost": cost_string
+            }
+        )
+        
+        logger.info("Fast Action Plan: %s", plan)
+
+        current_width, current_height = self.global_state.get_screen_size()
+        self.grounding.reset_screen_size(current_width, current_height)
+        try:
+            code_pattern = r"```python\s*(.*?)\s*```"
+            import re
+            match = re.search(code_pattern, plan, re.DOTALL)
+            
+            if match:
+                action_code = match.group(1).strip()
+                logger.info("Extracted action code: %s", action_code)
+                
+                agent: FastGrounding = self.grounding # type: ignore
+                exec_code = eval(action_code) # type: ignore
+                actions = [exec_code]
+                self.latest_action = action_code
+            else:
+                logger.warning("No code block found, trying to parse the entire response")
+                action_code = plan.strip()
+                
+                if action_code.startswith("agent."):
+                    agent: FastGrounding = self.grounding # type: ignore
+                    exec_code = eval(action_code) # type: ignore
+                    actions = [exec_code]
+                    self.latest_action = action_code
+                else:
+                    logger.error("Could not parse action, using wait action")
+                    self.global_state.add_agent_log({
+                        "type": "Wrong action code format",
+                        "content": action_code
+                    })
+                    agent: FastGrounding = self.grounding # type: ignore
+                    exec_code = eval("agent.wait(1000)") # type: ignore
+                    actions = [exec_code]
+                    self.latest_action = "agent.wait(1000)"
+        except Exception as e:
+            logger.error("Error in parsing action code: %s", e)
+            self.global_state.add_agent_log({
+                    "type": "Error in parsing action code",
+                    "content": str(e)  # Convert Exception to string
+                })
+            agent: FastGrounding = self.grounding # type: ignore
+            exec_code = eval("agent.wait(1000)") # type: ignore
+            actions = [exec_code]
+            self.latest_action = "agent.wait(1000)"
+            
+            self.global_state.log_operation(
+                module="agent",
+                operation="fast_action_error",
+                data={
+                    "content": str(e),
+                    "fallback_action": "agent.wait(1000)"
+                }
+            )
+
+        self.step_count += 1
+        self.turn_count += 1
+        
+        executor_info = {
+            "executor_plan": plan,
+            "reflection": reflection or "",
+            "plan_code": self.latest_action
+        }
+        
+        predict_total_time = time.time() - predict_start_time
+        self.global_state.log_operation(
+            module="agent",
+            operation="predict_execution_fast_direct",
+            data={
+                "duration": predict_total_time,
+                "step_count": self.step_count,
+                "turn_count": self.turn_count
+            }
+        )
+
+        return executor_info, actions