Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
27 changes: 25 additions & 2 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -159,7 +159,7 @@ from lf_toolkit.chat import ChatRequest, ChatResponse, Message

## Image Upload

Upload PIL images to S3 using AWS SigV4 authentication:
Upload PIL images to AWS S3 or Google Cloud Storage:

```python
from PIL import Image
Expand All @@ -169,7 +169,17 @@ img = Image.open("diagram.png")
url = upload_image(img, folder_name="my-eval-function")
```

Required environment variables:
### Backend selection

`upload_image` picks a backend at call time:

1. `IMAGE_UPLOAD_BACKEND` env var, when set to `s3` or `gcs`, wins.
2. Otherwise, if `GCS_BUCKET` is set the GCS backend is used.
3. Otherwise S3 is used (the default — existing deployments are unaffected).

### S3 backend

Uses AWS SigV4-signed `PUT` requests.

| Variable | Description |
|---|---|
Expand All @@ -179,6 +189,19 @@ Required environment variables:
| `AWS_SESSION_TOKEN` | (optional) Session token |
| `AWS_REGION` | AWS region (default: `eu-west-2`) |

### GCS backend

Requires the `gcs` extra (`poetry install --extras gcs`, or
`lf_toolkit = { ..., extras = ["gcs"] }`). Authenticates with Application
Default Credentials (the runtime service account on Cloud Run / GKE / GCE) — no
static keys. The target bucket / prefix must be readable by whoever consumes the
returned URL (e.g. `roles/storage.objectViewer` for `allUsers`).

| Variable | Description |
|---|---|
| `GCS_BUCKET` | Target bucket name |
| `GCS_PUBLIC_BASE_URL` | (optional) URL host for the returned link (default: `https://storage.googleapis.com`) |

## Set Notation Parser

Parse and evaluate set expressions (requires `parsing` extra):
Expand Down
116 changes: 87 additions & 29 deletions lf_toolkit/evaluation/image_upload.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,11 @@
from botocore.awsrequest import AWSRequest
from botocore.credentials import Credentials

try:
from google.cloud import storage as _gcs_storage
except ImportError: # pragma: no cover - optional dependency, install the "gcs" extra
_gcs_storage = None

load_dotenv()

MIME_TO_FORMAT: Dict[str, List[str]] = {
Expand Down Expand Up @@ -69,6 +74,19 @@ def get_s3_bucket_uri() -> str:
return s3_uri


def resolve_upload_backend() -> str:
"""Decide which storage backend upload_image() should use.

``IMAGE_UPLOAD_BACKEND`` ("s3" or "gcs") wins when set to a known value.
Otherwise the presence of ``GCS_BUCKET`` selects GCS, and S3 is the default
so existing deployments keep working unchanged.
"""
backend = (os.getenv('IMAGE_UPLOAD_BACKEND') or '').strip().lower()
if backend in ('s3', 'gcs'):
return backend
return 'gcs' if os.getenv('GCS_BUCKET') else 's3'


def get_aws_signed_request(full_url, buffer, mime_type):
credentials = Credentials(
access_key=os.environ['AWS_ACCESS_KEY_ID'],
Expand Down Expand Up @@ -113,57 +131,97 @@ def get_aws_signed_request(full_url, buffer, mime_type):
return aws_request


def _upload_s3(folder_name: str, filename: str, data: bytes, mime_type: str) -> str:
"""Upload bytes to S3 with a SigV4-signed PUT and return the object URL."""
base_url: str = get_s3_bucket_uri()
full_url = os.path.join(base_url, folder_name, filename)

aws_request = get_aws_signed_request(full_url, data, mime_type).prepare()

response: requests.Response = requests.request(
method=aws_request.method,
url=aws_request.url,
data=aws_request.body,
headers=aws_request.headers,
timeout=30
)

if response.status_code != 200:
raise ImageUploadError(
f"Upload failed with status code {response.status_code}: {response.text}"
)

return full_url


def _upload_gcs(folder_name: str, filename: str, data: bytes, mime_type: str) -> str:
"""Upload bytes to Google Cloud Storage and return the object URL.

Authenticates with Application Default Credentials (the runtime service
account on Cloud Run / GKE / GCE) -- no static keys. Set ``GCS_BUCKET`` to
the target bucket and, optionally, ``GCS_PUBLIC_BASE_URL`` to override the
returned URL's host (e.g. a CDN or custom domain).
"""
if _gcs_storage is None:
raise ImageUploadError(
"google-cloud-storage is not installed; install lf_toolkit with the "
"'gcs' extra to use IMAGE_UPLOAD_BACKEND=gcs"
)

bucket_name: Optional[str] = os.getenv('GCS_BUCKET')
if not bucket_name:
raise MissingEnvironmentVariableError(
"GCS_BUCKET environment variable is not set"
)

blob_name = f"{folder_name}/{filename}"
client = _gcs_storage.Client()
blob = client.bucket(bucket_name).blob(blob_name)
blob.upload_from_string(data, content_type=mime_type)

base_url = os.getenv('GCS_PUBLIC_BASE_URL', 'https://storage.googleapis.com').rstrip('/')
return f"{base_url}/{bucket_name}/{blob_name}"


def upload_image(img: Image.Image, folder_name: str) -> str:
"""Upload PIL image with comprehensive MIME type validation
"""Upload a PIL image to the configured storage backend.

The backend is chosen by :func:`resolve_upload_backend` (env var
``IMAGE_UPLOAD_BACKEND=s3|gcs``, else auto-detected from ``GCS_BUCKET`` /
``S3_BUCKET_URI``, defaulting to S3).

Args:
folder_name: name of folder to save image
img: PIL Image object to upload
folder_name: name of the folder/prefix to store the image under

Returns:
JSON response from the server as a dictionary
The public URL of the uploaded object

Raises:
InvalidMimeTypeError: If MIME type validation fails
MissingEnvironmentVariableError: If S3_BUCKET_URI is not set
MissingEnvironmentVariableError: If required env vars are not set
ImageUploadError: If upload fails for any reason
"""
try:
# Get URL from environment variable
base_url: str = get_s3_bucket_uri()

filename: str = generate_file_name(img)

full_url = os.path.join(base_url, folder_name, filename)

if img.format is None:
img.format = 'PNG'

mime_type = FORMAT_TO_MIME[img.format.upper()]

buffer: BytesIO = BytesIO()
img_format: str = img.format if img.format else 'PNG'
img.save(buffer, format=img_format)
buffer.seek(0)

aws_request = get_aws_signed_request(full_url, buffer, mime_type).prepare()

response: requests.Response = requests.request(
method=aws_request.method,
url=aws_request.url,
data=aws_request.body,
headers=aws_request.headers,
timeout=30
)

if response.status_code != 200:
raise ImageUploadError(
f"Upload failed with status code {response.status_code}: {response.text}"
)
img.save(buffer, format=img.format)
data: bytes = buffer.getvalue()

return full_url
backend = resolve_upload_backend()
if backend == 'gcs':
return _upload_gcs(folder_name, filename, data, mime_type)
return _upload_s3(folder_name, filename, data, mime_type)

except (InvalidMimeTypeError, MissingEnvironmentVariableError):
except ImageUploadError:
# InvalidMimeTypeError / MissingEnvironmentVariableError / backend errors
# already carry a useful message -- propagate as-is.
raise
except requests.exceptions.RequestException as e:
raise ImageUploadError(f"Network error: {str(e)}")
Expand Down
Loading
Loading