chore: import upstream snapshot with attribution

gui_agents/s3/utils/common_utils.py

@@ -0,0 +1,198 @@
+import re
+import time
+from io import BytesIO
+from PIL import Image
+
+from typing import Tuple, Dict
+
+from gui_agents.s3.memory.procedural_memory import PROCEDURAL_MEMORY
+
+import logging
+
+logger = logging.getLogger("desktopenv.agent")
+
+
+def create_pyautogui_code(agent, code: str, obs: Dict) -> str:
+    """
+    Attempts to evaluate the code into a pyautogui code snippet with grounded actions using the observation screenshot.
+
+    Args:
+        agent (ACI): The grounding agent to use for evaluation.
+        code (str): The code string to evaluate.
+        obs (Dict): The current observation containing the screenshot.
+
+    Returns:
+        exec_code (str): The pyautogui code to execute the grounded action.
+
+    Raises:
+        Exception: If there is an error in evaluating the code.
+    """
+    agent.assign_screenshot(obs)  # Necessary for grounding
+    exec_code = eval(code)
+    return exec_code
+
+
+def call_llm_safe(
+    agent, temperature: float = 0.0, use_thinking: bool = False, **kwargs
+) -> str:
+    # Retry if fails
+    max_retries = 3  # Set the maximum number of retries
+    attempt = 0
+    response = ""
+    while attempt < max_retries:
+        try:
+            response = agent.get_response(
+                temperature=temperature, use_thinking=use_thinking, **kwargs
+            )
+            assert response is not None, "Response from agent should not be None"
+            print("Response success!")
+            break  # If successful, break out of the loop
+        except Exception as e:
+            attempt += 1
+            print(f"Attempt {attempt} failed: {e}")
+            if attempt == max_retries:
+                print("Max retries reached. Handling failure.")
+        time.sleep(1.0)
+    return response if response is not None else ""
+
+
+def call_llm_formatted(generator, format_checkers, **kwargs):
+    """
+    Calls the generator agent's LLM and ensures correct formatting.
+
+    Args:
+        generator (ACI): The generator agent to call.
+        obs (Dict): The current observation containing the screenshot.
+        format_checkers (Callable): Functions that take the response and return a tuple of (success, feedback).
+        **kwargs: Additional keyword arguments for the LLM call.
+
+    Returns:
+        response (str): The formatted response from the generator agent.
+    """
+    max_retries = 3  # Set the maximum number of retries
+    attempt = 0
+    response = ""
+    if kwargs.get("messages") is None:
+        messages = (
+            generator.messages.copy()
+        )  # Copy messages to avoid modifying the original
+    else:
+        messages = kwargs["messages"]
+        del kwargs["messages"]  # Remove messages from kwargs to avoid passing it twice
+    while attempt < max_retries:
+        response = call_llm_safe(generator, messages=messages, **kwargs)
+
+        # Prepare feedback messages for incorrect formatting
+        feedback_msgs = []
+        for format_checker in format_checkers:
+            success, feedback = format_checker(response)
+            if not success:
+                feedback_msgs.append(feedback)
+        if not feedback_msgs:
+            # logger.info(f"Response formatted correctly on attempt {attempt} for {generator.engine.model}")
+            break
+        logger.error(
+            f"Response formatting error on attempt {attempt} for {generator.engine.model}. Response: {response} {', '.join(feedback_msgs)}"
+        )
+        messages.append(
+            {
+                "role": "assistant",
+                "content": [{"type": "text", "text": response}],
+            }
+        )
+        logger.info(f"Bad response: {response}")
+        delimiter = "\n- "
+        formatting_feedback = f"- {delimiter.join(feedback_msgs)}"
+        messages.append(
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "text",
+                        "text": PROCEDURAL_MEMORY.FORMATTING_FEEDBACK_PROMPT.replace(
+                            "FORMATTING_FEEDBACK", formatting_feedback
+                        ),
+                    }
+                ],
+            }
+        )
+        logger.info("Feedback:\n%s", formatting_feedback)
+
+        attempt += 1
+        if attempt == max_retries:
+            logger.error(
+                "Max retries reached when formatting response. Handling failure."
+            )
+        time.sleep(1.0)
+    return response
+
+
+def split_thinking_response(full_response: str) -> Tuple[str, str]:
+    try:
+        # Extract thoughts section
+        thoughts = full_response.split("<thoughts>")[-1].split("</thoughts>")[0].strip()
+
+        # Extract answer section
+        answer = full_response.split("<answer>")[-1].split("</answer>")[0].strip()
+
+        return answer, thoughts
+    except Exception as e:
+        return full_response, ""
+
+
+def parse_code_from_string(input_string):
+    """Parses a string to extract each line of code enclosed in triple backticks (```)
+
+    Args:
+        input_string (str): The input string containing code snippets.
+
+    Returns:
+        str: The last code snippet found in the input string, or an empty string if no code is found.
+    """
+    input_string = input_string.strip()
+
+    # This regular expression will match both ```code``` and ```python code```
+    # and capture the `code` part. It uses a non-greedy match for the content inside.
+    pattern = r"```(?:\w+\s+)?(.*?)```"
+
+    # Find all non-overlapping matches in the string
+    matches = re.findall(pattern, input_string, re.DOTALL)
+    if len(matches) == 0:
+        # return []
+        return ""
+    relevant_code = matches[
+        -1
+    ]  # We only care about the last match given it is the grounded action
+    return relevant_code
+
+
+def extract_agent_functions(code):
+    """Extracts all agent function calls from the given code.
+
+    Args:
+        code (str): The code string to search for agent function calls.
+
+    Returns:
+        list: A list of all agent function calls found in the code.
+    """
+    pattern = r"(agent\.\w+\(\s*.*\))"  # Matches
+    return re.findall(pattern, code)
+
+
+def compress_image(image_bytes: bytes = None, image: Image = None) -> bytes:
+    """Compresses an image represented as bytes.
+
+    Compression involves resizing image into half its original size and saving to webp format.
+
+    Args:
+        image_bytes (bytes): The image data to compress.
+
+    Returns:
+        bytes: The compressed image data.
+    """
+    if not image:
+        image = Image.open(BytesIO(image_bytes))
+    output = BytesIO()
+    image.save(output, format="WEBP")
+    compressed_image_bytes = output.getvalue()
+    return compressed_image_bytes

gui_agents/s3/utils/formatters.py

@@ -0,0 +1,58 @@
+"""This file contains various formatting checks used to reprompt an agent for correctly formatted responses."""
+
+from gui_agents.s3.utils.common_utils import (
+    extract_agent_functions,
+    parse_code_from_string,
+    create_pyautogui_code,
+    split_thinking_response,
+)
+
+single_action_check = (
+    lambda response: len(extract_agent_functions(parse_code_from_string(response))) == 1
+)
+single_action_error_msg = (
+    "Incorrect code: There must be a single agent action in the code response."
+)
+SINGLE_ACTION_FORMATTER = lambda response: (
+    single_action_check(response),
+    single_action_error_msg,
+)
+
+
+def _attempt_code_creation(agent, code, obs):
+    """Attempts to create a pyautogui code snippet from the response code"""
+    try:
+        return create_pyautogui_code(agent, code, obs)
+    except Exception as e:
+        return None
+
+
+code_valid_check = (
+    lambda agent, obs, response: _attempt_code_creation(
+        agent, parse_code_from_string(response), obs
+    )
+    is not None
+)
+code_valid_error_msg = "Incorrect code: The agent action must be a valid function and use valid parameters from the docstring list."
+CODE_VALID_FORMATTER = lambda agent, obs, response: (
+    code_valid_check(agent, obs, response),
+    code_valid_error_msg,
+)
+
+thoughts_answer_tag_check = lambda response: split_thinking_response(response)[1] != ""
+thoughts_answer_tag_error_msg = "Incorrect response: The response must contain both <thoughts>...</thoughts> and <answer>...</answer> tags."
+THOUGHTS_ANSWER_TAG_FORMATTER = lambda response: (
+    thoughts_answer_tag_check(response),
+    thoughts_answer_tag_error_msg,
+)
+
+integer_answer_check = (
+    lambda response: split_thinking_response(response)[0].strip().isdigit()
+)
+integer_answer_error_msg = (
+    "Incorrect response: The <answer>...</answer> tag must contain a single integer."
+)
+INTEGER_ANSWER_FORMATTER = lambda response: (
+    integer_answer_check(response),
+    integer_answer_error_msg,
+)

gui_agents/s3/utils/local_env.py

@@ -0,0 +1,77 @@
+import subprocess
+import sys
+from typing import Dict
+
+
+class LocalController:
+    """Minimal controller to execute bash and python code locally.
+
+    WARNING: Executing arbitrary code is dangerous. Only enable/use this in trusted
+    environments and with trusted inputs.
+    """
+
+    def run_bash_script(self, code: str, timeout: int = 30) -> Dict:
+        try:
+            proc = subprocess.run(
+                ["/bin/bash", "-lc", code],
+                capture_output=True,
+                text=True,
+                timeout=timeout,
+            )
+            output = (proc.stdout or "") + (proc.stderr or "")
+
+            print("BASH OUTPUT =======================================")
+            print(output)
+            print("BASH OUTPUT =======================================")
+
+            return {
+                "status": "ok" if proc.returncode == 0 else "error",
+                "returncode": proc.returncode,
+                "output": output,
+                "error": "",
+            }
+        except subprocess.TimeoutExpired as e:
+            return {
+                "status": "error",
+                "returncode": -1,
+                "output": e.stdout or "",
+                "error": f"TimeoutExpired: {str(e)}",
+            }
+        except Exception as e:
+            return {
+                "status": "error",
+                "returncode": -1,
+                "output": "",
+                "error": str(e),
+            }
+
+    def run_python_script(self, code: str) -> Dict:
+        try:
+            proc = subprocess.run(
+                [sys.executable, "-c", code],
+                capture_output=True,
+                text=True,
+            )
+            print("PYTHON OUTPUT =======================================")
+            print(proc.stdout or "")
+            print("PYTHON OUTPUT =======================================")
+            return {
+                "status": "ok" if proc.returncode == 0 else "error",
+                "return_code": proc.returncode,
+                "output": proc.stdout or "",
+                "error": proc.stderr or "",
+            }
+        except Exception as e:
+            return {
+                "status": "error",
+                "return_code": -1,
+                "output": "",
+                "error": str(e),
+            }
+
+
+class LocalEnv:
+    """Simple environment that provides a controller compatible with CodeAgent."""
+
+    def __init__(self):
+        self.controller = LocalController()