mirror of https://github.com/kortix-ai/suna.git
s3 url save instead of base64
This commit is contained in:
parent
70755d3074
commit
f22412b963
|
@ -5,6 +5,7 @@ from agentpress.tool import ToolResult, openapi_schema, xml_schema
|
||||||
from agentpress.thread_manager import ThreadManager
|
from agentpress.thread_manager import ThreadManager
|
||||||
from sandbox.tool_base import SandboxToolsBase
|
from sandbox.tool_base import SandboxToolsBase
|
||||||
from utils.logger import logger
|
from utils.logger import logger
|
||||||
|
from utils.s3_upload_utils import upload_base64_image
|
||||||
|
|
||||||
|
|
||||||
class SandboxBrowserTool(SandboxToolsBase):
|
class SandboxBrowserTool(SandboxToolsBase):
|
||||||
|
@ -59,6 +60,17 @@ class SandboxBrowserTool(SandboxToolsBase):
|
||||||
|
|
||||||
logger.info("Browser automation request completed successfully")
|
logger.info("Browser automation request completed successfully")
|
||||||
|
|
||||||
|
if "screenshot_base64" in result:
|
||||||
|
try:
|
||||||
|
image_url = await upload_base64_image(result["screenshot_base64"])
|
||||||
|
result["image_url"] = image_url
|
||||||
|
# Remove base64 data from result to keep it clean
|
||||||
|
del result["screenshot_base64"]
|
||||||
|
logger.debug(f"Uploaded screenshot to {image_url}")
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(f"Failed to upload screenshot: {e}")
|
||||||
|
result["image_upload_error"] = str(e)
|
||||||
|
|
||||||
added_message = await self.thread_manager.add_message(
|
added_message = await self.thread_manager.add_message(
|
||||||
thread_id=self.thread_id,
|
thread_id=self.thread_id,
|
||||||
type="browser_state",
|
type="browser_state",
|
||||||
|
@ -83,6 +95,8 @@ class SandboxBrowserTool(SandboxToolsBase):
|
||||||
success_response["scrollable_content"] = result["pixels_below"] > 0
|
success_response["scrollable_content"] = result["pixels_below"] > 0
|
||||||
if result.get("ocr_text"):
|
if result.get("ocr_text"):
|
||||||
success_response["ocr_text"] = result["ocr_text"]
|
success_response["ocr_text"] = result["ocr_text"]
|
||||||
|
if result.get("image_url"):
|
||||||
|
success_response["image_url"] = result["image_url"]
|
||||||
|
|
||||||
return self.success_response(success_response)
|
return self.success_response(success_response)
|
||||||
|
|
||||||
|
|
|
@ -6,6 +6,9 @@ from typing import Optional
|
||||||
from supabase import create_async_client, AsyncClient
|
from supabase import create_async_client, AsyncClient
|
||||||
from utils.logger import logger
|
from utils.logger import logger
|
||||||
from utils.config import config
|
from utils.config import config
|
||||||
|
import base64
|
||||||
|
import uuid
|
||||||
|
from datetime import datetime
|
||||||
|
|
||||||
class DBConnection:
|
class DBConnection:
|
||||||
"""Singleton database connection manager using Supabase."""
|
"""Singleton database connection manager using Supabase."""
|
||||||
|
@ -66,4 +69,45 @@ class DBConnection:
|
||||||
raise RuntimeError("Database not initialized")
|
raise RuntimeError("Database not initialized")
|
||||||
return self._client
|
return self._client
|
||||||
|
|
||||||
|
async def upload_base64_image(self, base64_data: str, bucket_name: str = "browser-screenshots") -> str:
|
||||||
|
"""Upload a base64 encoded image to Supabase storage and return the URL.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
base64_data (str): Base64 encoded image data (with or without data URL prefix)
|
||||||
|
bucket_name (str): Name of the storage bucket to upload to
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
str: Public URL of the uploaded image
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
# Remove data URL prefix if present
|
||||||
|
if base64_data.startswith('data:'):
|
||||||
|
base64_data = base64_data.split(',')[1]
|
||||||
|
|
||||||
|
# Decode base64 data
|
||||||
|
image_data = base64.b64decode(base64_data)
|
||||||
|
|
||||||
|
# Generate unique filename
|
||||||
|
timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
|
||||||
|
unique_id = str(uuid.uuid4())[:8]
|
||||||
|
filename = f"image_{timestamp}_{unique_id}.png"
|
||||||
|
|
||||||
|
# Upload to Supabase storage
|
||||||
|
client = await self.client
|
||||||
|
storage_response = await client.storage.from_(bucket_name).upload(
|
||||||
|
filename,
|
||||||
|
image_data,
|
||||||
|
{"content-type": "image/png"}
|
||||||
|
)
|
||||||
|
|
||||||
|
# Get public URL
|
||||||
|
public_url = await client.storage.from_(bucket_name).get_public_url(filename)
|
||||||
|
|
||||||
|
logger.debug(f"Successfully uploaded image to {public_url}")
|
||||||
|
return public_url
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(f"Error uploading base64 image: {e}")
|
||||||
|
raise RuntimeError(f"Failed to upload image: {str(e)}")
|
||||||
|
|
||||||
|
|
||||||
|
|
|
@ -0,0 +1,51 @@
|
||||||
|
"""
|
||||||
|
Utility functions for handling image operations.
|
||||||
|
"""
|
||||||
|
|
||||||
|
import base64
|
||||||
|
import uuid
|
||||||
|
from datetime import datetime
|
||||||
|
from utils.logger import logger
|
||||||
|
from services.supabase import DBConnection
|
||||||
|
|
||||||
|
async def upload_base64_image(base64_data: str, bucket_name: str = "browser-screenshots") -> str:
|
||||||
|
"""Upload a base64 encoded image to Supabase storage and return the URL.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
base64_data (str): Base64 encoded image data (with or without data URL prefix)
|
||||||
|
bucket_name (str): Name of the storage bucket to upload to
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
str: Public URL of the uploaded image
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
# Remove data URL prefix if present
|
||||||
|
if base64_data.startswith('data:'):
|
||||||
|
base64_data = base64_data.split(',')[1]
|
||||||
|
|
||||||
|
# Decode base64 data
|
||||||
|
image_data = base64.b64decode(base64_data)
|
||||||
|
|
||||||
|
# Generate unique filename
|
||||||
|
timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
|
||||||
|
unique_id = str(uuid.uuid4())[:8]
|
||||||
|
filename = f"image_{timestamp}_{unique_id}.png"
|
||||||
|
|
||||||
|
# Upload to Supabase storage
|
||||||
|
db = DBConnection()
|
||||||
|
client = await db.client
|
||||||
|
storage_response = await client.storage.from_(bucket_name).upload(
|
||||||
|
filename,
|
||||||
|
image_data,
|
||||||
|
{"content-type": "image/png"}
|
||||||
|
)
|
||||||
|
|
||||||
|
# Get public URL
|
||||||
|
public_url = await client.storage.from_(bucket_name).get_public_url(filename)
|
||||||
|
|
||||||
|
logger.debug(f"Successfully uploaded image to {public_url}")
|
||||||
|
return public_url
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(f"Error uploading base64 image: {e}")
|
||||||
|
raise RuntimeError(f"Failed to upload image: {str(e)}")
|
Loading…
Reference in New Issue