fix(sb_image_edit_tool): update image processing to use base64 data instead of URL

This commit is contained in:
sharath 2025-07-05 12:39:55 +00:00
parent 9cfaac080c
commit 3d28721350
No known key found for this signature in database
1 changed files with 6 additions and 8 deletions

View File

@ -6,6 +6,7 @@ import httpx
from io import BytesIO from io import BytesIO
import uuid import uuid
from litellm import aimage_generation, aimage_edit from litellm import aimage_generation, aimage_edit
import base64
class SandboxImageEditTool(SandboxToolsBase): class SandboxImageEditTool(SandboxToolsBase):
@ -92,7 +93,7 @@ class SandboxImageEditTool(SandboxToolsBase):
) )
response = await aimage_edit( response = await aimage_edit(
image=image_io, image=[image_io], # Type in the LiteLLM SDK is wrong
prompt=prompt, prompt=prompt,
model="gpt-image-1", model="gpt-image-1",
n=1, n=1,
@ -155,19 +156,16 @@ class SandboxImageEditTool(SandboxToolsBase):
async def _process_image_response(self, response) -> str | ToolResult: async def _process_image_response(self, response) -> str | ToolResult:
"""Download generated image and save to sandbox with random name.""" """Download generated image and save to sandbox with random name."""
try: try:
original_url = response.data[0].url original_b64_str = response.data[0].b64_json
# Decode base64 image data
async with httpx.AsyncClient() as client: image_data = base64.b64decode(original_b64_str)
img_response = await client.get(original_url)
img_response.raise_for_status()
# Generate random filename # Generate random filename
random_filename = f"generated_image_{uuid.uuid4().hex[:8]}.png" random_filename = f"generated_image_{uuid.uuid4().hex[:8]}.png"
sandbox_path = f"{self.workspace_path}/{random_filename}" sandbox_path = f"{self.workspace_path}/{random_filename}"
# Save image to sandbox # Save image to sandbox
await self.sandbox.fs.upload_file(sandbox_path, img_response.content) await self.sandbox.fs.upload_file(image_data, sandbox_path)
return random_filename return random_filename
except Exception as e: except Exception as e: