first commit

jailbreak_toolbox/models/__init__.py

@@ -0,0 +1,2 @@
+# FILE: /jailbreak-toolbox/jailbreak-toolbox/jailbreak_toolbox/models/__init__.py
+# This file is intentionally left blank.

jailbreak_toolbox/models/base_image_generator.py

@@ -0,0 +1,36 @@
+from abc import ABC, abstractmethod
+from typing import Optional, Union, Any
+from PIL import Image
+
+class BaseImageGenerator(ABC):
+    """
+    Abstract base class for image generation models.
+    Defines the interface that all image generation models must implement.
+    """
+    def __init__(self, **kwargs):
+        """Initialize with model-specific parameters."""
+        pass
+
+    @abstractmethod
+    def generate(self, prompt: Any, output_path: Optional[str] = None) -> Image.Image:
+        """
+        Generate an image from the given prompt.
+        
+        Args:
+            prompt: Input for image generation (can be text, image, or other data)
+            output_path: Optional path to save the generated image
+            
+        Returns:
+            PIL.Image: Generated image
+        """
+        pass
+
+    def save_image(self, image: Image.Image, output_path: str) -> None:
+        """
+        Helper method to save the generated image.
+        
+        Args:
+            image: PIL Image to save
+            output_path: Path where to save the image
+        """
+        image.save(output_path)

jailbreak_toolbox/models/base_model.py

@@ -0,0 +1,34 @@
+from abc import ABC, abstractmethod
+from typing import Any, Dict
+
+class BaseModel(ABC):
+    """
+    所有模型的抽象基类。
+    定义了所有模型（无论是本地白盒还是远程API黑盒）都必须遵守的接口。
+    """
+    def __init__(self, **kwargs):
+        """允许在配置文件中传入任意模型特定的参数。"""
+        pass
+
+    @abstractmethod
+    def query(self, text_input: str, image_input: Any = None) -> str:
+        """
+        向模型发送查询并获取响应的核心方法。
+        对于纯文本模型，image_input 将被忽略。
+        这是所有模型都必须实现的功能。
+        """
+        pass
+
+    def get_gradients(self, inputs) -> Dict:
+        """
+        (可选) 获取梯度，用于白盒攻击。
+        如果模型不支持（如黑盒API模型），则直接抛出异常。
+        """
+        raise NotImplementedError(f"{self.__class__.__name__} does not support gradient access.")
+
+    def get_embeddings(self, inputs) -> Any:
+        """
+        (可选) 获取嵌入向量，用于白盒或灰盒攻击。
+        如果模型不支持，则直接抛出异常。
+        """
+        raise NotImplementedError(f"{self.__class__.__name__} does not support embedding access.")

jailbreak_toolbox/models/implementations/__init__.py

@@ -0,0 +1 @@
+from . import mock_model, openai_model, huggingface_model

jailbreak_toolbox/models/implementations/multithreaded_openai_model.py

@@ -0,0 +1,206 @@
+# 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}]

jailbreak_toolbox/models/implementations/openai_model.py

@@ -0,0 +1,286 @@
+from ..base_model import BaseModel
+from ...core.registry import model_registry
+import openai
+import time
+import torch
+import base64
+from io import BytesIO
+from typing import Any, Optional, List, Dict, Union
+from PIL import Image
+
+@model_registry.register("openai")
+class OpenAIModel(BaseModel):
+    """
+    OpenAI API model wrapper for jailbreak toolbox.
+    Supports various OpenAI models like gpt-3.5-turbo and gpt-4.
+    """
+    def __init__(self,
+                 api_key: str,
+                 base_url: Optional[str] = None,
+                 model_name: str = "gpt-3.5-turbo",
+                 temperature: float = 0.7,
+                 max_tokens: int = None,
+                 retry_attempts: int = 3,
+                 retry_delay: float = 2.0,
+                 system_message: str = "You are a helpful assistant.",
+                 embedding_model: str = "text-embedding-3-small",
+                 **kwargs):
+        """
+        Initialize the OpenAI model.
+        
+        Args:
+            api_key: OpenAI API key
+            base_url: Optional base URL for OpenAI API
+            model_name: Name of the OpenAI model to use (e.g., "gpt-3.5-turbo", "gpt-4")
+            temperature: Sampling temperature (0.0-2.0)
+            max_tokens: Maximum tokens in the response
+            retry_attempts: Number of retry attempts for API calls
+            retry_delay: Delay between retry attempts in seconds
+            system_message: System message to set the assistant's behavior
+            embedding_model: Model to use for generating embeddings
+        """
+        super().__init__(**kwargs)
+        self.model_name = model_name
+        self.temperature = temperature
+        self.max_tokens = max_tokens
+        self.retry_attempts = retry_attempts
+        self.retry_delay = retry_delay
+        self.system_message = system_message
+        self.embedding_model = embedding_model
+
+        # Initialize conversation history
+        self.conversation_history = [{"role": "system", "content": system_message}]
+
+        # Filter kwargs for different OpenAI API functions
+        self.client_kwargs = self._filter_client_kwargs(kwargs)
+        self.chat_kwargs = self._filter_chat_kwargs(kwargs)
+        self.embedding_kwargs = self._filter_embedding_kwargs(kwargs)
+
+        # Configure OpenAI client
+        openai.api_key = api_key
+        if base_url:
+            openai.base_url = base_url
+        self.client = openai.OpenAI(api_key=api_key, base_url=base_url, **self.client_kwargs)
+        print(f"Initialized OpenAI model: {model_name}")
+
+    def _filter_client_kwargs(self, kwargs: Dict[str, Any]) -> Dict[str, Any]:
+        """Filter kwargs suitable for OpenAI client initialization"""
+        # Valid parameters for OpenAI() client
+        valid_client_params = {
+            'timeout', 'max_retries', 'default_headers', 'default_query',
+            'http_client', 'api_key', 'base_url', 'organization', 'project'
+        }
+        return {k: v for k, v in kwargs.items() if k in valid_client_params}
+
+    def _filter_chat_kwargs(self, kwargs: Dict[str, Any]) -> Dict[str, Any]:
+        """Filter kwargs suitable for chat.completions.create()"""
+        # Valid parameters for chat completions
+        valid_chat_params = {
+            'frequency_penalty', 'logit_bias', 'logprobs', 'top_logprobs',
+            'max_tokens', 'n', 'presence_penalty', 'response_format',
+            'seed', 'stop', 'stream', 'temperature', 'top_p', 'tools', 'tool_choice',
+            'parallel_tool_calls', 'user'
+        }
+        return {k: v for k, v in kwargs.items() if k in valid_chat_params}
+
+    def _filter_embedding_kwargs(self, kwargs: Dict[str, Any]) -> Dict[str, Any]:
+        """Filter kwargs suitable for embeddings.create()"""
+        # Valid parameters for embeddings
+        valid_embedding_params = {
+            'encoding_format', 'dimensions', 'user'
+        }
+        return {k: v for k, v in kwargs.items() if k in valid_embedding_params}
+
+    def _encode_image_to_base64(self, image_input: Union[str, Any]) -> str:
+        """Encode an image to base64 string. Supports both file paths and PIL Image objects.
+        
+        Args:
+            image_input: Path to image file or PIL Image object
+            
+        Returns:
+            Base64 encoded string of the image
+        """
+        # Check if it's a file path (string)
+        if isinstance(image_input, str):
+            with open(image_input, "rb") as image_file:
+                return base64.b64encode(image_file.read()).decode('utf-8')
+        
+        # Check if it's a PIL Image object
+        if isinstance(image_input, Image.Image):
+            buffered = BytesIO()
+            image_input.save(buffered, format="JPEG")
+            return base64.b64encode(buffered.getvalue()).decode('utf-8')
+        
+        raise ValueError("image_input must be either a file path (string) or a PIL Image object")
+    
+    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
+        """
+        #print("text input: ",text_input)
+        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 text_input is not None:
+                text_input = str(text_input)
+                # 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}]
+            #print(type(text_input))
+        
+        for attempt in range(self.retry_attempts):
+            try:
+                #print(messages)
+                # Prepare chat completion parameters
+                chat_params = {
+                    'model': self.model_name,
+                    'messages': messages,
+                    **self.chat_kwargs
+                }
+
+                # Add explicit parameters if not in kwargs
+                if self.temperature is not None:
+                    chat_params['temperature'] = self.temperature
+                if self.max_tokens is not None:
+                    chat_params['max_tokens'] = self.max_tokens
+
+                response = self.client.chat.completions.create(**chat_params)
+                
+                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
+            
+            except Exception as e:
+                print(f"Unexpected error: {str(e)}")
+                if attempt == self.retry_attempts - 1:
+                    return f"Error: {str(e)}"
+                # Note: time.sleep removed to prevent blocking in asyncio environments
+                # For retry delays, the calling async code should handle timing
+        
+        return "Error: Failed to get response from model"
+    
+    def add_user_message(self, content: Union[str, List[Dict]]) -> None:
+        """Add a user message to the conversation history."""
+        self.conversation_history.append({"role": "user", "content": content})
+    
+    def add_assistant_message(self, content: Union[str, List[Dict]]) -> None:
+        """Add an assistant message to the conversation history."""
+        self.conversation_history.append({"role": "assistant", "content": content})
+    
+    def add_system_message(self, content: str) -> None:
+        """Add a system message to the conversation history."""
+        self.conversation_history.append({"role": "system", "content": content})
+        
+    def remove_last_turn(self) -> None:
+        """
+        Remove the last turn of conversation (last user message and its corresponding assistant message).
+        Assumes conversation is stored sequentially as messages.
+        """
+        if not self.conversation_history:
+            return
+
+        for idx in range(len(self.conversation_history) - 1, -1, -1):
+            if self.conversation_history[idx]["role"] == "user":
+                self.conversation_history = self.conversation_history[:idx]
+                break
+    
+    def reset_conversation(self) -> None:
+        """Reset the conversation history to only include the initial system message."""
+        self.conversation_history = [{"role": "system", "content": self.system_message}]
+    
+    def get_conversation_history(self) -> List[Dict[str, str]]:
+        """Get the current conversation history."""
+        return self.conversation_history
+    
+    def set_system_message(self, system_message: str) -> None:
+        """Set a new system message and reset the conversation."""
+        self.system_message = system_message
+        self.reset_conversation()
+        
+    def reset_system_message(self) -> None:
+        """Reset the system message to the default."""
+        self.system_message = "You are a helpful assistant."
+        self.reset_conversation()
+        
+    def get_embedding(self, text_input: str, model: str = None) -> list[float]:
+        """
+        Get embedding vector for the given text using OpenAI embedding model.
+
+        Args:
+            text_input: The input text to embed
+            model: The embedding model to use (default: uses self.embedding_model)
+
+        Returns:
+            A list of floats representing the embedding vector
+        """
+        try:
+            clean_text = text_input.replace("\n", " ")
+
+            # Use specified model or default embedding model
+            embedding_model = model or self.embedding_model
+
+            # Prepare embedding parameters
+            embedding_params = {
+                'input': clean_text,
+                'model': embedding_model,
+                **self.embedding_kwargs
+            }
+
+            response = self.client.embeddings.create(**embedding_params)
+            embedding = response.data[0].embedding
+            return torch.tensor(embedding)
+
+        except Exception as e:
+            print(f"Error while generating embedding: {str(e)}")
+            return []

jailbreak_toolbox/models/multithreaded_model.py

@@ -0,0 +1,136 @@
+# jailbreak_toolbox/models/implementations/multithreaded_model.py
+import threading
+import queue
+import time
+from typing import List, Dict, Any, Optional, Callable
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from .base_model import BaseModel
+from ..core.registry import model_registry
+
+@model_registry.register("multithreaded_model")
+class MultiThreadedModel(BaseModel):
+    """
+    Base class for models that support multi-threaded API calls.
+    
+    This provides an efficient way to make multiple API calls in parallel,
+    with rate limiting to avoid API throttling.
+    """
+    
+    def __init__(
+        self,
+        max_workers: int = 5,
+        requests_per_minute: int = 60,
+        retry_attempts: int = 3,
+        retry_delay: float = 1.0,
+        **kwargs
+    ):
+        """
+        Initialize the multi-threaded model.
+        
+        Args:
+            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__(**kwargs)
+        self.max_workers = max_workers
+        self.requests_per_minute = requests_per_minute
+        self.retry_attempts = retry_attempts
+        self.retry_delay = retry_delay
+        
+        # Rate limiting
+        self.request_interval = 60.0 / requests_per_minute
+        self.last_request_time = 0
+        self.request_lock = threading.Lock()
+    
+    def query_batch(
+        self, 
+        inputs: List[Dict[str, Any]], 
+        callback: Optional[Callable[[int, str], None]] = None
+    ) -> List[str]:
+        """
+        Send multiple queries in parallel using a thread pool.
+        
+        Args:
+            inputs: List of input dictionaries, each containing:
+                - 'text': Text input for the model
+                - 'image': Optional image input (path or PIL Image)
+                - Other model-specific parameters
+            callback: Optional callback function to call when each result is ready
+                     Function signature: callback(index, response)
+                     
+        Returns:
+            List of model responses in the same order as inputs
+        """
+        results = [None] * len(inputs)
+        
+        with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
+            # Submit all jobs
+            future_to_index = {
+                executor.submit(
+                    self._thread_safe_query, 
+                    input_dict.get('text', ''),
+                    input_dict.get('image', None),
+                    input_dict.get('maintain_history', False)
+                ): i 
+                for i, input_dict in enumerate(inputs)
+            }
+            
+            # Process results as they complete
+            for future in as_completed(future_to_index):
+                index = future_to_index[future]
+                try:
+                    response = future.result()
+                    results[index] = response
+                    
+                    # Call callback if provided
+                    if callback:
+                        callback(index, response)
+                        
+                except Exception as e:
+                    # Record error message as response
+                    results[index] = f"Error: {str(e)}"
+                    if callback:
+                        callback(index, results[index])
+        
+        return results
+    
+    def _thread_safe_query(self, text_input: str = "", image_input: Any = None, maintain_history: bool = False) -> str:
+        """
+        Thread-safe wrapper around the query method with rate limiting.
+        
+        Args:
+            text_input: The text input for the model
+            image_input: Optional image input
+            maintain_history: Whether to maintain conversation history
+            
+        Returns:
+            The model's response
+        """
+        # Apply rate limiting
+        with self.request_lock:
+            current_time = time.time()
+            time_since_last = current_time - self.last_request_time
+            
+            if time_since_last < self.request_interval:
+                sleep_time = self.request_interval - time_since_last
+                time.sleep(sleep_time)
+            
+            self.last_request_time = time.time()
+        
+        # Make the actual query with retries
+        for attempt in range(self.retry_attempts):
+            try:
+                return self.query(text_input, image_input, maintain_history)
+            except Exception as e:
+                if attempt < self.retry_attempts - 1:
+                    time.sleep(self.retry_delay * (attempt + 1))  # Exponential backoff
+                else:
+                    raise e  # Re-raise the exception on the last attempt
+    
+    def query(self, text_input: str = "", image_input: Any = None, maintain_history: bool = False) -> str:
+        """
+        Subclasses must implement this method.
+        """
+        raise NotImplementedError("Subclasses must implement the query method")