cua.py

import asyncio
import base64
import inspect
import io
import json
import re

import openai
import PIL


class Scaler:
    """Wrapper for a computer that performs resizing and coordinate translation."""

    def __init__(self, computer, dimensions: tuple[int, int] | None = None):
        self.computer = computer
        self.size = dimensions
        self.screen_width = -1
        self.screen_height = -1

    @property
    def dimensions(self):
        if not self.size:
            # Scale to fit within 1440x900 while preserving aspect ratio
            # 1440x900 recommended by OpenAI for computer use
            # https://developers.openai.com/api/docs/guides/tools-computer-use
            width, height = self.computer.dimensions
            max_width, max_height = 1440, 900
            scale = min(max_width / width, max_height / height)
            if scale >= 1:
                self.size = (width, height)
            else:
                self.size = (int(width * scale), int(height * scale))
        return self.size

    async def screenshot(self) -> str:
        # Take a screenshot from the actual computer
        screenshot = await self.computer.screenshot()
        screenshot = base64.b64decode(screenshot)
        buffer = io.BytesIO(screenshot)
        image = PIL.Image.open(buffer)
        # Scale the screenshot
        self.screen_width, self.screen_height = image.size
        width, height = self.dimensions
        ratio = min(width / self.screen_width, height / self.screen_height)
        new_width = int(self.screen_width * ratio)
        new_height = int(self.screen_height * ratio)
        new_size = (new_width, new_height)
        resized_image = image.resize(new_size, PIL.Image.Resampling.LANCZOS)
        image = PIL.Image.new("RGB", (width, height), (0, 0, 0))
        image.paste(resized_image, (0, 0))
        buffer = io.BytesIO()
        image.save(buffer, format="PNG")
        buffer.seek(0)
        data = bytearray(buffer.getvalue())
        return base64.b64encode(data).decode("utf-8")

    async def click(self, x: int, y: int, button: str = "left") -> None:
        x, y = self._point_to_screen_coords(x, y)
        await self.computer.click(x, y, button=button)

    async def double_click(self, x: int, y: int) -> None:
        x, y = self._point_to_screen_coords(x, y)
        await self.computer.double_click(x, y)

    async def scroll(self, x: int, y: int, scroll_x: int, scroll_y: int) -> None:
        x, y = self._point_to_screen_coords(x, y)
        await self.computer.scroll(x, y, scroll_x, scroll_y)

    async def type(self, text: str) -> None:
        await self.computer.type(text)

    async def wait(self, ms: int = 1000) -> None:
        await self.computer.wait(ms)

    async def move(self, x: int, y: int) -> None:
        x, y = self._point_to_screen_coords(x, y)
        await self.computer.move(x, y)

    async def keypress(self, keys: list[str]) -> None:
        await self.computer.keypress(keys)

    async def drag(self, path: list[tuple[int, int]]) -> None:
        path = [self._point_to_screen_coords(*point) for point in path]
        await self.computer.drag(path)

    def _point_to_screen_coords(self, x, y):
        width, height = self.dimensions
        ratio = min(width / self.screen_width, height / self.screen_height)
        x = x / ratio
        y = y / ratio
        return int(x), int(y)


class Agent:
    """CUA agent to start and continue task execution"""

    def __init__(self, client, model: str, computer, logger=None):
        self.client = client
        self.model = model
        self.computer = computer
        self.logger = logger
        self.tools = {}
        self.extra_headers = None
        self.reasoning = {"effort": "medium", "summary": "detailed"}
        self.parallel_tool_calls = False
        self.start_task()

    def add_tool(self, tool: dict, func):
        name = tool["name"]
        self.tools[name] = (tool, func)

    @property
    def requires_user_input(self) -> bool:
        if self.response is None or len(self.response.output) == 0:
            return True
        item = self.response.output[-1]
        return item.type == "message" and item.role == "assistant"

    @property
    def messages(self) -> list[str]:
        result: list[str] = []
        if self.response:
            for item in self.response.output:
                if item.type == "message":
                    for content in item.content:
                        if content.type == "output_text":
                            result.append(content.text)
        return result

    def start_task(self):
        self.response = None

    async def continue_task(
        self,
        input: str | openai.types.responses.response_input_param.ResponseInputParam,
        temperature=None,
    ):
        inputs = []
        screenshot = ""
        previous_response = self.response
        previous_response_id = None
        if previous_response:
            previous_response_id = previous_response.id
            for item in previous_response.output:
                if item.type == "computer_call":
                    for action in item.actions:
                        action_args = vars(action) | {}
                        action_type = action_args.pop("type")
                        if action_type == "drag":
                            action_args["path"] = [(p.x, p.y) for p in action.path]
                        if action_type != "screenshot":
                            method = getattr(self.computer, action_type)
                            if inspect.iscoroutinefunction(method):
                                result = await method(**action_args)
                            else:
                                result = method(**action_args)
                    screenshot = await self.computer.screenshot()
                    output = openai.types.responses.response_input_param.ComputerCallOutput(
                        type="computer_call_output",
                        call_id=item.call_id,
                        output=openai.types.responses.response_input_param.ResponseComputerToolCallOutputScreenshotParam(
                            type="computer_screenshot",
                            image_url=f"data:image/png;base64,{screenshot}",
                            detail="original",
                        )
                    )
                    inputs.append(output)
                elif item.type == "function_call":
                    tool_name = item.name
                    kwargs = json.loads(item.arguments)
                    if tool_name not in self.tools:
                        raise ValueError(f"Unsupported tool '{tool_name}'.")
                    _, func = self.tools[tool_name]
                    if inspect.iscoroutinefunction(func):
                        result = await func(**kwargs)
                    else:
                        result = func(**kwargs)
                    output = (
                        openai.types.responses.response_input_param.FunctionCallOutput(
                            type="function_call_output",
                            call_id=item.call_id,
                            output=json.dumps(result),
                        )
                    )
                    inputs.append(output)
                elif item.type == "reasoning" or item.type == "message":
                    pass
                else:
                    message = f"Unsupported response output type '{item.type}'."
                    raise NotImplementedError(message)
        if isinstance(input, str):
            inputs.append(
                openai.types.responses.response_input_param.Message(
                    role="user",
                    content=input,
                )
            )
        else:
            inputs.extend(input)
        self.response = None
        wait = 0
        retry = 10
        while retry > 0:
            retry -= 1
            try:
                kwargs = {
                    "model": self.model,
                    "input": inputs,
                    "previous_response_id": previous_response_id,
                    "tools": self.get_tools(),
                    "reasoning": self.reasoning,
                    "truncation": "auto",
                    "extra_headers": self.extra_headers,
                    "parallel_tool_calls": self.parallel_tool_calls,
                    **({} if temperature is None else {"temperature": temperature}),
                }
                if isinstance(self.client, openai.AsyncOpenAI):
                    self.response = await self.client.responses.create(**kwargs)
                else:
                    self.response = self.client.responses.create(**kwargs)
                assert self.response.status == "completed"
                return
            except openai.RateLimitError as e:
                if retry <= 0:
                    if self.logger:
                        self.logger.exception("Rate limit exceeded.", exc_info=e)
                    raise
                match = re.search(r"Please try again in (\d+)s", e.message)
                wait = int(match.group(1)) if match else 10
                if self.logger:
                    self.logger.warning(
                        f"Rate limit exceeded. Waiting for {wait} seconds.",
                        exc_info=e,
                    )
                await asyncio.sleep(wait)
            except openai.InternalServerError as e:
                if retry <= 0:
                    if self.logger:
                        self.logger.exception(
                            f"Internal server error: {e.message}",
                            exc_info=e,
                        )
                    raise
                wait = max(wait, 10)
                if self.logger:
                    self.logger.warning(
                        f"Internal server error: {e.message}. Waiting for {wait} seconds.",
                        exc_info=e,
                    )
                await asyncio.sleep(wait)

    def get_tools(self) -> list[openai.types.responses.tool_param.ToolParam]:
        tools = [entry[0] for entry in self.tools.values()]
        computer_tool = openai.types.responses.ComputerToolParam(type="computer")
        return [computer_tool, *tools]