EvoSynth/jailbreak_toolbox/models/implementations/multithreaded_openai_model.py

# jailbreak_toolbox/models/implementations/multithreaded_openai_model.py
from typing import Union, List, Dict, Any, Optional
import time
import base64
import os
import io
from PIL import Image
from ..multithreaded_model import MultiThreadedModel
from ...core.registry import model_registry

@model_registry.register("multithreaded_openai")
class MultiThreadedOpenAIModel(MultiThreadedModel):
    """
    Multi-threaded implementation of OpenAI API client.
    
    Supports concurrent API calls with rate limiting and retries.
    """
    
    def __init__(
        self,
        api_key: str,
        model_name: str = "gpt-4",
        base_url: Optional[str] = None,
        system_message: str = "You are a helpful assistant.",
        temperature: float = 0.7,
        max_tokens: int = 1024,
        max_workers: int = 5,
        requests_per_minute: int = 60,
        retry_attempts: int = 3,
        retry_delay: float = 1.0,
        **kwargs
    ):
        """
        Initialize the multi-threaded OpenAI model.
        
        Args:
            api_key: OpenAI API key
            model_name: Model name/identifier
            base_url: Optional base URL for API requests
            system_message: System message for conversation
            temperature: Temperature parameter for generation
            max_tokens: Maximum tokens to generate
            max_workers: Maximum number of worker threads
            requests_per_minute: Maximum number of requests per minute
            retry_attempts: Number of retry attempts on failure
            retry_delay: Delay between retries in seconds
        """
        super().__init__(
            max_workers=max_workers,
            requests_per_minute=requests_per_minute,
            retry_attempts=retry_attempts,
            retry_delay=retry_delay,
            **kwargs
        )
        
        self.api_key = api_key
        self.model_name = model_name
        self.base_url = base_url
        self.system_message = system_message
        self.temperature = temperature
        self.max_tokens = max_tokens
        
        # Initialize OpenAI client
        try:
            from openai import OpenAI
            self.client = OpenAI(api_key=self.api_key, base_url=self.base_url)
            self.conversation_history = [{"role": "system", "content": self.system_message}]
        except ImportError:
            raise ImportError("OpenAI package is required. Install it using: pip install openai")
    
    def query(self, text_input: Union[str, List[Dict]] = "", image_input: Any = None, maintain_history: bool = False) -> str:
        """
        Send a query to the OpenAI API and return the response.
        
        Args:
            text_input: The prompt text to send (can be string or list of message dicts)
            image_input: Path to image file or PIL Image object for vision models
            maintain_history: Whether to add this exchange to conversation history
            
        Returns:
            The model's response as a string
        """
        messages = []
        
        # Handle image input
        if image_input is not None:
            try:
                image_base64 = self._encode_image_to_base64(image_input)
                if isinstance(text_input, list):
                    # If text_input is already a list of messages, use it directly
                    messages = text_input
                elif isinstance(text_input, str):
                    # Create a message with both text and image
                    messages = [{"role": "system", "content": self.system_message}] if not maintain_history else []
                    messages.append({
                        "role": "user",
                        "content": [
                            {
                                "type": "text",
                                "text": text_input
                            },
                            {
                                "type": "image_url",
                                "image_url": {
                                    "url": f"data:image/jpeg;base64,{image_base64}"
                                },
                            },
                        ],
                    })
            except Exception as e:
                print(f"Warning: Failed to process image input: {str(e)}")
                # Fall back to text-only input
                image_input = None
        
        # Handle text-only input or fallback from image processing failure
        if image_input is None:
            if isinstance(text_input, list):
                # If text_input is already a list of messages, use it directly
                messages = text_input
            elif isinstance(text_input, str):
                # Add user message to history if maintaining history
                if maintain_history:
                    self.add_user_message(text_input)
                    messages = self.conversation_history
                else:
                    # For single-turn interactions without affecting history
                    messages = [{"role": "system", "content": self.system_message},
                                {"role": "user", "content": text_input}]
        
        # Make API call with retries built into MultiThreadedModel base class
        response = self.client.chat.completions.create(
            model=self.model_name,
            messages=messages,
            temperature=self.temperature,
            max_tokens=self.max_tokens
        )
        
        response_text = response.choices[0].message.content
        
        # Add assistant response to history if maintaining history
        if maintain_history:
            self.add_assistant_message(response_text)
        
        return response_text
    
    def _encode_image_to_base64(self, image_input) -> str:
        """
        Convert image input (file path or PIL Image) to base64 encoding.
        
        Args:
            image_input: Path to image file or PIL Image object
            
        Returns:
            Base64 encoded image string
        """
        # If input is a string, treat it as a file path
        if isinstance(image_input, str):
            if not os.path.exists(image_input):
                raise FileNotFoundError(f"Image file not found: {image_input}")
            
            with open(image_input, "rb") as image_file:
                return base64.b64encode(image_file.read()).decode('utf-8')
                
        # If input is a PIL Image
        elif isinstance(image_input, Image.Image):
            buffer = io.BytesIO()
            image_input.save(buffer, format="JPEG")
            return base64.b64encode(buffer.getvalue()).decode('utf-8')
            
        # If input is bytes
        elif isinstance(image_input, bytes):
            return base64.b64encode(image_input).decode('utf-8')
            
        else:
            raise TypeError(f"Unsupported image input type: {type(image_input)}")
    
    def add_user_message(self, message: Union[str, List]) -> None:
        """
        Add a user message to the conversation history.
        
        Args:
            message: The message content (string or list of content dicts)
        """
        self.conversation_history.append({"role": "user", "content": message})
    
    def add_assistant_message(self, message: str) -> None:
        """
        Add an assistant message to the conversation history.
        
        Args:
            message: The message content
        """
        self.conversation_history.append({"role": "assistant", "content": message})
    
    def get_conversation_history(self) -> List[Dict]:
        """
        Get the current conversation history.
        
        Returns:
            List of message dictionaries
        """
        return self.conversation_history
    
    def reset_conversation(self) -> None:
        """Reset the conversation history to only include the system message."""
        self.conversation_history = [{"role": "system", "content": self.system_message}]