app.py

import os
import base64
import gradio as gr
from typing import List, Tuple, Optional, Dict, Any
from enum import Enum
from datetime import datetime
import httpx
from dataclasses import dataclass
from anthropic import Anthropic
import json

# Get API key from environment variable
ANTHROPIC_API_KEY = os.getenv('ANTHROPIC_API_KEY')
if not ANTHROPIC_API_KEY:
    raise ValueError("ANTHROPIC_API_KEY environment variable must be set")

# Initialize Anthropic client
anthropic = Anthropic(api_key=ANTHROPIC_API_KEY)

@dataclass
class ToolResult:
    output: Optional[str] = None
    error: Optional[str] = None
    base64_image: Optional[str] = None

    def to_dict(self) -> Dict[str, Any]:
        return {
            "type": "tool_result",
            "output": self.output,
            "error": self.error,
            "base64_image": self.base64_image
        }

# Define the computer use tools
COMPUTER_USE_TOOLS = [
    {
        "type": "computer_20241022",
        "name": "computer",
        "display_width_px": 1024,
        "display_height_px": 768,
        "display_number": 1,
    },
    {
        "type": "text_editor_20241022",
        "name": "str_replace_editor"
    },
    {
        "type": "bash_20241022",
        "name": "bash"
    }
]

class ComputerUseDemo:
    def __init__(self):
        self.messages: List[Dict[str, Any]] = []
        self.model = "claude-3-5-sonnet-20241022"
        self.max_tokens = 4096
        
    def validate_auth(self) -> Optional[str]:
        if not ANTHROPIC_API_KEY:
            return "ANTHROPIC_API_KEY environment variable is not set."
        return None

    def format_messages(self, chat_history: List[Tuple[str, str]]) -> List[Dict[str, Any]]:
        """Convert chat history to Anthropic message format"""
        formatted_messages = []
        for user_msg, assistant_msg in chat_history:
            formatted_messages.append({"role": "user", "content": user_msg})
            if assistant_msg:
                formatted_messages.append({"role": "assistant", "content": assistant_msg})
        return formatted_messages

    def execute_tool(self, tool_name: str, tool_input: str) -> ToolResult:
        """Execute the requested tool and return results"""
        try:
            if tool_name == "computer":
                return ToolResult(
                    output=f"Executed computer command: {tool_input}",
                    base64_image=None
                )
            elif tool_name == "str_replace_editor":
                return ToolResult(
                    output=f"Executed text editor command: {tool_input}"
                )
            elif tool_name == "bash":
                return ToolResult(
                    output=f"Executed bash command: {tool_input}"
                )
            else:
                return ToolResult(
                    error=f"Unknown tool: {tool_name}"
                )
        except Exception as e:
            return ToolResult(error=str(e))

    def process_message(
        self, 
        message: str, 
        chat_history: List[Tuple[str, str]],
        system_prompt: str
    ) -> Tuple[List[Tuple[str, str]], Optional[str]]:
        """Process a message with proper error handling"""
        if error := self.validate_auth():
            return chat_history, error

        try:
            # Format messages for Claude
            messages = self.format_messages(chat_history)
            messages.append({"role": "user", "content": message})

            # Create system prompt for computer use
            base_system_prompt = """You have access to a computing environment through specialized tools.
            Use the computer tool to interact with the GUI, the text editor for file operations, 
            and bash for command-line tasks. Take screenshots after significant actions and 
            verify the results carefully. Always handle errors gracefully and inform the user 
            of any issues."""
            
            full_system_prompt = f"{base_system_prompt}\n{system_prompt}" if system_prompt else base_system_prompt

            while True:
                # Call Claude API with computer use enabled
                response = anthropic.beta.messages.create(
                    model=self.model,
                    max_tokens=self.max_tokens,
                    messages=messages,
                    system=full_system_prompt,
                    tools=COMPUTER_USE_TOOLS,
                    betas=["computer-use-2024-10-22"]
                )

                # Handle tool calls
                if response.stop_reason == "tool_use" and response.tool_calls:
                    tool_call = response.tool_calls[0]
                    tool_result = self.execute_tool(tool_call.name, tool_call.arguments)
                    
                    # Add tool interactions to message history
                    messages.append({
                        "role": "assistant",
                        "content": [{"type": "tool_use", "name": tool_call.name, "input": tool_call.arguments}]
                    })
                    messages.append({
                        "role": "user",
                        "content": [tool_result.to_dict()]
                    })
                else:
                    # Final response received
                    assistant_response = response.content[0].text
                    chat_history.append((message, assistant_response))
                    break

            return chat_history, None
            
        except Exception as e:
            return chat_history, f"Error processing message: {str(e)}"

def create_demo() -> gr.Interface:
    demo = ComputerUseDemo()
    
    with gr.Blocks(title="Claude Computer Use Demo") as interface:
        gr.Markdown("# Claude Computer Use Demo")
        gr.Markdown("""## Security Notice
        ⚠️ This demo runs in a sandboxed environment. Never provide access to sensitive 
        information or credentials, as Claude's behavior could be influenced by malicious content.""")
        
        with gr.Row():
            with gr.Column(scale=4):
                chatbot = gr.Chatbot(
                    label="Chat History",
                    height=600,
                    container=True
                )
                with gr.Row():
                    msg = gr.Textbox(
                        label="Message",
                        placeholder="Type your instructions for Claude...",
                        lines=3,
                        scale=8
                    )
                    with gr.Column(scale=1):
                        submit = gr.Button("Send", variant="primary")
                        clear = gr.Button("Reset")
                
            with gr.Column(scale=1):
                system_prompt = gr.Textbox(
                    label="System Instructions",
                    lines=3,
                    placeholder="Add custom instructions for Claude..."
                )
                with gr.Accordion("Advanced Settings", open=False):
                    model = gr.Textbox(
                        label="Model",
                        value=demo.model,
                        interactive=False
                    )
                    max_tokens = gr.Slider(
                        label="Max Tokens",
                        minimum=1024,
                        maximum=8192,
                        value=demo.max_tokens,
                        step=1024
                    )

        def on_message(
            message: str,
            chat_history: List[Tuple[str, str]],
            system_prompt: str
        ) -> Tuple[List[Tuple[str, str]], str]:
            if not message.strip():
                return chat_history, ""
            new_history, error = demo.process_message(message, chat_history or [], system_prompt)
            if error:
                gr.Warning(error)
            return new_history, ""

        # Set up event handlers
        msg.submit(
            on_message,
            inputs=[msg, chatbot, system_prompt],
            outputs=[chatbot, msg]
        )
        
        submit.click(
            on_message,
            inputs=[msg, chatbot, system_prompt],
            outputs=[chatbot, msg]
        )

        clear.click(
            lambda: ([], ""),
            outputs=[chatbot, msg]
        )

    return interface

# Create and launch the demo
demo = create_demo()

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)