baqr/computer_use_ootb

1import asyncio2from typing import Any, Dict, cast3from collections.abc import Callable4from anthropic.types.beta import (5    BetaContentBlock,6    BetaContentBlockParam,7    BetaImageBlockParam,8    BetaMessage,9    BetaMessageParam,10    BetaTextBlockParam,11    BetaToolResultBlockParam,12)13from anthropic.types import TextBlock14from anthropic.types.beta import BetaMessage, BetaTextBlock, BetaToolUseBlock15from ..tools import BashTool, ComputerTool, EditTool, ToolCollection, ToolResult16 17 18class AnthropicExecutor:19    def __init__(20        self, 21        output_callback: Callable[[BetaContentBlockParam], None], 22        tool_output_callback: Callable[[Any, str], None],23        selected_screen: int = 024    ):25        self.tool_collection = ToolCollection(26            ComputerTool(selected_screen=selected_screen),27            BashTool(),28            EditTool(),29        )30        self.output_callback = output_callback31        self.tool_output_callback = tool_output_callback32 33    def __call__(self, response: BetaMessage, messages: list[BetaMessageParam]):34        new_message = {35            "role": "assistant",36            "content": cast(list[BetaContentBlockParam], response.content),37        }38        if new_message not in messages:39            messages.append(new_message)40        else:41            print("new_message already in messages, there are duplicates.")42        43        tool_result_content: list[BetaToolResultBlockParam] = []44        for content_block in cast(list[BetaContentBlock], response.content):45            46            self.output_callback(content_block, sender="bot")47            # Execute the tool48            if content_block.type == "tool_use":49                # Run the asynchronous tool execution in a synchronous context50                result = asyncio.run(self.tool_collection.run(51                    name=content_block.name,52                    tool_input=cast(dict[str, Any], content_block.input),53                ))54                55                self.output_callback(result, sender="bot")56                57                tool_result_content.append(58                    _make_api_tool_result(result, content_block.id)59                )60                self.tool_output_callback(result, content_block.id)61 62            # Craft messages based on the content_block63            # Note: to display the messages in the gradio, you should organize the messages in the following way (user message, bot message)64            65            display_messages = _message_display_callback(messages)66            # display_messages = []67            68            # Send the messages to the gradio69            for user_msg, bot_msg in display_messages:70                yield [user_msg, bot_msg], tool_result_content71 72        if not tool_result_content:73            return messages74        75        return tool_result_content76 77def _message_display_callback(messages):78    display_messages = []79    for msg in messages:80        try:81            if isinstance(msg["content"][0], TextBlock):82                display_messages.append((msg["content"][0].text, None))  # User message83            elif isinstance(msg["content"][0], BetaTextBlock):84                display_messages.append((None, msg["content"][0].text))  # Bot message85            elif isinstance(msg["content"][0], BetaToolUseBlock):86                display_messages.append((None, f"Tool Use: {msg['content'][0].name}\nInput: {msg['content'][0].input}"))  # Bot message87            elif isinstance(msg["content"][0], Dict) and msg["content"][0]["content"][-1]["type"] == "image":88                display_messages.append((None, f'<img src="data:image/png;base64,{msg["content"][0]["content"][-1]["source"]["data"]}">'))  # Bot message89            else:90                print(msg["content"][0])91        except Exception as e:92            print("error", e)93            pass94    return display_messages95 96def _make_api_tool_result(97    result: ToolResult, tool_use_id: str98) -> BetaToolResultBlockParam:99    """Convert an agent ToolResult to an API ToolResultBlockParam."""100    tool_result_content: list[BetaTextBlockParam | BetaImageBlockParam] | str = []101    is_error = False102    if result.error:103        is_error = True104        tool_result_content = _maybe_prepend_system_tool_result(result, result.error)105    else:106        if result.output:107            tool_result_content.append(108                {109                    "type": "text",110                    "text": _maybe_prepend_system_tool_result(result, result.output),111                }112            )113        if result.base64_image:114            tool_result_content.append(115                {116                    "type": "image",117                    "source": {118                        "type": "base64",119                        "media_type": "image/png",120                        "data": result.base64_image,121                    },122                }123            )124    return {125        "type": "tool_result",126        "content": tool_result_content,127        "tool_use_id": tool_use_id,128        "is_error": is_error,129    }130 131 132def _maybe_prepend_system_tool_result(result: ToolResult, result_text: str):133    if result.system:134        result_text = f"<system>{result.system}</system>\n{result_text}"135    return result_text