s3 url save instead of base64

This commit is contained in:
marko-kraemer 2025-05-19 01:28:14 +02:00
parent 70755d3074
commit f22412b963
3 changed files with 109 additions and 0 deletions

View File

@ -5,6 +5,7 @@ from agentpress.tool import ToolResult, openapi_schema, xml_schema
from agentpress.thread_manager import ThreadManager
from sandbox.tool_base import SandboxToolsBase
from utils.logger import logger
from utils.s3_upload_utils import upload_base64_image
class SandboxBrowserTool(SandboxToolsBase):
@ -59,6 +60,17 @@ class SandboxBrowserTool(SandboxToolsBase):
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(
thread_id=self.thread_id,
type="browser_state",
@ -83,6 +95,8 @@ class SandboxBrowserTool(SandboxToolsBase):
success_response["scrollable_content"] = result["pixels_below"] > 0
if result.get("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)

View File

@ -6,6 +6,9 @@ from typing import Optional
from supabase import create_async_client, AsyncClient
from utils.logger import logger
from utils.config import config
import base64
import uuid
from datetime import datetime
class DBConnection:
"""Singleton database connection manager using Supabase."""
@ -66,4 +69,45 @@ class DBConnection:
raise RuntimeError("Database not initialized")
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)}")

View File

@ -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)}")