[Feature Request] Integrate Grok image (#3040)

Co-authored-by: Wendong-Fan <133094783+Wendong-Fan@users.noreply.github.com>

Author: jajajalalalala

.env

@@ -130,3 +130,7 @@
 # E2B
 # E2B_API_KEY="Fill your e2b or e2b-compatible sandbox provider API Key here"
 # E2B_DOMAIN="Fill your custom e2b domain here"
+
+# Grok API key
+# XAI_API_KEY="Fill your Grok API Key here"
+# XAI_API_BASE_URL="Fill your Grok API Base URL here"

camel/toolkits/__init__.py

@@ -23,7 +23,7 @@
 from .math_toolkit import MathToolkit
 from .search_toolkit import SearchToolkit
 from .weather_toolkit import WeatherToolkit
-from .openai_image_toolkit import OpenAIImageToolkit
+from .image_generation_toolkit import ImageGenToolkit, OpenAIImageToolkit
 from .ask_news_toolkit import AskNewsToolkit, AsyncAskNewsToolkit
 from .linkedin_toolkit import LinkedInToolkit
 from .reddit_toolkit import RedditToolkit
@@ -102,7 +102,7 @@
     'SearchToolkit',
     'SlackToolkit',
     'WhatsAppToolkit',
-    'OpenAIImageToolkit',
+    'ImageGenToolkit',
     'TwitterToolkit',
     'WeatherToolkit',
     'RetrievalToolkit',
@@ -151,7 +151,7 @@
     'PlaywrightMCPToolkit',
     'WolframAlphaToolkit',
     'BohriumToolkit',
-    'OpenAIImageToolkit',
+    'OpenAIImageToolkit',  # Backward compatibility
     'TaskPlanningToolkit',
     'HybridBrowserToolkit',
     'EdgeOnePagesMCPToolkit',

camel/toolkits/openai_image_toolkit.py renamed to camel/toolkits/image_generation_toolkit.py

@@ -15,7 +15,7 @@
 import base64
 import os
 from io import BytesIO
-from typing import List, Literal, Optional, Union
+from typing import ClassVar, List, Literal, Optional, Tuple, Union
 
 from openai import OpenAI
 from PIL import Image
@@ -29,21 +29,32 @@
 
 
 @MCPServer()
-class OpenAIImageToolkit(BaseToolkit):
-    r"""A class toolkit for image generation using OpenAI's
-    Image Generation API.
-    """
-
-    @api_keys_required(
-        [
-            ("api_key", "OPENAI_API_KEY"),
-        ]
-    )
+class ImageGenToolkit(BaseToolkit):
+    r"""A class toolkit for image generation using Grok and OpenAI models."""
+
+    GROK_MODELS: ClassVar[List[str]] = [
+        "grok-2-image",
+        "grok-2-image-latest",
+        "grok-2-image-1212",
+    ]
+    OPENAI_MODELS: ClassVar[List[str]] = [
+        "gpt-image-1",
+        "dall-e-3",
+        "dall-e-2",
+    ]
+
     def __init__(
         self,
         model: Optional[
-            Literal["gpt-image-1", "dall-e-3", "dall-e-2"]
-        ] = "gpt-image-1",
+            Literal[
+                "gpt-image-1",
+                "dall-e-3",
+                "dall-e-2",
+                "grok-2-image",
+                "grok-2-image-latest",
+                "grok-2-image-1212",
+            ]
+        ] = "dall-e-3",
         timeout: Optional[float] = None,
         api_key: Optional[str] = None,
         url: Optional[str] = None,
@@ -72,12 +83,12 @@ def __init__(
         # NOTE: Some arguments are set in the constructor to prevent the agent
         # from making invalid API calls with model-specific parameters. For
         # example, the 'style' argument is only supported by 'dall-e-3'.
-        r"""Initializes a new instance of the OpenAIImageToolkit class.
+        r"""Initializes a new instance of the ImageGenToolkit class.
 
         Args:
             api_key (Optional[str]): The API key for authenticating
-                with the OpenAI service. (default: :obj:`None`)
-            url (Optional[str]): The url to the OpenAI service.
+                with the image model service. (default: :obj:`None`)
+            url (Optional[str]): The url to the image model service.
                 (default: :obj:`None`)
             model (Optional[str]): The model to use.
                 (default: :obj:`"dall-e-3"`)
@@ -103,9 +114,23 @@ def __init__(
                 image.(default: :obj:`"image_save"`)
         """
         super().__init__(timeout=timeout)
-        api_key = api_key or os.environ.get("OPENAI_API_KEY")
-        url = url or os.environ.get("OPENAI_API_BASE_URL")
-        self.client = OpenAI(api_key=api_key, base_url=url)
+        if model not in self.GROK_MODELS + self.OPENAI_MODELS:
+            available_models = sorted(self.OPENAI_MODELS + self.GROK_MODELS)
+            raise ValueError(
+                f"Unsupported model: {model}. "
+                f"Supported models are: {available_models}"
+            )
+
+        # Set default url for Grok models
+        url = "https://api.x.ai/v1" if model in self.GROK_MODELS else url
+
+        api_key, base_url = (
+            self.get_openai_credentials(url, api_key)
+            if model in self.OPENAI_MODELS
+            else self.get_grok_credentials(url, api_key)
+        )
+
+        self.client = OpenAI(api_key=api_key, base_url=base_url)
         self.model = model
         self.size = size
         self.quality = quality
@@ -139,7 +164,7 @@ def base64_to_image(self, base64_string: str) -> Optional[Image.Image]:
             return None
 
     def _build_base_params(self, prompt: str, n: Optional[int] = None) -> dict:
-        r"""Build base parameters dict for OpenAI API calls.
+        r"""Build base parameters dict for Image Model API calls.
 
         Args:
             prompt (str): The text prompt for the image operation.
@@ -153,6 +178,10 @@ def _build_base_params(self, prompt: str, n: Optional[int] = None) -> dict:
         # basic parameters supported by all models
         if n is not None:
             params["n"] = n  # type: ignore[assignment]
+
+        if self.model in self.GROK_MODELS:
+            return params
+
         if self.size is not None:
             params["size"] = self.size
 
@@ -179,16 +208,18 @@ def _build_base_params(self, prompt: str, n: Optional[int] = None) -> dict:
                 params["quality"] = self.quality
             if self.background is not None:
                 params["background"] = self.background
-
         return params
 
     def _handle_api_response(
-        self, response, image_name: Union[str, List[str]], operation: str
+        self,
+        response,
+        image_name: Union[str, List[str]],
+        operation: str,
     ) -> str:
-        r"""Handle API response from OpenAI image operations.
+        r"""Handle API response from image operations.
 
         Args:
-            response: The response object from OpenAI API.
+            response: The response object from image model API.
             image_name (Union[str, List[str]]): Name(s) for the saved image
                 file(s). If str, the same name is used for all images (will
                 cause error for multiple images). If list, must have exactly
@@ -198,8 +229,9 @@ def _handle_api_response(
         Returns:
             str: Success message with image path/URL or error message.
         """
+        source = "Grok" if self.model in self.GROK_MODELS else "OpenAI"
         if response.data is None or len(response.data) == 0:
-            error_msg = "No image data returned from OpenAI API."
+            error_msg = f"No image data returned from {source} API."
             logger.error(error_msg)
             return error_msg
 
@@ -283,7 +315,7 @@ def generate_image(
         image_name: Union[str, List[str]] = "image.png",
         n: int = 1,
     ) -> str:
-        r"""Generate an image using OpenAI's Image Generation models.
+        r"""Generate an image using image models.
         The generated image will be saved locally (for ``b64_json`` response
         formats) or an image URL will be returned (for ``url`` response
         formats).
@@ -309,15 +341,50 @@ def generate_image(
             logger.error(error_msg)
             return error_msg
 
+    @api_keys_required([("api_key", "XAI_API_KEY")])
+    def get_grok_credentials(self, url, api_key) -> Tuple[str, str]:  # type: ignore[return-value]
+        r"""Get API credentials for the specified Grok model.
+
+        Args:
+            url (str): The base URL for the Grok API.
+            api_key (str): The API key for the Grok API.
+
+        Returns:
+            tuple: (api_key, base_url)
+        """
+
+        # Get credentials based on model type
+        api_key = api_key or os.getenv("XAI_API_KEY")
+        return api_key, url
+
+    @api_keys_required([("api_key", "OPENAI_API_KEY")])
+    def get_openai_credentials(self, url, api_key) -> Tuple[str, str | None]:  # type: ignore[return-value]
+        r"""Get API credentials for the specified OpenAI model.
+
+        Args:
+            url (str): The base URL for the OpenAI API.
+            api_key (str): The API key for the OpenAI API.
+
+        Returns:
+            Tuple[str, str | None]: (api_key, base_url)
+        """
+
+        api_key = api_key or os.getenv("OPENAI_API_KEY")
+        base_url = url or os.getenv("OPENAI_API_BASE_URL")
+        return api_key, base_url
+
     def get_tools(self) -> List[FunctionTool]:
-        r"""Returns a list of FunctionTool objects representing the
-        functions in the toolkit.
+        r"""Returns a list of FunctionTool objects representing the functions
+            in the toolkit.
 
         Returns:
-            List[FunctionTool]: A list of FunctionTool objects
-                representing the functions in the toolkit.
+            List[FunctionTool]: A list of FunctionTool objects representing the
+                functions in the toolkit.
         """
         return [
             FunctionTool(self.generate_image),
-            # could add edit_image function later
         ]
+
+
+# Backward compatibility alias
+OpenAIImageToolkit = ImageGenToolkit

camel/utils/commons.py

@@ -354,6 +354,8 @@ def wrapper(*args: Any, **kwargs: Any) -> Any:
                 key_way = "https://www.zhipuai.cn/"
             elif env_var_name == 'KLAVIS_API_KEY':
                 key_way = "https://www.klavis.ai/docs"
+            elif env_var_name == 'XAI_API_KEY':
+                key_way = "https://api.x.ai/v1"
 
             if missing_keys:
                 raise ValueError(

examples/toolkits/grok_image_example.py

@@ -0,0 +1,49 @@
+# ========= Copyright 2023-2024 @ CAMEL-AI.org. All Rights Reserved. =========
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ========= Copyright 2023-2024 @ CAMEL-AI.org. All Rights Reserved. =========
+from camel.agents import ChatAgent
+from camel.models import ModelFactory
+from camel.toolkits import ImageGenToolkit
+from camel.types import ModelPlatformType, ModelType
+
+# Define system message
+sys_msg = "You are a helpful assistant that can generate images."
+
+# Create Image Generation Toolkit with Grok-2 model and base64 response format
+tools = [
+    *ImageGenToolkit(
+        model="grok-2-image-1212",
+        response_format="b64_json",
+    ).get_tools()
+]
+
+model = ModelFactory.create(
+    model_platform=ModelPlatformType.DEFAULT,
+    model_type=ModelType.DEFAULT,
+)
+
+# Set agent
+camel_agent = ChatAgent(
+    system_message=sys_msg,
+    model=model,
+    tools=tools,
+)
+
+# Define a user message
+usr_msg = "Generate 1 image of a camel working out in a gym."
+
+# Get response information
+response = camel_agent.step(usr_msg)
+
+print(f"Tool calls made: {len(response.info['tool_calls'])}")
+print(f"\nAgent response: {response.msg}")

examples/toolkits/openai_image_example.py

@@ -13,15 +13,15 @@
 # ========= Copyright 2023-2024 @ CAMEL-AI.org. All Rights Reserved. =========
 from camel.agents import ChatAgent
 from camel.models import ModelFactory
-from camel.toolkits import OpenAIImageToolkit
+from camel.toolkits import ImageGenToolkit
 from camel.types import ModelPlatformType, ModelType
 
 # Define system message
 sys_msg = "You are a helpful assistant that can generate images."
 
-# Create OpenAI Image Toolkit with DALL-E 3 model and base64 response format
+# Create Image Generation Toolkit with DALL-E 3 model and base64 response
 tools = [
-    *OpenAIImageToolkit(
+    *ImageGenToolkit(
         model="dall-e-3",
         response_format="b64_json",
         size="1024x1024",

test/toolkits/test_openai_image.py renamed to test/toolkits/test_image_generation.py

@@ -14,12 +14,12 @@
 import pytest
 from PIL import Image
 
-from camel.toolkits import OpenAIImageToolkit
+from camel.toolkits import ImageGenToolkit
 
 
 @pytest.fixture
 def image_toolkit():
-    return OpenAIImageToolkit()
+    return ImageGenToolkit()
 
 
 def test_base64_to_image_valid(image_toolkit):