From ed0a8763d279c6c979f4531dc4bc23fabc57db70 Mon Sep 17 00:00:00 2001
From: vyokky <7678676@qq.com>
Date: Thu, 4 Jul 2024 19:51:03 +0800
Subject: [PATCH] fix prompt and safeguard

---
 ufo/agents/processors/app_agent_processor.py | 9 +++++++++
 ufo/agents/processors/basic.py               | 5 +++++
 ufo/agents/states/app_agent_state.py         | 3 +++
 ufo/prompts/share/base/app_agent.yaml        | 4 +++-
 4 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/ufo/agents/processors/app_agent_processor.py b/ufo/agents/processors/app_agent_processor.py
index 3e12528f..fc07affe 100644
--- a/ufo/agents/processors/app_agent_processor.py
+++ b/ufo/agents/processors/app_agent_processor.py
@@ -406,6 +406,15 @@ def update_memory(self) -> None:
             key: self._memory_data.to_dict().get(key) for key in configs["HISTORY_KEYS"]
         }
 
+        if self.is_confirm():
+
+            if self._is_resumed:
+                self._memory_data.set_values_from_dict({"UserConfirm": "Yes"})
+                memorized_action["UserConfirm"] = "Yes"
+            else:
+                self._memory_data.set_values_from_dict({"UserConfirm": "No"})
+                memorized_action["UserConfirm"] = "No"
+
         # Save the screenshot to the blackboard if the SaveScreenshot flag is set to True by the AppAgent.
         self._update_image_blackboard()
         self.host_agent.blackboard.add_trajectories(memorized_action)
diff --git a/ufo/agents/processors/basic.py b/ufo/agents/processors/basic.py
index f36ada70..fea39101 100644
--- a/ufo/agents/processors/basic.py
+++ b/ufo/agents/processors/basic.py
@@ -53,6 +53,7 @@ def __init__(self, agent: BasicAgent, context: Context) -> None:
         self._results = None
         self._question_list = []
         self._agent_status_manager = self.agent.status_manager
+        self._is_resumed = False
 
     def process(self) -> None:
         """
@@ -120,6 +121,8 @@ def resume(self) -> None:
         Resume the process of action execution after the session is paused.
         """
 
+        self._is_resumed = True
+
         # Step 1: Execute the action.
         self.execute_action()
 
@@ -132,6 +135,8 @@ def resume(self) -> None:
         # Step 4: Update the step.
         self.update_step()
 
+        self._is_resumed = False
+
     @abstractmethod
     def print_step_info(self) -> None:
         """
diff --git a/ufo/agents/states/app_agent_state.py b/ufo/agents/states/app_agent_state.py
index 4af12d6d..7f38554f 100644
--- a/ufo/agents/states/app_agent_state.py
+++ b/ufo/agents/states/app_agent_state.py
@@ -356,6 +356,9 @@ def handle(self, agent: "AppAgent", context: Optional["Context"] = None) -> None
         # If the safe guard is not enabled, the agent should resume the task.
         if not configs["SAFE_GUARD"]:
             agent.process_resume()
+            self._confirm = True
+
+            return
 
         agent_processor = agent.processor
 
diff --git a/ufo/prompts/share/base/app_agent.yaml b/ufo/prompts/share/base/app_agent.yaml
index 3de39930..317065fc 100644
--- a/ufo/prompts/share/base/app_agent.yaml
+++ b/ufo/prompts/share/base/app_agent.yaml
@@ -110,7 +110,8 @@ system: |-
   - The 'Copilot' Add-in can help you with some special requests, such as creating a slide in PowerPoint from a Word document, or summarizing the entire ppt.
   - Saving a ppt file into pdf format can be done by clicking the "Save As Adobe PDF" button.
   - You are given the help documents of the application or/and the online search results for completing the user request. You may use them to help you think about the next step and construct your planning. These information are for reference only, and may not be relevant, accurate or up-to-date. You must rely more on the current screenshots and control item list to make the decision.
-
+  - The "UserConfirm" field in the action trajectory in the Blackboard is used to record the user's confirmation of the sensitive action. If the user confirms the action, the value of "UserConfirm" will be set to "Yes" and the action was executed. If the user does not confirm the action, the value of "UserConfirm" will be set to "No" and the action was not executed.
+  
   ## Here are other tips for you to complete the user request:
   {tips}
 
@@ -208,6 +209,7 @@ system_nonvisual: |-
   - When inputting the searched text on Google, you must use the Search Box, which is a ComboBox type of control item. Do not use the address bar to input the searched text.
   - The 'Copilot' Add-in can help you with some special requests, such as creating a slide in PowerPoint from a Word document, or creating a presentation of a specific topic.
   - You are given the help documents of the application or/and the online search results for completing the user request. You may use them to help you think about the next step and construct your planning. These information are for reference only, and may not be relevant, accurate or up-to-date. You must rely more on the current control item list to make the decision.
+  - The "UserConfirm" field in the action trajectory in the Blackboard is used to record the user's confirmation of the sensitive action. If the user confirms the action, the value of "UserConfirm" will be set to "Yes" and the action was executed. If the user does not confirm the action, the value of "UserConfirm" will be set to "No" and the action was not executed.
 
   ## Here are other tips for you to complete the user request:
   {tips}