""" 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'' ), 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