/
OS-Worldf593f35from ast import parse
import logging
import json
from typing import List, Dict, Any, Optional, Tuple
from mm_agents.os_symphony.utils.common_utils import (
call_llm_formatted,
enhance_observation,
parse_code_from_string
)
from functools import partial
from mm_agents.os_symphony.utils.formatters import JSON_ANSWER_FORMATTER
from mm_agents.os_symphony.core.mllm import LMMAgent
from mm_agents.os_symphony.memory.procedural_memory import PROCEDURAL_MEMORY
import imagehash
import io
import os
from PIL import Image
import numpy as np
from skimage.metrics import structural_similarity as ssim
logger = logging.getLogger("desktopenv.agent")
class StepBehavior:
"""
Narrative Step Behavior.
Description of each step, cosists of generative agent (main agent)'s output, screenshot (if this step is milestone), and textual description.
The textual description shows that how the agent thought and did, and how the state changes.
"""
def __init__(self, is_milestone: bool, gen_output: str, summary: str, obs: Dict, action_dict: Dict):
self.is_milestone = is_milestone
self.gen_output = gen_output
self.obs = obs
self.summary = summary
self.action_dict = action_dict
# Variants for opyimizing the time complexity of loop detection
# --- 1. pHash ---
self.phash = None
# --- 2. SSIM ---
self.ssim_list = []
def _update_phash_ssim(self, history: List):
# Calculate the ssim_list of current obs
# Update pHash
cur_img = Image.open(io.BytesIO(self.obs["screenshot"]))
cur_img_gray = cur_img.convert('L')
cur_img_np = np.array(cur_img_gray)
self.phash = imagehash.phash(cur_img)
# Update ssim_list
for hs in history:
compare_img = Image.open(io.BytesIO(hs.obs["screenshot"]))
compare_img_gray = compare_img.convert('L')
compare_img_np = np.array(compare_img_gray)
self.ssim_list.append(ssim(cur_img_np, compare_img_np, data_range=cur_img_np.max() - compare_img_np.min()))
class ReflectionMemoryAgent:
"""
Reflection Memory Agent (RMA).
Responsible for maintaining long-term memory, extracting narratives from trajectories,
providing reflections to the Main Agent, and validating task completion status.
"""
def __init__(self, engine_params: Dict):
"""
Initialize the RMA.
Args:
- engine_params:
{
"engine_type": args.provider,
"model": args.model,
"base_url": args.model_url,
"api_key": args.model_api_key,
"temperature": getattr(args, "model_temperature", None),
}
- max_img_len: max image number to use in reflection process
"""
self.engine_params = engine_params
self.max_images = engine_params.get('max_images', 8)
self.memoryer_level = engine_params.get('memoryer_level', 3)
self.reset()
logger.info(f"ReflectionMemoryAgent initialized with:\n {self.engine_params}")
def reset(self):
"""Reset the code agent state."""
logger.debug("Resetting RMA state")
self.instruction = None
self.trajectory: List[StepBehavior] = []
self.knowledge_base: List[str] = []
self.last_code_step_idx = -1
'''
Control the count of images, we only use the maximum number of max_img_len images.
The update logic: the 0-th screenshot is always retained. If the total number of screenshots is less than max_img_len, all are kept; otherwise, starting from index 1, milestone screenshots are managed via FIFO.
'''
self.active_img_idx = []
self.reflection_agent = LMMAgent(
engine_params=self.engine_params,
system_prompt=PROCEDURAL_MEMORY.REFLECTION_SYSTEM_PROMPT,
)
self.behavior_agent = LMMAgent(
engine_params=self.engine_params,
system_prompt=PROCEDURAL_MEMORY.SUMMARIZE_STEP_SYSTEM_PROMPT
)
def add_instruction(self, instruction):
"""
[Interface] Main -> RMA
Main agent set the instruction to RMA.
"""
self.instruction = instruction
def _update_trajectory(self, step_behavior):
self.trajectory.append(step_behavior)
if len(self.active_img_idx) >= self.max_images:
if step_behavior.is_milestone:
self.active_img_idx.append(len(self.trajectory) - 1) # over max_img_len，only milestone image
del self.active_img_idx[1] # FIFO starts from index 1
else:
self.active_img_idx.append(len(self.trajectory) - 1) # less than max_img_len, feed all images
assert len(self.active_img_idx) <= self.max_images, "[RMA] Errors in updating StepBehavior!!"
def _summarize_step_behavior(
self,
generator_output: str,
cur_obs: Dict,
enhanced_obs: bytes | None,
is_milestone: bool,
mode: str = "gui",
code_exec_summary: str = "",
action_dict: Dict = {}
) -> Tuple[StepBehavior, bool]:
"""
[Interface] Main -> RMA
The Main Agent (MA) calls this method to "feed" the information of the just-completed step to the RMA.
RMA will internally process and store this step.
"""
if mode == "search":
is_success = "successful"
# summary is fixed
step_behavior = StepBehavior(
False,
generator_output,
"Search Agent was called last step, and a tutorial has been generated.",
cur_obs,
action_dict
)
elif mode == "code":
self.last_code_step_idx = len(self.trajectory)
is_success = "successful"
# the summary returned by the code agent
step_behavior = StepBehavior(
False,
generator_output,
f"Code Agent was called last step, and the summary of its trajectory is: \n---\n{code_exec_summary}\n---",
cur_obs,
action_dict
)
else: # common gui operation, use LLM to summarize
prev_obs = self.trajectory[-1].obs
text_content = f"""Computer Use Agent's Output: \n{generator_output}"""
self.behavior_agent.reset() # don't need history messages
updated_sys_prompt = (
self.behavior_agent.system_prompt + "\n" + text_content
)
self.behavior_agent.add_system_prompt(updated_sys_prompt)
self.behavior_agent.add_message(
text_content="This is the observation before executing action (attached below).",
image_content=prev_obs['screenshot'],
role="user",
put_text_last=False
)
self.behavior_agent.add_message(
text_content="This is the zoom-in view, which may help you to identify the operational region (attached below).",
image_content=enhanced_obs,
role="user",
put_text_last=False
)
self.behavior_agent.add_message(
text_content="This is the observation after executing action (attached below).",
image_content=cur_obs['screenshot'],
role="user",
put_text_last=False
)
required_fields = ["summary", "evaluation"]
format_checkers = [
partial(JSON_ANSWER_FORMATTER, required_fields)
]
full_response = call_llm_formatted(
self.behavior_agent,
format_checkers,
temperature=self.engine_params.get("temperture", 0.1),
)
response = parse_code_from_string(full_response)
try:
data = json.loads(response)
behavior_summary = data['summary']
is_success = data["evaluation"]
except Exception as e:
print("[RMA] Errors in generating step summary: ", e)
logger.info("Response is not a JSON object or miss required keys!")
behavior_summary = response
is_success = "successful"
step_behavior = StepBehavior(is_milestone, generator_output, behavior_summary, cur_obs, action_dict)
return step_behavior, is_success == "successful"
def get_reflection(
self,
cur_obs: Dict,
generator_output: str,
coordinates: List,
mode: str="gui",
code_exec_summary: str = "",
action_dict: Dict = {}
) -> Dict:
"""
[Interface] RMA -> Main
The Main Agent (MA) calls this method to get RMA's reflection before deciding the next action.
Args:
- cur_obs (Dict): The Main Agent's current observation (o_k).
- generator_output (str): The thoughts, screen analysis and action of Main Agent.
- coordinates (List): coordinates in the last operation step of Main Agent.
- mode(str): [gui, code, search]. Indicate which agent that main agent called last step.
- code_exec_summary: execution summary for code agent.
- action_dict: extracted action from generator output.
Returns:
- reflection_info(Dict): all the info related to reflection
"""
if self.memoryer_level == 0:
return {
"reflection": None,
"reflection_thoughts": None,
"existing_knowledge": None,
"is_milestone": False,
"new_knowledge": None,
"step_summary": None,
"hint": {
"gui_operation_error": False,
"lack_of_tutorial": False,
"code_error": False,
"loop_detection": None,
}
}
reflection = None
reflection_thought = None
if len(self.trajectory) == 0:
step_behavior = StepBehavior(
True,
"The initial screen is provided. No action has been taken yet.",
"The initial screen is provided. No action has been taken yet.",
cur_obs,
action_dict
)
step_behavior._update_phash_ssim(self.trajectory)
self._update_trajectory(step_behavior)
reflection_info = {
"reflection": reflection,
"reflection_thoughts": reflection_thought,
"existing_knowledge": "\n".join(self.knowledge_base),
"is_milestone": True,
"new_knowledge": "",
"step_summary": "",
"loop_detection": None
}
else:
### Step Summary
prev_obs = self.trajectory[-1].obs
enhanced_obs = None
if coordinates:
enhanced_obs, _, _, _, _ = enhance_observation(
prev_obs["screenshot"],
coordinates,
draw=True
)
# generate step behavior
step_behavior, last_gui_check = self._summarize_step_behavior(
generator_output,
cur_obs,
enhanced_obs,
False,
mode,
code_exec_summary,
action_dict
)
step_behavior._update_phash_ssim(self.trajectory)
### make additional hints
additional_hints = []
if not last_gui_check:
additional_hints.append(f"\t- Warning: The last GUI operation is unsuccessful. Careful review is required to avoid GUI Operation Error.")
code_error_hint = False
if self.last_code_step_idx != -1 and len(self.trajectory) - self.last_code_step_idx < 0:
code_error_hint = True
additional_hints.append(f"\t- Warning: The Computer Use Agent might in the verification stage of Code Agent. Careful review is required to avoid Code Error.")
# loop detection
from mm_agents.os_symphony.utils.loop_detection import detect_loop
is_loop, loop_details = detect_loop(full_trajectory=self.trajectory, N=3)
if is_loop and loop_details:
match_sequence_indices = loop_details["match_sequence_indices"]
loop_hint_message = f"\t- Warning: A potential LOOP has been detected between Step {match_sequence_indices[0]} and Step {match_sequence_indices[-1]}. Careful review is required to avoid Repetitive Behavior Error."
additional_hints.append(loop_hint_message)
self.reflection_agent.reset()
updated_sys_prompt = (
PROCEDURAL_MEMORY.REFLECTION_SYSTEM_PROMPT + "\n\n" +
f"---\n- **user instruction**: {self.instruction}\n" +
"- **existing knowledge**: \n" + "\n".join(self.knowledge_base) +
"\n- **additional_hints**: " + "\n".join(additional_hints) + "\n---"
)
# update system prompt
self.reflection_agent.add_system_prompt(updated_sys_prompt)
for i, step in enumerate(self.trajectory):
text_content = f"""### (Step {i}) history:\nsummary: '''\n{step.summary}\n'''"""
if i in self.active_img_idx:
if i == 0:
text_content += f"\ninitial screenshot:"
else:
text_content += f"\nscreenshot (after executing action): (attached below)"
self.reflection_agent.add_message(
text_content=text_content,
image_content=step.obs['screenshot'] if i in self.active_img_idx else None,
role="user",
put_text_last=False
)
text_content = f"""### (Last Step) CUA's output (has been finished):\n---\n{generator_output}\n---\nStep Summary:\n---\n{step_behavior.summary}\n---\nlatest_screenshot: (attached below)"""
self.reflection_agent.add_message(
text_content=text_content,
image_content=cur_obs['screenshot'],
role="user",
put_text_last=False
)
required_fields = ["is_milestone", "reflection", "knowledge"]
format_checkers = [
partial(JSON_ANSWER_FORMATTER, required_fields)
]
full_response = call_llm_formatted(
self.reflection_agent,
format_checkers
)
reflection_thought = full_response
response = parse_code_from_string(full_response)
try:
data = json.loads(response)
reflection = data['reflection']
is_milestone = data["is_milestone"]
knowledge = data['knowledge']
except Exception as e:
print("[RMA] Errors in dealing with reflection: ", e)
logger.info("Response is not a JSON object or miss required keys!")
reflection = response
is_milestone = False
knowledge = ""
if len(knowledge) > 0:
self.knowledge_base.append(knowledge)
if isinstance(is_milestone, str):
is_milestone = True if "true" in is_milestone.lower() else False
# update trajectory and is_milestone
self._update_trajectory(step_behavior)
if mode == "gui": # only gui opration can be considered as milestone
self.trajectory[-1].is_milestone = is_milestone
reflection_info = {
"reflection": reflection,
"reflection_thoughts": reflection_thought,
"existing_knowledge": "\n".join(self.knowledge_base),
"is_milestone": is_milestone,
"new_knowledge": knowledge,
"step_summary": step_behavior.summary,
"hint": {
"gui_operation_error": not last_gui_check,
"lack_of_tutorial": is_loop,
"code_error": code_error_hint,
"loop_detection": loop_details,
}
}
return reflection_info