computer_use_demo/loop.py

"""
Agentic sampling loop that calls the Anthropic API and local implementation of computer use tools.
"""
import time
import json
from collections.abc import Callable
from enum import StrEnum

from anthropic import APIResponse
from anthropic.types.beta import BetaContentBlock, BetaMessage, BetaMessageParam
from computer_use_demo.tools import ToolResult


import torch

from computer_use_demo.gui_agent.planner.anthropic_agent import AnthropicActor
from computer_use_demo.executor.anthropic_executor import AnthropicExecutor
from computer_use_demo.gui_agent.planner.api_vlm_planner import APIVLMPlanner
from computer_use_demo.gui_agent.planner.local_vlm_planner import LocalVLMPlanner
from computer_use_demo.gui_agent.actor.showui_agent import ShowUIActor
from computer_use_demo.executor.showui_executor import ShowUIExecutor
from computer_use_demo.gui_agent.actor.uitars_agent import UITARS_Actor
from computer_use_demo.tools.colorful_text import colorful_text_showui, colorful_text_vlm
from computer_use_demo.tools.screen_capture import get_screenshot
from computer_use_demo.gui_agent.llm_utils.oai import encode_image

from computer_use_demo.tools.logger import logger



class APIProvider(StrEnum):
    ANTHROPIC = "anthropic"
    BEDROCK = "bedrock"
    VERTEX = "vertex"
    OPENAI = "openai"
    QWEN = "qwen"
    SSH = "ssh"


PROVIDER_TO_DEFAULT_MODEL_NAME: dict[APIProvider, str] = {
    APIProvider.ANTHROPIC: "claude-3-5-sonnet-20241022",
    APIProvider.BEDROCK: "anthropic.claude-3-5-sonnet-20241022-v2:0",
    APIProvider.VERTEX: "claude-3-5-sonnet-v2@20241022",
    APIProvider.OPENAI: "gpt-4o",
    APIProvider.QWEN: "qwen2vl",
    APIProvider.SSH: "qwen2-vl-2b",
}


def sampling_loop_sync(
    *,
    planner_model: str,
    planner_provider: APIProvider | None,
    actor_model: str,
    actor_provider: APIProvider | None,
    system_prompt_suffix: str,
    messages: list[BetaMessageParam],
    output_callback: Callable[[BetaContentBlock], None],
    tool_output_callback: Callable[[ToolResult, str], None],
    api_response_callback: Callable[[APIResponse[BetaMessage]], None],
    api_key: str,
    only_n_most_recent_images: int | None = None,
    max_tokens: int = 4096,
    selected_screen: int = 0,
    showui_max_pixels: int = 1344,
    showui_awq_4bit: bool = False,
    ui_tars_url: str = ""
):
    """
    Synchronous agentic sampling loop for the assistant/tool interaction of computer use.
    """

    # ---------------------------
    # Initialize Planner
    # ---------------------------
    if planner_model == "claude-3-5-sonnet-20241022":
        # Register Actor and Executor
        actor = AnthropicActor(
            model=planner_model, 
            provider=actor_provider, 
            system_prompt_suffix=system_prompt_suffix, 
            api_key=api_key, 
            api_response_callback=api_response_callback,
            max_tokens=max_tokens,
            only_n_most_recent_images=only_n_most_recent_images,
            selected_screen=selected_screen
        )

        executor = AnthropicExecutor(
            output_callback=output_callback,
            tool_output_callback=tool_output_callback,
            selected_screen=selected_screen
        )

        loop_mode = "unified"

    elif planner_model in ["gpt-4o", "gpt-4o-mini", "qwen2-vl-max"]:

        if torch.cuda.is_available(): device = torch.device("cuda")
        elif torch.backends.mps.is_available(): device = torch.device("mps")
        else: device = torch.device("cpu") # support: 'cpu', 'mps', 'cuda'
        logger.info(f"Model inited on device: {device}.")

        planner = APIVLMPlanner(
            model=planner_model,
            provider=planner_provider,
            system_prompt_suffix=system_prompt_suffix,
            api_key=api_key,
            api_response_callback=api_response_callback,
            selected_screen=selected_screen,
            output_callback=output_callback,
            device=device
        )
        loop_mode = "planner + actor"

    elif planner_model == "qwen2-vl-7b-instruct":
        planner = LocalVLMPlanner(
            model=planner_model,
            provider=planner_provider,
            system_prompt_suffix=system_prompt_suffix,
            api_key=api_key,
            api_response_callback=api_response_callback,
            selected_screen=selected_screen,
            output_callback=output_callback,
            device=device
        )
        loop_mode = "planner + actor"
    elif "ssh" in planner_model:
        if torch.cuda.is_available(): device = torch.device("cuda")
        elif torch.backends.mps.is_available(): device = torch.device("mps")
        else: device = torch.device("cpu") # support: 'cpu', 'mps', 'cuda'
        logger.info(f"Model inited on device: {device}.")
        planner = APIVLMPlanner(
            model=planner_model,
            provider=planner_provider,
            system_prompt_suffix=system_prompt_suffix,
            api_key=api_key,
            api_response_callback=api_response_callback,
            selected_screen=selected_screen,
            output_callback=output_callback,
            device=device
        )
        loop_mode = "planner + actor"
    else:
        logger.error(f"Planner Model {planner_model} not supported")
        raise ValueError(f"Planner Model {planner_model} not supported")
        

    # ---------------------------
    # Initialize Actor
    # ---------------------------
    if actor_model == "ShowUI":
        if showui_awq_4bit:
            showui_model_path = "./showui-2b-awq-4bit/"
        else:
            showui_model_path = "./showui-2b/"

        actor = ShowUIActor(
            model_path=showui_model_path,  
            device=device,  
            split='desktop',  # 'desktop' or 'phone'
            selected_screen=selected_screen,
            output_callback=output_callback,
            max_pixels=showui_max_pixels,
            awq_4bit=showui_awq_4bit
        )

        executor = ShowUIExecutor(
            output_callback=output_callback,
            tool_output_callback=tool_output_callback,
            selected_screen=selected_screen
        )
    elif actor_model == "UI-TARS":
        actor = UITARS_Actor(
            ui_tars_url=ui_tars_url,
            output_callback=output_callback,
            selected_screen=selected_screen
        )

    else:
        raise ValueError(f"Actor Model {actor_model} not supported")


    tool_result_content = None
    showui_loop_count = 0
    
    logger.info(f"Start the message loop. User messages: {messages}")

    if loop_mode == "unified":
        # ------------------------------
        # Unified loop: repeatedly call actor -> executor -> check tool_result -> maybe end
        # ------------------------------
        while True:
            # Call the actor with current messages
            response = actor(messages=messages)

            # Let the executor process that response, yielding any intermediate messages
            for message, tool_result_content in executor(response, messages):
                yield message

            # If executor didn't produce further content, we're done
            if not tool_result_content:
                return messages

            # If there is more tool content, treat that as user input
            messages.append({
                "content": tool_result_content,
                "role": "user"
            })

    elif loop_mode == "planner + actor":
        # ------------------------------------------------------
        # Planner + actor loop: 
        # 1) planner => get next_action
        # 2) If no next_action -> end 
        # 3) Otherwise actor => executor
        # 4) repeat
        # ------------------------------------------------------
        while True:
            # Step 1: Planner (VLM) response
            vlm_response = planner(messages=messages)

            # Step 2: Extract the "Next Action" from the planner output
            next_action = json.loads(vlm_response).get("Next Action")

            # Yield the next_action string, in case the UI or logs want to show it
            yield next_action

            # Step 3: Check if there are no further actions
            if not next_action or next_action in ("None", ""):
                final_sc, final_sc_path = get_screenshot(selected_screen=selected_screen)
                final_image_b64 = encode_image(str(final_sc_path))

                output_callback(
                    (
                        f"No more actions from {colorful_text_vlm}. End of task. Final State:\n"
                        f'<img src="data:image/png;base64,{final_image_b64}">'
                    ),
                    sender="bot"
                )
                yield None
                break

            # Step 4: Output an action message
            output_callback(
                f"{colorful_text_vlm} sending action to {colorful_text_showui}:\n{next_action}",
                sender="bot"
            )

            # Step 5: Actor response
            actor_response = actor(messages=next_action)
            yield actor_response

            # Step 6: Execute the actor response
            for message, tool_result_content in executor(actor_response, messages):
                time.sleep(0.5)  # optional small delay
                yield message

            # Step 7: Update conversation with embedding history of plan and actions
            messages.append({
                "role": "user",
                "content": [
                    "History plan:" + str(json.loads(vlm_response)),
                    "History actions:" + str(actor_response["content"])
                ]
            })

            logger.info(
                f"End of loop {showui_loop_count + 1}. "
                f"Messages: {str(messages)[:100000]}. "
                f"Total cost: $USD{planner.total_cost:.5f}"
            )


            # Increment loop counter
            showui_loop_count += 1