diff --git a/backend/agent/prompt.py b/backend/agent/prompt.py
index 9bd0dc22..1fae53c1 100644
--- a/backend/agent/prompt.py
+++ b/backend/agent/prompt.py
@@ -89,7 +89,31 @@ You have the ability to execute operations using both Python and CLI tools:
* Supported formats include JPG, PNG, GIF, WEBP, and other common image formats.
* Maximum file size limit is 10 MB.
-### 2.2.7 DATA PROVIDERS
+### 2.2.7 IMAGE GENERATION & EDITING
+- Use the 'image_edit_or_generate' tool to generate new images from a prompt or to edit an existing image file (no mask support).
+ * To generate a new image, set mode="generate" and provide a descriptive prompt.
+ * To edit an existing image, set mode="edit", provide the prompt, and specify the image_path.
+ * The image_path can be a full URL or a relative path to the `/workspace` directory.
+ * Example (generate):
+
+
+ generate
+ A futuristic cityscape at sunset
+
+
+ * Example (edit):
+
+
+ edit
+ Add a red hat to the person in the image
+ http://example.com/images/person.png
+
+
+ * ALWAYS use this tool for any image creation or editing tasks. Do not attempt to generate or edit images by any other means.
+ * You must use edit mode when the user asks you to edit an image or change an existing image in any way.
+ * Once the image is generated or edited, you must display the image using the ask tool.
+
+### 2.2.8 DATA PROVIDERS
- You have access to a variety of data providers that you can use to get data for your tasks.
- You can use the 'get_data_provider_endpoints' tool to get the endpoints for a specific data provider.
- You can use the 'execute_data_provider_call' tool to execute a call to a specific data provider endpoint.
diff --git a/backend/agent/run.py b/backend/agent/run.py
index e97ef729..365676a6 100644
--- a/backend/agent/run.py
+++ b/backend/agent/run.py
@@ -25,6 +25,7 @@ from utils.logger import logger
from utils.auth_utils import get_account_id_from_thread
from services.billing import check_billing_status
from agent.tools.sb_vision_tool import SandboxVisionTool
+from agent.tools.sb_image_edit_tool import SandboxImageEditTool
from services.langfuse import langfuse
from langfuse.client import StatefulTraceClient
from services.langfuse import langfuse
@@ -107,6 +108,7 @@ async def run_agent(
thread_manager.add_tool(MessageTool)
thread_manager.add_tool(SandboxWebSearchTool, project_id=project_id, thread_manager=thread_manager)
thread_manager.add_tool(SandboxVisionTool, project_id=project_id, thread_id=thread_id, thread_manager=thread_manager)
+ thread_manager.add_tool(SandboxImageEditTool, project_id=project_id, thread_id=thread_id, thread_manager=thread_manager)
if config.RAPID_API_KEY:
thread_manager.add_tool(DataProvidersTool)
else:
diff --git a/backend/agent/tools/sb_image_edit_tool.py b/backend/agent/tools/sb_image_edit_tool.py
new file mode 100644
index 00000000..c79596b5
--- /dev/null
+++ b/backend/agent/tools/sb_image_edit_tool.py
@@ -0,0 +1,167 @@
+from typing import Optional
+from agentpress.tool import ToolResult, openapi_schema, xml_schema
+from sandbox.tool_base import SandboxToolsBase
+from agentpress.thread_manager import ThreadManager
+from openai import OpenAI
+import httpx
+import os
+from io import BytesIO
+import uuid
+
+
+class SandboxImageEditTool(SandboxToolsBase):
+ """Tool for generating or editing images using OpenAI DALL-E via OpenAI SDK (no mask support)."""
+
+ def __init__(self, project_id: str, thread_id: str, thread_manager: ThreadManager):
+ super().__init__(project_id, thread_manager)
+ self.thread_id = thread_id
+ self.thread_manager = thread_manager
+ self.client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
+
+ @openapi_schema(
+ {
+ "type": "function",
+ "function": {
+ "name": "image_edit_or_generate",
+ "description": "Generate a new image from a prompt, or edit an existing image (no mask support) using OpenAI DALL-E via OpenAI SDK. Stores the result in the thread context.",
+ "parameters": {
+ "type": "object",
+ "properties": {
+ "mode": {
+ "type": "string",
+ "enum": ["generate", "edit"],
+ "description": "'generate' to create a new image from a prompt, 'edit' to edit an existing image.",
+ },
+ "prompt": {
+ "type": "string",
+ "description": "Text prompt describing the desired image or edit.",
+ },
+ "image_path": {
+ "type": "string",
+ "description": "(edit mode only) Path to the image file to edit, relative to /workspace. Required for 'edit'.",
+ },
+ },
+ "required": ["mode", "prompt"],
+ },
+ },
+ }
+ )
+ @xml_schema(
+ tag_name="image-edit-or-generate",
+ mappings=[
+ {"param_name": "mode", "node_type": "attribute", "path": "."},
+ {"param_name": "prompt", "node_type": "attribute", "path": "."},
+ {"param_name": "image_path", "node_type": "attribute", "path": "."},
+ ],
+ example="""
+
+
+ generate
+ A futuristic cityscape at sunset
+
+
+ """,
+ )
+ async def image_edit_or_generate(
+ self,
+ mode: str,
+ prompt: str,
+ image_path: Optional[str] = None,
+ ) -> ToolResult:
+ """Generate or edit images using OpenAI DALL-E via OpenAI SDK (no mask support)."""
+ try:
+ await self._ensure_sandbox()
+
+ if mode == "generate":
+ response = self.client.images.generate(
+ prompt=prompt, n=1, size="1024x1024"
+ )
+ elif mode == "edit":
+ if not image_path:
+ return self.fail_response("'image_path' is required for edit mode.")
+
+ image_bytes = await self._get_image_bytes(image_path)
+ if isinstance(image_bytes, ToolResult): # Error occurred
+ return image_bytes
+
+ # Create BytesIO object with proper filename to set MIME type
+ image_io = BytesIO(image_bytes)
+ image_io.name = "image.png" # Set filename to ensure proper MIME type detection
+
+ response = self.client.images.edit(
+ image=image_io, prompt=prompt, n=1, size="1024x1024"
+ )
+ else:
+ return self.fail_response("Invalid mode. Use 'generate' or 'edit'.")
+
+ # Download and save the generated image to sandbox
+ image_filename = await self._process_image_response(response)
+ if isinstance(image_filename, ToolResult): # Error occurred
+ return image_filename
+
+ return self.success_response(
+ f"Successfully generated image using mode '{mode}'. Image saved as: {image_filename}. You can use the ask tool to display the image."
+ )
+
+ except Exception as e:
+ return self.fail_response(
+ f"An error occurred during image generation/editing: {str(e)}"
+ )
+
+ async def _get_image_bytes(self, image_path: str) -> bytes | ToolResult:
+ """Get image bytes from URL or local file path."""
+ if image_path.startswith(("http://", "https://")):
+ return await self._download_image_from_url(image_path)
+ else:
+ return await self._read_image_from_sandbox(image_path)
+
+ async def _download_image_from_url(self, url: str) -> bytes | ToolResult:
+ """Download image from URL."""
+ try:
+ async with httpx.AsyncClient() as client:
+ response = await client.get(url)
+ response.raise_for_status()
+ return response.content
+ except Exception:
+ return self.fail_response(f"Could not download image from URL: {url}")
+
+ async def _read_image_from_sandbox(self, image_path: str) -> bytes | ToolResult:
+ """Read image from sandbox filesystem."""
+ try:
+ cleaned_path = self.clean_path(image_path)
+ full_path = f"{self.workspace_path}/{cleaned_path}"
+
+ # Check if file exists and is not a directory
+ file_info = self.sandbox.fs.get_file_info(full_path)
+ if file_info.is_dir:
+ return self.fail_response(
+ f"Path '{cleaned_path}' is a directory, not an image file."
+ )
+
+ return self.sandbox.fs.download_file(full_path)
+
+ except Exception as e:
+ return self.fail_response(
+ f"Could not read image file from sandbox: {image_path} - {str(e)}"
+ )
+
+ async def _process_image_response(self, response) -> str | ToolResult:
+ """Download generated image and save to sandbox with random name."""
+ try:
+ original_url = response.data[0].url
+
+ async with httpx.AsyncClient() as client:
+ img_response = await client.get(original_url)
+ img_response.raise_for_status()
+
+ # Generate random filename
+
+ random_filename = f"generated_image_{uuid.uuid4().hex[:8]}.png"
+ sandbox_path = f"{self.workspace_path}/{random_filename}"
+
+ # Save image to sandbox
+ self.sandbox.fs.upload_file(sandbox_path, img_response.content)
+ return random_filename
+
+ except Exception as e:
+ return self.fail_response(f"Failed to download and save image: {str(e)}")
diff --git a/backend/poetry.lock b/backend/poetry.lock
index dbf884d8..14202f37 100644
--- a/backend/poetry.lock
+++ b/backend/poetry.lock
@@ -1428,14 +1428,14 @@ openai = ["openai (>=0.27.8)"]
[[package]]
name = "litellm"
-version = "1.66.1"
+version = "1.72.2"
description = "Library to easily interface with LLM API providers"
optional = false
python-versions = "!=2.7.*,!=3.0.*,!=3.1.*,!=3.2.*,!=3.3.*,!=3.4.*,!=3.5.*,!=3.6.*,!=3.7.*,>=3.8"
groups = ["main"]
files = [
- {file = "litellm-1.66.1-py3-none-any.whl", hash = "sha256:1f601fea3f086c1d2d91be60b9db115082a2f3a697e4e0def72f8b9c777c7232"},
- {file = "litellm-1.66.1.tar.gz", hash = "sha256:98f7add913e5eae2131dd412ee27532d9a309defd9dbb64f6c6c42ea8a2af068"},
+ {file = "litellm-1.72.2-py3-none-any.whl", hash = "sha256:51e70f5cd98748a603d725ef29ede0ecad3d55e1a89cbbcec8d12d6fff55bff4"},
+ {file = "litellm-1.72.2.tar.gz", hash = "sha256:b50c7f7a0df67117889479264a12b0dea9c566a02173d4c3159540a13760d38b"},
]
[package.dependencies]
@@ -1453,7 +1453,8 @@ tokenizers = "*"
[package.extras]
extra-proxy = ["azure-identity (>=1.15.0,<2.0.0)", "azure-keyvault-secrets (>=4.8.0,<5.0.0)", "google-cloud-kms (>=2.21.3,<3.0.0)", "prisma (==0.11.0)", "redisvl (>=0.4.1,<0.5.0) ; python_version >= \"3.9\" and python_version < \"3.14\"", "resend (>=0.8.0,<0.9.0)"]
-proxy = ["PyJWT (>=2.8.0,<3.0.0)", "apscheduler (>=3.10.4,<4.0.0)", "backoff", "boto3 (==1.34.34)", "cryptography (>=43.0.1,<44.0.0)", "fastapi (>=0.115.5,<0.116.0)", "fastapi-sso (>=0.16.0,<0.17.0)", "gunicorn (>=23.0.0,<24.0.0)", "litellm-proxy-extras (==0.1.7)", "mcp (==1.5.0) ; python_version >= \"3.10\"", "orjson (>=3.9.7,<4.0.0)", "pynacl (>=1.5.0,<2.0.0)", "python-multipart (>=0.0.18,<0.0.19)", "pyyaml (>=6.0.1,<7.0.0)", "rq", "uvicorn (>=0.29.0,<0.30.0)", "uvloop (>=0.21.0,<0.22.0)", "websockets (>=13.1.0,<14.0.0)"]
+proxy = ["PyJWT (>=2.8.0,<3.0.0)", "apscheduler (>=3.10.4,<4.0.0)", "backoff", "boto3 (==1.34.34)", "cryptography (>=43.0.1,<44.0.0)", "fastapi (>=0.115.5,<0.116.0)", "fastapi-sso (>=0.16.0,<0.17.0)", "gunicorn (>=23.0.0,<24.0.0)", "litellm-enterprise (==0.1.7)", "litellm-proxy-extras (==0.2.3)", "mcp (==1.5.0) ; python_version >= \"3.10\"", "orjson (>=3.9.7,<4.0.0)", "pynacl (>=1.5.0,<2.0.0)", "python-multipart (>=0.0.18,<0.0.19)", "pyyaml (>=6.0.1,<7.0.0)", "rich (==13.7.1)", "rq", "uvicorn (>=0.29.0,<0.30.0)", "uvloop (>=0.21.0,<0.22.0) ; sys_platform != \"win32\"", "websockets (>=13.1.0,<14.0.0)"]
+utils = ["numpydoc"]
[[package]]
name = "mailtrap"
@@ -3904,4 +3905,4 @@ testing = ["coverage[toml]", "zope.event", "zope.testing"]
[metadata]
lock-version = "2.1"
python-versions = "^3.11"
-content-hash = "09a851f3db2d0b1f130405a69c1661c453f82ce23e078256bc6749662af897a7"
+content-hash = "3b983fbe8614f4e59280b2087fa4bcc574502d58fc75aa73a44426279f99e3d2"
diff --git a/backend/pyproject.toml b/backend/pyproject.toml
index 05af3385..029e05f3 100644
--- a/backend/pyproject.toml
+++ b/backend/pyproject.toml
@@ -19,7 +19,7 @@ classifiers = [
[tool.poetry.dependencies]
python = "^3.11"
python-dotenv = "1.0.1"
-litellm = "1.66.1"
+litellm = "1.72.2"
click = "8.1.7"
questionary = "2.0.1"
requests = "^2.31.0"
diff --git a/backend/requirements.txt b/backend/requirements.txt
index 395edcae..2350d88e 100644
--- a/backend/requirements.txt
+++ b/backend/requirements.txt
@@ -1,5 +1,5 @@
python-dotenv==1.0.1
-litellm==1.66.1
+litellm==1.72.2
click==8.1.7
questionary==2.0.1
requests>=2.31.0