baqr/computer_use_ootb
0
1"""2Entrypoint for Gradio, see https://gradio.app/3"""4 5import platform6import asyncio7import base648import os9import io10import json11from datetime import datetime12from enum import StrEnum13from functools import partial14from pathlib import Path15from typing import cast, Dict16from PIL import Image17 18import gradio as gr19from anthropic import APIResponse20from anthropic.types import TextBlock21from anthropic.types.beta import BetaMessage, BetaTextBlock, BetaToolUseBlock22from anthropic.types.tool_use_block import ToolUseBlock23 24from screeninfo import get_monitors25from computer_use_demo.tools.logger import logger, truncate_string26 27logger.info("Starting the gradio app")28 29screens = get_monitors()30logger.info(f"Found {len(screens)} screens")31 32from computer_use_demo.loop import APIProvider, sampling_loop_sync33 34from computer_use_demo.tools import ToolResult35from computer_use_demo.tools.computer import get_screen_details36SCREEN_NAMES, SELECTED_SCREEN_INDEX = get_screen_details()37 38API_KEY_FILE = "./api_keys.json"39 40WARNING_TEXT = "⚠️ Security Alert: Do not provide access to sensitive accounts or data, as malicious web content can hijack Agent's behavior. Keep monitor on the Agent's actions."41 42 43def setup_state(state):44 45 if "messages" not in state:46 state["messages"] = []47 # -------------------------------48 if "planner_model" not in state:49 state["planner_model"] = "gpt-4o" # default50 if "actor_model" not in state:51 state["actor_model"] = "ShowUI" # default52 53 if "planner_provider" not in state:54 state["planner_provider"] = "openai" # default55 if "actor_provider" not in state:56 state["actor_provider"] = "local" # default57 58 # Fetch API keys from environment variables59 if "openai_api_key" not in state: 60 state["openai_api_key"] = os.getenv("OPENAI_API_KEY", "")61 if "anthropic_api_key" not in state:62 state["anthropic_api_key"] = os.getenv("ANTHROPIC_API_KEY", "") 63 if "qwen_api_key" not in state:64 state["qwen_api_key"] = os.getenv("QWEN_API_KEY", "")65 if "ui_tars_url" not in state:66 state["ui_tars_url"] = ""67 68 # Set the initial api_key based on the provider69 if "planner_api_key" not in state:70 if state["planner_provider"] == "openai":71 state["planner_api_key"] = state["openai_api_key"]72 elif state["planner_provider"] == "anthropic":73 state["planner_api_key"] = state["anthropic_api_key"]74 elif state["planner_provider"] == "qwen":75 state["planner_api_key"] = state["qwen_api_key"]76 else:77 state["planner_api_key"] = ""78 79 logger.info(f"loaded initial api_key for {state['planner_provider']}: {state['planner_api_key']}")80 81 if not state["planner_api_key"]:82 logger.warning("Planner API key not found. Please set it in the environment or paste in textbox.")83 84 85 if "selected_screen" not in state:86 state['selected_screen'] = SELECTED_SCREEN_INDEX if SCREEN_NAMES else 087 88 if "auth_validated" not in state:89 state["auth_validated"] = False90 if "responses" not in state:91 state["responses"] = {}92 if "tools" not in state:93 state["tools"] = {}94 if "only_n_most_recent_images" not in state:95 state["only_n_most_recent_images"] = 10 # 1096 if "custom_system_prompt" not in state:97 state["custom_system_prompt"] = ""98 # remove if want to use default system prompt99 device_os_name = "Windows" if platform.system() == "Windows" else "Mac" if platform.system() == "Darwin" else "Linux"100 state["custom_system_prompt"] += f"\n\nNOTE: you are operating a {device_os_name} machine"101 if "hide_images" not in state:102 state["hide_images"] = False103 if 'chatbot_messages' not in state:104 state['chatbot_messages'] = []105 106 if "showui_config" not in state:107 state["showui_config"] = "Default"108 if "max_pixels" not in state:109 state["max_pixels"] = 1344110 if "awq_4bit" not in state:111 state["awq_4bit"] = False112 113 114async def main(state):115 """Render loop for Gradio"""116 setup_state(state)117 return "Setup completed"118 119 120def validate_auth(provider: APIProvider, api_key: str | None):121 if provider == APIProvider.ANTHROPIC:122 if not api_key:123 return "Enter your Anthropic API key to continue."124 if provider == APIProvider.BEDROCK:125 import boto3126 127 if not boto3.Session().get_credentials():128 return "You must have AWS credentials set up to use the Bedrock API."129 if provider == APIProvider.VERTEX:130 import google.auth131 from google.auth.exceptions import DefaultCredentialsError132 133 if not os.environ.get("CLOUD_ML_REGION"):134 return "Set the CLOUD_ML_REGION environment variable to use the Vertex API."135 try:136 google.auth.default(scopes=["https://www.googleapis.com/auth/cloud-platform"])137 except DefaultCredentialsError:138 return "Your google cloud credentials are not set up correctly."139 140 141def _api_response_callback(response: APIResponse[BetaMessage], response_state: dict):142 response_id = datetime.now().isoformat()143 response_state[response_id] = response144 145 146def _tool_output_callback(tool_output: ToolResult, tool_id: str, tool_state: dict):147 tool_state[tool_id] = tool_output148 149 150def chatbot_output_callback(message, chatbot_state, hide_images=False, sender="bot"):151 152 def _render_message(message: str | BetaTextBlock | BetaToolUseBlock | ToolResult, hide_images=False):153 154 logger.info(f"_render_message: {str(message)[:100]}")155 156 if isinstance(message, str):157 return message158 159 is_tool_result = not isinstance(message, str) and (160 isinstance(message, ToolResult)161 or message.__class__.__name__ == "ToolResult"162 or message.__class__.__name__ == "CLIResult"163 )164 if not message or (165 is_tool_result166 and hide_images167 and not hasattr(message, "error")168 and not hasattr(message, "output")169 ): # return None if hide_images is True170 return171 # render tool result172 if is_tool_result:173 message = cast(ToolResult, message)174 if message.output:175 return message.output176 if message.error:177 return f"Error: {message.error}"178 if message.base64_image and not hide_images:179 # somehow can't display via gr.Image180 # image_data = base64.b64decode(message.base64_image)181 # return gr.Image(value=Image.open(io.BytesIO(image_data)))182 return f'<img src="data:image/png;base64,{message.base64_image}">'183 184 elif isinstance(message, BetaTextBlock) or isinstance(message, TextBlock):185 return message.text186 elif isinstance(message, BetaToolUseBlock) or isinstance(message, ToolUseBlock):187 return f"Tool Use: {message.name}\nInput: {message.input}"188 else: 189 return message190 191 192 # processing Anthropic messages193 message = _render_message(message, hide_images)194 195 if sender == "bot":196 chatbot_state.append((None, message))197 else:198 chatbot_state.append((message, None))199 200 # Create a concise version of the chatbot state for logging201 concise_state = [(truncate_string(user_msg), truncate_string(bot_msg)) for user_msg, bot_msg in chatbot_state]202 logger.info(f"chatbot_output_callback chatbot_state: {concise_state} (truncated)")203 204 205def process_input(user_input, state):206 207 setup_state(state)208 209 # Append the user message to state["messages"]210 state["messages"].append(211 {212 "role": "user",213 "content": [TextBlock(type="text", text=user_input)],214 }215 )216 217 # Append the user's message to chatbot_messages with None for the assistant's reply218 state['chatbot_messages'].append((user_input, None))219 yield state['chatbot_messages'] # Yield to update the chatbot UI with the user's message220 221 # Run sampling_loop_sync with the chatbot_output_callback222 for loop_msg in sampling_loop_sync(223 system_prompt_suffix=state["custom_system_prompt"],224 planner_model=state["planner_model"],225 planner_provider=state["planner_provider"],226 actor_model=state["actor_model"],227 actor_provider=state["actor_provider"],228 messages=state["messages"],229 output_callback=partial(chatbot_output_callback, chatbot_state=state['chatbot_messages'], hide_images=state["hide_images"]),230 tool_output_callback=partial(_tool_output_callback, tool_state=state["tools"]),231 api_response_callback=partial(_api_response_callback, response_state=state["responses"]),232 api_key=state["planner_api_key"],233 only_n_most_recent_images=state["only_n_most_recent_images"],234 selected_screen=state['selected_screen'],235 showui_max_pixels=state['max_pixels'],236 showui_awq_4bit=state['awq_4bit']237 ): 238 if loop_msg is None:239 yield state['chatbot_messages']240 logger.info("End of task. Close the loop.")241 break242 243 244 yield state['chatbot_messages'] # Yield the updated chatbot_messages to update the chatbot UI245 246 247with gr.Blocks(theme=gr.themes.Soft()) as demo:248 249 state = gr.State({}) # Use Gradio's state management250 setup_state(state.value) # Initialize the state251 252 # Retrieve screen details253 gr.Markdown("# Computer Use OOTB")254 255 if not os.getenv("HIDE_WARNING", False):256 gr.Markdown(WARNING_TEXT)257 258 with gr.Accordion("Settings", open=True): 259 with gr.Row():260 with gr.Column():261 # --------------------------262 # Planner263 planner_model = gr.Dropdown(264 label="Planner Model",265 choices=["gpt-4o", 266 "gpt-4o-mini", 267 "qwen2-vl-max", 268 "qwen2-vl-2b (local)", 269 "qwen2-vl-7b (local)",270 "qwen2-vl-2b (ssh)", 271 "qwen2-vl-7b (ssh)",272 "qwen2.5-vl-7b (ssh)", 273 "claude-3-5-sonnet-20241022"],274 value="gpt-4o",275 interactive=True,276 )277 with gr.Column():278 planner_api_provider = gr.Dropdown(279 label="API Provider",280 choices=[option.value for option in APIProvider],281 value="openai",282 interactive=False,283 )284 with gr.Column():285 planner_api_key = gr.Textbox(286 label="Planner API Key",287 type="password",288 value=state.value.get("planner_api_key", ""),289 placeholder="Paste your planner model API key",290 interactive=True,291 )292 293 with gr.Column():294 actor_model = gr.Dropdown(295 label="Actor Model",296 choices=["ShowUI", "UI-TARS"],297 value="ShowUI",298 interactive=True,299 )300 301 with gr.Column():302 custom_prompt = gr.Textbox(303 label="System Prompt Suffix",304 value="",305 interactive=True,306 )307 with gr.Column():308 screen_options, primary_index = get_screen_details()309 SCREEN_NAMES = screen_options310 SELECTED_SCREEN_INDEX = primary_index311 screen_selector = gr.Dropdown(312 label="Select Screen",313 choices=screen_options,314 value=screen_options[primary_index] if screen_options else None,315 interactive=True,316 )317 with gr.Column():318 only_n_images = gr.Slider(319 label="N most recent screenshots",320 minimum=0,321 maximum=10,322 step=1,323 value=2,324 interactive=True,325 )326 327 with gr.Accordion("ShowUI Advanced Settings", open=False): 328 329 gr.Markdown("""330 **Note:** Adjust these settings to fine-tune the resource (**memory** and **infer time**) and performance trade-offs of ShowUI. \\331 Quantization model requires additional download. Please refer to [Computer Use OOTB - #ShowUI Advanced Settings guide](https://github.com/showlab/computer_use_ootb?tab=readme-ov-file#showui-advanced-settings) for preparation for this feature.332 """)333 334 # New configuration for ShowUI335 with gr.Row():336 with gr.Column():337 showui_config = gr.Dropdown(338 label="ShowUI Preset Configuration",339 choices=["Default (Maximum)", "Medium", "Minimal", "Custom"],340 value="Default (Maximum)",341 interactive=True,342 )343 with gr.Column():344 max_pixels = gr.Slider(345 label="Max Visual Tokens",346 minimum=720,347 maximum=1344,348 step=16,349 value=1344,350 interactive=False,351 )352 with gr.Column():353 awq_4bit = gr.Checkbox(354 label="Enable AWQ-4bit Model",355 value=False,356 interactive=False357 )358 359 # Define the merged dictionary with task mappings360 merged_dict = json.load(open("assets/examples/ootb_examples.json", "r"))361 362 def update_only_n_images(only_n_images_value, state):363 state["only_n_most_recent_images"] = only_n_images_value364 365 # Callback to update the second dropdown based on the first selection366 def update_second_menu(selected_category):367 return gr.update(choices=list(merged_dict.get(selected_category, {}).keys()))368 369 # Callback to update the third dropdown based on the second selection370 def update_third_menu(selected_category, selected_option):371 return gr.update(choices=list(merged_dict.get(selected_category, {}).get(selected_option, {}).keys()))372 373 # Callback to update the textbox based on the third selection374 def update_textbox(selected_category, selected_option, selected_task):375 task_data = merged_dict.get(selected_category, {}).get(selected_option, {}).get(selected_task, {})376 prompt = task_data.get("prompt", "")377 preview_image = task_data.get("initial_state", "")378 task_hint = "Task Hint: " + task_data.get("hint", "")379 return prompt, preview_image, task_hint380 381 # Function to update the global variable when the dropdown changes382 def update_selected_screen(selected_screen_name, state):383 global SCREEN_NAMES384 global SELECTED_SCREEN_INDEX385 SELECTED_SCREEN_INDEX = SCREEN_NAMES.index(selected_screen_name)386 logger.info(f"Selected screen updated to: {SELECTED_SCREEN_INDEX}")387 state['selected_screen'] = SELECTED_SCREEN_INDEX388 389 390 def update_planner_model(model_selection, state):391 state["model"] = model_selection392 # Update planner_model393 state["planner_model"] = model_selection394 logger.info(f"Model updated to: {state['planner_model']}")395 396 if model_selection == "qwen2-vl-max":397 provider_choices = ["qwen"]398 provider_value = "qwen"399 provider_interactive = False400 api_key_interactive = True401 api_key_placeholder = "qwen API key"402 actor_model_choices = ["ShowUI", "UI-TARS"]403 actor_model_value = "ShowUI"404 actor_model_interactive = True405 api_key_type = "password" # Display API key in password form406 407 elif model_selection == "qwen2-vl-2b (local)" or model_selection == "qwen2-vl-7b (local)":408 # Set provider to "openai", make it unchangeable409 provider_choices = ["local"]410 provider_value = "local"411 provider_interactive = False412 api_key_interactive = False413 api_key_placeholder = "not required"414 actor_model_choices = ["ShowUI", "UI-TARS"]415 actor_model_value = "ShowUI"416 actor_model_interactive = True417 api_key_type = "password" # Maintain consistency418 419 elif "ssh" in model_selection:420 provider_choices = ["ssh"]421 provider_value = "ssh"422 provider_interactive = False423 api_key_interactive = True424 api_key_placeholder = "ssh host and port (e.g. localhost:8000)"425 actor_model_choices = ["ShowUI", "UI-TARS"]426 actor_model_value = "ShowUI"427 actor_model_interactive = True428 api_key_type = "text" # Display SSH connection info in plain text429 # If SSH connection info already exists, keep it430 if "planner_api_key" in state and state["planner_api_key"]:431 state["api_key"] = state["planner_api_key"]432 else:433 state["api_key"] = ""434 435 elif model_selection == "gpt-4o" or model_selection == "gpt-4o-mini":436 # Set provider to "openai", make it unchangeable437 provider_choices = ["openai"]438 provider_value = "openai"439 provider_interactive = False440 api_key_interactive = True441 api_key_type = "password" # Display API key in password form442 443 api_key_placeholder = "openai API key"444 actor_model_choices = ["ShowUI", "UI-TARS"]445 actor_model_value = "ShowUI"446 actor_model_interactive = True447 448 elif model_selection == "claude-3-5-sonnet-20241022":449 # Provider can be any of the current choices except 'openai'450 provider_choices = [option.value for option in APIProvider if option.value != "openai"]451 provider_value = "anthropic" # Set default to 'anthropic'452 provider_interactive = True453 api_key_interactive = True454 api_key_placeholder = "claude API key"455 actor_model_choices = ["claude-3-5-sonnet-20241022"]456 actor_model_value = "claude-3-5-sonnet-20241022"457 actor_model_interactive = False458 api_key_type = "password" # Display API key in password form459 460 else:461 raise ValueError(f"Model {model_selection} not supported")462 463 # Update the provider in state464 state["planner_api_provider"] = provider_value465 466 # Update api_key in state based on the provider467 if provider_value == "openai":468 state["api_key"] = state.get("openai_api_key", "")469 elif provider_value == "anthropic":470 state["api_key"] = state.get("anthropic_api_key", "")471 elif provider_value == "qwen":472 state["api_key"] = state.get("qwen_api_key", "")473 elif provider_value == "local":474 state["api_key"] = ""475 # SSH的情况已经在上面处理过了，这里不需要重复处理476 477 provider_update = gr.update(478 choices=provider_choices,479 value=provider_value,480 interactive=provider_interactive481 )482 483 # Update the API Key textbox484 api_key_update = gr.update(485 placeholder=api_key_placeholder,486 value=state["api_key"],487 interactive=api_key_interactive,488 type=api_key_type # 添加 type 参数的更新489 )490 491 actor_model_update = gr.update(492 choices=actor_model_choices,493 value=actor_model_value,494 interactive=actor_model_interactive495 )496 497 logger.info(f"Updated state: model={state['planner_model']}, provider={state['planner_api_provider']}, api_key={state['api_key']}")498 return provider_update, api_key_update, actor_model_update499 500 def update_actor_model(actor_model_selection, state):501 state["actor_model"] = actor_model_selection502 logger.info(f"Actor model updated to: {state['actor_model']}")503 504 def update_api_key_placeholder(provider_value, model_selection):505 if model_selection == "claude-3-5-sonnet-20241022":506 507 if provider_value == "anthropic":508 return gr.update(placeholder="anthropic API key")509 elif provider_value == "bedrock":510 return gr.update(placeholder="bedrock API key")511 elif provider_value == "vertex":512 return gr.update(placeholder="vertex API key")513 else:514 return gr.update(placeholder="")515 elif model_selection == "gpt-4o + ShowUI":516 return gr.update(placeholder="openai API key")517 else:518 return gr.update(placeholder="")519 520 def update_system_prompt_suffix(system_prompt_suffix, state):521 state["custom_system_prompt"] = system_prompt_suffix522 523 # When showui_config changes, we set the max_pixels and awq_4bit accordingly.524 def handle_showui_config_change(showui_config_val, state):525 if showui_config_val == "Default (Maximum)":526 state["max_pixels"] = 1344527 state["awq_4bit"] = False528 return (529 gr.update(value=1344, interactive=False), 530 gr.update(value=False, interactive=False)531 )532 elif showui_config_val == "Medium":533 state["max_pixels"] = 1024534 state["awq_4bit"] = False535 return (536 gr.update(value=1024, interactive=False), 537 gr.update(value=False, interactive=False)538 )539 elif showui_config_val == "Minimal":540 state["max_pixels"] = 1024541 state["awq_4bit"] = True542 return (543 gr.update(value=1024, interactive=False), 544 gr.update(value=True, interactive=False)545 )546 elif showui_config_val == "Custom":547 # Do not overwrite the current user values, just make them interactive548 return (549 gr.update(interactive=True), 550 gr.update(interactive=True)551 )552 553 def update_api_key(api_key_value, state):554 """Handle API key updates"""555 state["planner_api_key"] = api_key_value556 if state["planner_provider"] == "ssh":557 state["api_key"] = api_key_value558 logger.info(f"API key updated: provider={state['planner_provider']}, api_key={state['api_key']}")559 560 with gr.Accordion("Quick Start Prompt", open=False): # open=False 表示默认收561 # Initialize Gradio interface with the dropdowns562 with gr.Row():563 # Set initial values564 initial_category = "Game Play"565 initial_second_options = list(merged_dict[initial_category].keys())566 initial_third_options = list(merged_dict[initial_category][initial_second_options[0]].keys())567 initial_text_value = merged_dict[initial_category][initial_second_options[0]][initial_third_options[0]]568 569 with gr.Column(scale=2):570 # First dropdown for Task Category571 first_menu = gr.Dropdown(572 choices=list(merged_dict.keys()), label="Task Category", interactive=True, value=initial_category573 )574 575 # Second dropdown for Software576 second_menu = gr.Dropdown(577 choices=initial_second_options, label="Software", interactive=True, value=initial_second_options[0]578 )579 580 # Third dropdown for Task581 third_menu = gr.Dropdown(582 choices=initial_third_options, label="Task", interactive=True, value=initial_third_options[0]583 # choices=["Please select a task"]+initial_third_options, label="Task", interactive=True, value="Please select a task"584 )585 586 with gr.Column(scale=1):587 initial_image_value = "./assets/examples/init_states/honkai_star_rail_showui.png" # default image path588 image_preview = gr.Image(value=initial_image_value, label="Reference Initial State", height=260-(318.75-280))589 hintbox = gr.Markdown("Task Hint: Selected options will appear here.")590 591 # Textbox for displaying the mapped value592 # textbox = gr.Textbox(value=initial_text_value, label="Action")593 594 # api_key.change(fn=lambda key: save_to_storage(API_KEY_FILE, key), inputs=api_key)595 596 with gr.Row():597 # submit_button = gr.Button("Submit") # Add submit button598 with gr.Column(scale=8):599 chat_input = gr.Textbox(show_label=False, placeholder="Type a message to send to Computer Use OOTB...", container=False)600 with gr.Column(scale=1, min_width=50):601 submit_button = gr.Button(value="Send", variant="primary")602 603 chatbot = gr.Chatbot(label="Chatbot History", type="tuples", autoscroll=True, height=580)604 605 planner_model.change(fn=update_planner_model, inputs=[planner_model, state], outputs=[planner_api_provider, planner_api_key, actor_model])606 planner_api_provider.change(fn=update_api_key_placeholder, inputs=[planner_api_provider, planner_model], outputs=planner_api_key)607 actor_model.change(fn=update_actor_model, inputs=[actor_model, state], outputs=None)608 609 screen_selector.change(fn=update_selected_screen, inputs=[screen_selector, state], outputs=None)610 only_n_images.change(fn=update_only_n_images, inputs=[only_n_images, state], outputs=None)611 612 # When showui_config changes, we update max_pixels and awq_4bit automatically.613 showui_config.change(fn=handle_showui_config_change, 614 inputs=[showui_config, state], 615 outputs=[max_pixels, awq_4bit])616 617 # Link callbacks to update dropdowns based on selections618 first_menu.change(fn=update_second_menu, inputs=first_menu, outputs=second_menu)619 second_menu.change(fn=update_third_menu, inputs=[first_menu, second_menu], outputs=third_menu)620 third_menu.change(fn=update_textbox, inputs=[first_menu, second_menu, third_menu], outputs=[chat_input, image_preview, hintbox])621 622 # chat_input.submit(process_input, [chat_input, state], chatbot)623 submit_button.click(process_input, [chat_input, state], chatbot)624 625 planner_api_key.change(626 fn=update_api_key,627 inputs=[planner_api_key, state],628 outputs=None629 )630 631demo.launch(share=True,632 allowed_paths=["./"],633 server_port=7888) # TODO: allowed_paths634