Agent S3

#!/usr/bin/env python3

import os
import io
import sys
import time
import requests
from dotenv import load_dotenv
from gui_agents.s3.agents.agent_s import AgentS3
from gui_agents.s3.agents.grounding import OSWorldACI
from orgo import Computer
import pyautogui

load_dotenv()

CONFIG = {
    "model": os.getenv("AGENT_MODEL", "gpt-5-2025-08-07"),
    "model_type": os.getenv("AGENT_MODEL_TYPE", "openai"),
    "grounding_model": os.getenv("GROUNDING_MODEL", "ui-tars-1.5-7b"),
    "grounding_type": os.getenv("GROUNDING_MODEL_TYPE", "huggingface"),
    "grounding_url": os.getenv("GROUNDING_URL", "http://localhost:8080"),
    "grounding_width": int(os.getenv("GROUNDING_WIDTH", "1920")),
    "grounding_height": int(os.getenv("GROUNDING_HEIGHT", "1080")),
    "max_steps": int(os.getenv("MAX_STEPS", "10")),
    "step_delay": float(os.getenv("STEP_DELAY", "0.5")),
    "remote": os.getenv("USE_CLOUD_ENVIRONMENT", "false").lower() == "true"
}


class LocalExecutor:
    def __init__(self):
        self.pyautogui = pyautogui
        if sys.platform == "win32":
            self.platform = "windows"
        elif sys.platform == "darwin":
            self.platform = "darwin"
        else:
            self.platform = "linux"

    def screenshot(self):
        img = self.pyautogui.screenshot()
        buffer = io.BytesIO()
        img.save(buffer, format="PNG")
        buffer.seek(0)
        return buffer.getvalue()

    def exec(self, code):
        exec(code, {"pyautogui": self.pyautogui, "time": time})

    def destroy(self):
        # No cleanup needed for local executor
        pass


class RemoteExecutor:
    def __init__(self):
        # Pass a size: SDKs before 0.0.47 default to 2 GB, which the API rejects.
        self.computer = Computer(ram=4, cpu=1)
        self.platform = "linux"
        # The Linux image ships python3 without pip or pyautogui, and the
        # actions Agent S3 emits import pyautogui. Install it once.
        self.computer.bash(
            "apt-get update -qq && apt-get install -y -qq python3-pip && "
            "pip3 install --break-system-packages -q pyautogui"
        )

    def screenshot(self):
        # Agent S3 wants raw image bytes, and response_format=binary returns
        # the PNG itself. The SDK's screenshot methods cannot read the
        # stored-image path the API returns by default.
        response = requests.get(
            f"https://www.orgo.ai/api/computers/{self.computer.computer_id}/screenshot",
            params={"response_format": "binary"},
            headers={"Authorization": f"Bearer {os.environ['ORGO_API_KEY']}"},
        )
        response.raise_for_status()
        return response.content

    def exec(self, code):
        result = self.computer.exec(code)
        if not result['success']:
            raise Exception(result.get('error', 'Execution failed'))
        if result['output']:
            print(f"Output: {result['output']}")

    def destroy(self):
        self.computer.destroy()


def create_agent(executor):
    engine_params = {
        "engine_type": CONFIG["model_type"],
        "model": CONFIG["model"],
    }

    # grounding_width/height describe the grounding model's output space,
    # not the computer's screen. Leave them as that model's docs specify.
    engine_params_for_grounding = {
        "engine_type": CONFIG["grounding_type"],
        "model": CONFIG["grounding_model"],
        "base_url": CONFIG["grounding_url"],
        "grounding_width": CONFIG["grounding_width"],
        "grounding_height": CONFIG["grounding_height"],
    }

    grounding_agent = OSWorldACI(
        env=None,
        platform=executor.platform,
        engine_params_for_generation=engine_params,
        engine_params_for_grounding=engine_params_for_grounding,
    )

    return AgentS3(
        engine_params,
        grounding_agent,
        platform=executor.platform,
    )


def run_task(agent, executor, instruction):
    print(f"\nTask: {instruction}")
    print(f"Mode: {'Remote' if CONFIG['remote'] else 'Local'}\n")

    for step in range(CONFIG["max_steps"]):
        print(f"Step {step + 1}/{CONFIG['max_steps']}")

        obs = {"screenshot": executor.screenshot()}
        info, action = agent.predict(instruction=instruction, observation=obs)

        if info:
            print(f"[thinking] {info}")

        if not action or not action[0]:
            print("Complete")
            return True

        try:
            print(f"[action] {action[0]}")
            executor.exec(action[0])
        except Exception as e:
            print(f"Error: {e}")
            instruction = "The previous action failed. Try a different approach."

        time.sleep(CONFIG["step_delay"])

    print("Max steps reached")
    return False


def main():
    executor = RemoteExecutor() if CONFIG["remote"] else LocalExecutor()
    try:
        agent = create_agent(executor)

        if len(sys.argv) > 1:
            run_task(agent, executor, " ".join(sys.argv[1:]))
        else:
            print("Interactive mode (type 'exit' to quit)\n")
            while True:
                task = input("Task: ").strip()
                if task == "exit":
                    break
                elif task:
                    run_task(agent, executor, task)
    finally:
        # Clean up
        executor.destroy()


if __name__ == "__main__":
    main()