wehub-resource-sync
281 lines
9.5 KiB
Python
281 lines
9.5 KiB
Python
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
|