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"Created/modified file with content: {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. After each action, explicitly describe what you did and what happened.
            Always show the results of your actions."""
            
            full_system_prompt = f"{base_system_prompt}\n{system_prompt}" if system_prompt else base_system_prompt

            conversation_log = []
            conversation_log.append((message, ""))  # Add user message

            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"]
                )

                # Get the assistant's response text
                assistant_message = response.content[0].text if response.content else ""
                
                # Check for tool calls in the response
                tool_calls = []
                for content in response.content:
                    if hasattr(content, 'tool_calls') and content.tool_calls:
                        tool_calls.extend(content.tool_calls)

                # Add assistant's initial response to conversation
                if assistant_message:
                    conversation_log.append(("", assistant_message))

                # Handle tool calls if present
                if tool_calls:
                    tool_call = tool_calls[0]  # Handle first tool call
                    tool_result = self.execute_tool(tool_call.name, tool_call.parameters)
                    
                    # Add tool execution to conversation
                    if tool_result.error:
                        conversation_log.append(("", f"Error: {tool_result.error}"))
                    else:
                        conversation_log.append(("", f"Tool execution: {tool_result.output}"))

                    # Add tool interactions to message history
                    messages.append({
                        "role": "assistant",
                        "content": [{"type": "tool_use", "name": tool_call.name, "input": tool_call.parameters}]
                    })
                    messages.append({
                        "role": "user",
                        "content": [tool_result.to_dict()]
                    })
                else:
                    break

            return conversation_log, 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)