chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,254 @@
|
||||
import logging
|
||||
import re
|
||||
import textwrap
|
||||
from typing import Dict, List, Tuple
|
||||
import platform
|
||||
|
||||
from gui_agents.s2.agents.grounding import ACI
|
||||
from gui_agents.s2.core.module import BaseModule
|
||||
from gui_agents.s2.core.knowledge import KnowledgeBase
|
||||
from gui_agents.s2.memory.procedural_memory import PROCEDURAL_MEMORY
|
||||
from gui_agents.s2.utils.common_utils import (
|
||||
Node,
|
||||
calculate_tokens,
|
||||
call_llm_safe,
|
||||
parse_single_code_from_string,
|
||||
sanitize_code,
|
||||
extract_first_agent_function,
|
||||
)
|
||||
|
||||
logger = logging.getLogger("desktopenv.agent")
|
||||
|
||||
|
||||
class Worker(BaseModule):
|
||||
def __init__(
|
||||
self,
|
||||
engine_params: Dict,
|
||||
grounding_agent: ACI,
|
||||
local_kb_path: str,
|
||||
embedding_engine,
|
||||
platform: str = platform.system().lower(),
|
||||
enable_reflection: bool = True,
|
||||
use_subtask_experience: bool = True,
|
||||
):
|
||||
"""
|
||||
Worker receives a subtask list and active subtask and generates the next action for the to execute.
|
||||
Args:
|
||||
engine_params: Dict
|
||||
Parameters for the multimodal engine
|
||||
grounding_agent: Agent
|
||||
The grounding agent to use
|
||||
local_kb_path: str
|
||||
Path to knowledge base
|
||||
platform: str
|
||||
OS platform the agent runs on (darwin, linux, windows)
|
||||
enable_reflection: bool
|
||||
Whether to enable reflection
|
||||
use_subtask_experience: bool
|
||||
Whether to use subtask experience
|
||||
"""
|
||||
super().__init__(engine_params, platform)
|
||||
|
||||
self.grounding_agent = grounding_agent
|
||||
self.local_kb_path = local_kb_path
|
||||
self.embedding_engine = embedding_engine
|
||||
self.enable_reflection = enable_reflection
|
||||
self.use_subtask_experience = use_subtask_experience
|
||||
self.reset()
|
||||
|
||||
def reset(self):
|
||||
if self.platform != "linux":
|
||||
skipped_actions = ["set_cell_values"]
|
||||
else:
|
||||
skipped_actions = []
|
||||
|
||||
sys_prompt = PROCEDURAL_MEMORY.construct_worker_procedural_memory(
|
||||
type(self.grounding_agent), skipped_actions=skipped_actions
|
||||
).replace("CURRENT_OS", self.platform)
|
||||
|
||||
self.generator_agent = self._create_agent(sys_prompt)
|
||||
self.reflection_agent = self._create_agent(
|
||||
PROCEDURAL_MEMORY.REFLECTION_ON_TRAJECTORY
|
||||
)
|
||||
|
||||
self.knowledge_base = KnowledgeBase(
|
||||
embedding_engine=self.embedding_engine,
|
||||
local_kb_path=self.local_kb_path,
|
||||
platform=self.platform,
|
||||
engine_params=self.engine_params,
|
||||
)
|
||||
|
||||
self.turn_count = 0
|
||||
self.worker_history = []
|
||||
self.reflections = []
|
||||
self.cost_this_turn = 0
|
||||
self.screenshot_inputs = []
|
||||
self.planner_history = []
|
||||
self.max_trajector_length = 8
|
||||
|
||||
def flush_messages(self):
|
||||
# generator msgs are alternating [user, assistant], so 2 per round
|
||||
if len(self.generator_agent.messages) > 2 * self.max_trajector_length + 1:
|
||||
self.generator_agent.remove_message_at(1)
|
||||
self.generator_agent.remove_message_at(1)
|
||||
# reflector msgs are all [(user text, user image)], so 1 per round
|
||||
if len(self.reflection_agent.messages) > self.max_trajector_length + 1:
|
||||
self.reflection_agent.remove_message_at(1)
|
||||
|
||||
def generate_next_action(
|
||||
self,
|
||||
instruction: str,
|
||||
search_query: str,
|
||||
subtask: str,
|
||||
subtask_info: Dict,
|
||||
future_tasks: List[Node],
|
||||
done_task: List[Node],
|
||||
obs: Dict,
|
||||
) -> Tuple[Dict, List]:
|
||||
"""
|
||||
Predict the next action(s) based on the current observation.
|
||||
"""
|
||||
# Provide the top_app to the Grounding Agent to remove all other applications from the tree. At t=0, top_app is None
|
||||
agent = self.grounding_agent
|
||||
|
||||
# Get RAG knowledge, only update system message at t=0
|
||||
if self.turn_count == 0:
|
||||
if self.use_subtask_experience:
|
||||
subtask_query_key = (
|
||||
"Task:\n"
|
||||
+ search_query
|
||||
+ "\n\nSubtask: "
|
||||
+ subtask
|
||||
+ "\nSubtask Instruction: "
|
||||
+ subtask_info
|
||||
)
|
||||
retrieved_similar_subtask, retrieved_subtask_experience = (
|
||||
self.knowledge_base.retrieve_episodic_experience(subtask_query_key)
|
||||
)
|
||||
|
||||
# Dirty fix to replace id with element description during subtask retrieval
|
||||
pattern = r"\(\d+"
|
||||
retrieved_subtask_experience = re.sub(
|
||||
pattern, "(element_description", retrieved_subtask_experience
|
||||
)
|
||||
retrieved_subtask_experience = retrieved_subtask_experience.replace(
|
||||
"_id", "_description"
|
||||
)
|
||||
|
||||
logger.info(
|
||||
"SIMILAR SUBTASK EXPERIENCE: %s",
|
||||
retrieved_similar_subtask
|
||||
+ "\n"
|
||||
+ retrieved_subtask_experience.strip(),
|
||||
)
|
||||
instruction += "\nYou may refer to some similar subtask experience if you think they are useful. {}".format(
|
||||
retrieved_similar_subtask + "\n" + retrieved_subtask_experience
|
||||
)
|
||||
|
||||
self.generator_agent.add_system_prompt(
|
||||
self.generator_agent.system_prompt.replace(
|
||||
"SUBTASK_DESCRIPTION", subtask
|
||||
)
|
||||
.replace("TASK_DESCRIPTION", instruction)
|
||||
.replace("FUTURE_TASKS", ", ".join([f.name for f in future_tasks]))
|
||||
.replace("DONE_TASKS", ",".join(d.name for d in done_task))
|
||||
)
|
||||
|
||||
# Reflection generation does not add its own response, it only gets the trajectory
|
||||
reflection = None
|
||||
if self.enable_reflection:
|
||||
# Load the initial subtask info
|
||||
if self.turn_count == 0:
|
||||
text_content = textwrap.dedent(f"""
|
||||
Subtask Description: {subtask}
|
||||
Subtask Information: {subtask_info}
|
||||
Current Trajectory below:
|
||||
""")
|
||||
updated_sys_prompt = (
|
||||
self.reflection_agent.system_prompt + "\n" + text_content
|
||||
)
|
||||
self.reflection_agent.add_system_prompt(updated_sys_prompt)
|
||||
self.reflection_agent.add_message(
|
||||
text_content="The initial screen is provided. No action has been taken yet.",
|
||||
image_content=obs["screenshot"],
|
||||
role="user",
|
||||
)
|
||||
# Load the latest action
|
||||
else:
|
||||
text_content = self.clean_worker_generation_for_reflection(
|
||||
self.planner_history[-1]
|
||||
)
|
||||
self.reflection_agent.add_message(
|
||||
text_content=text_content,
|
||||
image_content=obs["screenshot"],
|
||||
role="user",
|
||||
)
|
||||
reflection = call_llm_safe(self.reflection_agent)
|
||||
self.reflections.append(reflection)
|
||||
logger.info("REFLECTION: %s", reflection)
|
||||
|
||||
generator_message = (
|
||||
f"\nYou may use this reflection on the previous action and overall trajectory: {reflection}\n"
|
||||
if reflection and self.turn_count > 0
|
||||
else ""
|
||||
) + f"Text Buffer = [{','.join(agent.notes)}]."
|
||||
|
||||
# Only provide subinfo in the very first message to avoid over influence and redundancy
|
||||
if self.turn_count == 0:
|
||||
generator_message += f"Remember only complete the subtask: {subtask}\n"
|
||||
generator_message += f"You can use this extra information for completing the current subtask: {subtask_info}.\n"
|
||||
|
||||
# logger.info("GENERATOR MESSAGE: %s", generator_message)
|
||||
|
||||
self.generator_agent.add_message(
|
||||
generator_message, image_content=obs["screenshot"], role="user"
|
||||
)
|
||||
|
||||
plan = call_llm_safe(self.generator_agent)
|
||||
self.planner_history.append(plan)
|
||||
logger.info("PLAN: %s", plan)
|
||||
self.generator_agent.add_message(plan, role="assistant")
|
||||
|
||||
# Calculate input/output tokens and gpt-4o cost
|
||||
input_tokens, output_tokens = calculate_tokens(self.generator_agent.messages)
|
||||
cost = input_tokens * (0.0050 / 1000) + output_tokens * (0.0150 / 1000)
|
||||
self.cost_this_turn += cost
|
||||
logger.info("EXECTUOR COST: %s", self.cost_this_turn)
|
||||
|
||||
# Use the DescriptionBasedACI to convert agent_action("desc") into agent_action([x, y])
|
||||
try:
|
||||
agent.assign_coordinates(plan, obs)
|
||||
plan_code = parse_single_code_from_string(plan.split("Grounded Action")[-1])
|
||||
plan_code = sanitize_code(plan_code)
|
||||
plan_code = extract_first_agent_function(plan_code)
|
||||
exec_code = eval(plan_code)
|
||||
except Exception as e:
|
||||
logger.error("Error in parsing plan code: %s", e)
|
||||
plan_code = "agent.wait(1.0)"
|
||||
exec_code = eval(plan_code)
|
||||
|
||||
executor_info = {
|
||||
"current_subtask": subtask,
|
||||
"current_subtask_info": subtask_info,
|
||||
"executor_plan": plan,
|
||||
"plan_code": plan_code,
|
||||
"reflection": reflection,
|
||||
"num_input_tokens_executor": input_tokens,
|
||||
"num_output_tokens_executor": output_tokens,
|
||||
}
|
||||
self.turn_count += 1
|
||||
|
||||
self.screenshot_inputs.append(obs["screenshot"])
|
||||
self.flush_messages()
|
||||
|
||||
return executor_info, [exec_code]
|
||||
|
||||
# Removes the previous action verification, and removes any extraneous grounded actions
|
||||
def clean_worker_generation_for_reflection(self, worker_generation: str) -> str:
|
||||
# Remove the previous action verification
|
||||
res = worker_generation[worker_generation.find("(Screenshot Analysis)") :]
|
||||
action = extract_first_agent_function(worker_generation)
|
||||
# Cut off extra grounded actions
|
||||
res = res[: res.find("(Grounded Action)")]
|
||||
res += f"(Grounded Action)\n```python\n{action}\n```\n"
|
||||
return res
|
||||
Reference in New Issue
Block a user