chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,280 @@
|
||||
import logging
|
||||
from collections import defaultdict
|
||||
from typing import Dict, List, Optional, Tuple
|
||||
import platform
|
||||
|
||||
from gui_agents.s1.aci.ACI import ACI
|
||||
from gui_agents.s1.core.BaseModule import BaseModule
|
||||
from gui_agents.s1.core.Knowledge import KnowledgeBase
|
||||
from gui_agents.s1.core.ProceduralMemory import PROCEDURAL_MEMORY
|
||||
from gui_agents.s1.utils.common_utils import (
|
||||
Dag,
|
||||
Node,
|
||||
calculate_tokens,
|
||||
call_llm_safe,
|
||||
parse_dag,
|
||||
)
|
||||
|
||||
logger = logging.getLogger("desktopenv.agent")
|
||||
|
||||
NUM_IMAGE_TOKEN = 1105 # Value set of screen of size 1920x1080 for openai vision
|
||||
|
||||
|
||||
class Manager(BaseModule):
|
||||
def __init__(
|
||||
self,
|
||||
engine_params: Dict,
|
||||
grounding_agent: ACI,
|
||||
local_kb_path: str,
|
||||
search_engine: Optional[str] = None,
|
||||
multi_round: bool = False,
|
||||
platform: str = platform.system().lower(),
|
||||
):
|
||||
# TODO: move the prompt to Procedural Memory
|
||||
super().__init__(engine_params, platform)
|
||||
|
||||
# Initialize the ACI
|
||||
self.grounding_agent = grounding_agent
|
||||
|
||||
# Initialize the submodules of the Manager
|
||||
self.generator_agent = self._create_agent(PROCEDURAL_MEMORY.MANAGER_PROMPT)
|
||||
self.dag_translator_agent = self._create_agent(
|
||||
PROCEDURAL_MEMORY.DAG_TRANSLATOR_PROMPT
|
||||
)
|
||||
self.narrative_summarization_agent = self._create_agent(
|
||||
PROCEDURAL_MEMORY.TASK_SUMMARIZATION_PROMPT
|
||||
)
|
||||
self.episode_summarization_agent = self._create_agent(
|
||||
PROCEDURAL_MEMORY.SUBTASK_SUMMARIZATION_PROMPT
|
||||
)
|
||||
|
||||
self.local_kb_path = local_kb_path
|
||||
|
||||
self.knowledge_base = KnowledgeBase(self.local_kb_path, platform, engine_params)
|
||||
|
||||
self.planner_history = []
|
||||
|
||||
self.turn_count = 0
|
||||
self.search_engine = search_engine
|
||||
self.multi_round = multi_round
|
||||
self.platform = platform
|
||||
|
||||
def summarize_episode(self, trajectory):
|
||||
"""Summarize the episode experience for lifelong learning reflection
|
||||
Args:
|
||||
trajectory: str: The episode experience to be summarized
|
||||
"""
|
||||
|
||||
# Create Reflection on whole trajectories for next round trial, keep earlier messages as exemplars
|
||||
self.episode_summarization_agent.add_message(trajectory)
|
||||
subtask_summarization = call_llm_safe(self.episode_summarization_agent)
|
||||
self.episode_summarization_agent.add_message(subtask_summarization)
|
||||
|
||||
return subtask_summarization
|
||||
|
||||
def summarize_narrative(self, trajectory):
|
||||
"""Summarize the narrative experience for lifelong learning reflection
|
||||
Args:
|
||||
trajectory: str: The narrative experience to be summarized
|
||||
"""
|
||||
# Create Reflection on whole trajectories for next round trial
|
||||
self.narrative_summarization_agent.add_message(trajectory)
|
||||
lifelong_learning_reflection = call_llm_safe(self.narrative_summarization_agent)
|
||||
|
||||
return lifelong_learning_reflection
|
||||
|
||||
def _generate_step_by_step_plan(
|
||||
self, observation: Dict, instruction: str, failure_feedback: str = ""
|
||||
) -> Tuple[Dict, str]:
|
||||
agent = self.grounding_agent
|
||||
|
||||
self.active_apps = agent.get_active_apps(observation)
|
||||
|
||||
tree_input = agent.linearize_and_annotate_tree(observation)
|
||||
observation["linearized_accessibility_tree"] = tree_input
|
||||
|
||||
# Perform Retrieval only at the first planning step
|
||||
if self.turn_count == 0:
|
||||
|
||||
self.search_query = self.knowledge_base.formulate_query(
|
||||
instruction, observation
|
||||
)
|
||||
|
||||
retrieved_experience = ""
|
||||
integrated_knowledge = ""
|
||||
# Retrieve most similar narrative (task) experience
|
||||
most_similar_task, retrieved_experience = (
|
||||
self.knowledge_base.retrieve_narrative_experience(instruction)
|
||||
)
|
||||
logger.info(
|
||||
"SIMILAR TASK EXPERIENCE: %s",
|
||||
most_similar_task + "\n" + retrieved_experience.strip(),
|
||||
)
|
||||
|
||||
# Retrieve knowledge from the web if search_engine is provided
|
||||
if self.search_engine is not None:
|
||||
retrieved_knowledge = self.knowledge_base.retrieve_knowledge(
|
||||
instruction=instruction,
|
||||
search_query=self.search_query,
|
||||
search_engine=self.search_engine,
|
||||
)
|
||||
logger.info("RETRIEVED KNOWLEDGE: %s", retrieved_knowledge)
|
||||
|
||||
if retrieved_knowledge is not None:
|
||||
# Fuse the retrieved knowledge and experience
|
||||
integrated_knowledge = self.knowledge_base.knowledge_fusion(
|
||||
observation=observation,
|
||||
instruction=instruction,
|
||||
web_knowledge=retrieved_knowledge,
|
||||
similar_task=most_similar_task,
|
||||
experience=retrieved_experience,
|
||||
)
|
||||
logger.info("INTEGRATED KNOWLEDGE: %s", integrated_knowledge)
|
||||
|
||||
integrated_knowledge = integrated_knowledge or retrieved_experience
|
||||
|
||||
# Add the integrated knowledge to the task instruction in the system prompt
|
||||
if integrated_knowledge:
|
||||
instruction += f"\nYou may refer to some retrieved knowledge if you think they are useful.{integrated_knowledge}"
|
||||
|
||||
self.generator_agent.add_system_prompt(
|
||||
self.generator_agent.system_prompt.replace(
|
||||
"TASK_DESCRIPTION", instruction
|
||||
)
|
||||
)
|
||||
|
||||
generator_message = (
|
||||
f"Accessibility Tree: {tree_input}\n"
|
||||
f"The clipboard contains: {agent.clipboard}."
|
||||
f"The current open applications are {agent.get_active_apps(observation)}"
|
||||
+ (
|
||||
f" Previous plan failed at step: {failure_feedback}"
|
||||
if failure_feedback
|
||||
else ""
|
||||
)
|
||||
)
|
||||
|
||||
self.generator_agent.add_message(
|
||||
generator_message, image_content=observation.get("screenshot", None)
|
||||
)
|
||||
|
||||
logger.info("GENERATING HIGH LEVEL PLAN")
|
||||
|
||||
plan = call_llm_safe(self.generator_agent)
|
||||
|
||||
if plan == "":
|
||||
raise Exception("Plan Generation Failed - Fix the Prompt")
|
||||
|
||||
logger.info("HIGH LEVEL STEP BY STEP PLAN: %s", plan)
|
||||
|
||||
self.generator_agent.add_message(plan)
|
||||
|
||||
self.planner_history.append(plan)
|
||||
|
||||
self.turn_count += 1
|
||||
|
||||
input_tokens, output_tokens = calculate_tokens(self.generator_agent.messages)
|
||||
|
||||
# Set Cost based on GPT-4o
|
||||
cost = input_tokens * (0.0050 / 1000) + output_tokens * (0.0150 / 1000)
|
||||
|
||||
planner_info = {
|
||||
"search_query": self.search_query,
|
||||
"goal_plan": plan,
|
||||
"num_input_tokens_plan": input_tokens,
|
||||
"num_output_tokens_plan": output_tokens,
|
||||
"goal_plan_cost": cost,
|
||||
}
|
||||
|
||||
assert type(plan) == str
|
||||
|
||||
return planner_info, plan
|
||||
|
||||
def _generate_dag(self, instruction: str, plan: str) -> Tuple[Dict, Dag]:
|
||||
# Add initial instruction and plan to the agent's message history
|
||||
self.dag_translator_agent.add_message(
|
||||
f"Instruction: {instruction}\nPlan: {plan}"
|
||||
)
|
||||
|
||||
logger.info("GENERATING DAG")
|
||||
|
||||
# Generate DAG
|
||||
dag_raw = call_llm_safe(self.dag_translator_agent)
|
||||
|
||||
dag = parse_dag(dag_raw)
|
||||
|
||||
logger.info("Generated DAG: %s", dag_raw)
|
||||
|
||||
self.dag_translator_agent.add_message(dag_raw)
|
||||
|
||||
input_tokens, output_tokens = calculate_tokens(
|
||||
self.dag_translator_agent.messages
|
||||
)
|
||||
|
||||
# Set Cost based on GPT-4o
|
||||
cost = input_tokens * (0.0050 / 1000) + output_tokens * (0.0150 / 1000)
|
||||
|
||||
dag_info = {
|
||||
"dag": dag_raw,
|
||||
"num_input_tokens_dag": input_tokens,
|
||||
"num_output_tokens_dag": output_tokens,
|
||||
"dag_cost": cost,
|
||||
}
|
||||
|
||||
assert type(dag) == Dag
|
||||
|
||||
return dag_info, dag
|
||||
|
||||
def _topological_sort(self, dag: Dag) -> List[Node]:
|
||||
"""Topological sort of the DAG using DFS
|
||||
dag: Dag: Object representation of the DAG with nodes and edges
|
||||
"""
|
||||
|
||||
def dfs(node_name, visited, stack):
|
||||
visited[node_name] = True
|
||||
for neighbor in adj_list[node_name]:
|
||||
if not visited[neighbor]:
|
||||
dfs(neighbor, visited, stack)
|
||||
stack.append(node_name)
|
||||
|
||||
# Convert edges to adjacency list
|
||||
adj_list = defaultdict(list)
|
||||
for u, v in dag.edges:
|
||||
adj_list[u.name].append(v.name)
|
||||
|
||||
visited = {node.name: False for node in dag.nodes}
|
||||
stack = []
|
||||
|
||||
for node in dag.nodes:
|
||||
if not visited[node.name]:
|
||||
dfs(node.name, visited, stack)
|
||||
|
||||
# Return the nodes in topologically sorted order
|
||||
sorted_nodes = [
|
||||
next(n for n in dag.nodes if n.name == name) for name in stack[::-1]
|
||||
]
|
||||
return sorted_nodes
|
||||
|
||||
def get_action_queue(
|
||||
self,
|
||||
instruction: str,
|
||||
observation: Dict,
|
||||
failure_feedback: str = None,
|
||||
):
|
||||
"""Generate the action list based on the instruction
|
||||
instruction:str: Instruction for the task
|
||||
"""
|
||||
# Generate the high level plan
|
||||
planner_info, plan = self._generate_step_by_step_plan(
|
||||
observation, instruction, failure_feedback
|
||||
)
|
||||
|
||||
# Generate the DAG
|
||||
dag_info, dag = self._generate_dag(instruction, plan)
|
||||
|
||||
# Topological sort of the DAG
|
||||
action_queue = self._topological_sort(dag)
|
||||
|
||||
planner_info.update(dag_info)
|
||||
|
||||
return planner_info, action_queue
|
||||
Reference in New Issue
Block a user