From 28b597f1e70cbcbdc9bf86b9fbfd4ac50774d3e8 Mon Sep 17 00:00:00 2001 From: Marcus Messer Date: Tue, 1 Sep 2026 15:46:13 +0100 Subject: [PATCH] Added support for Google Cloud Storage in image uploads, including environment-based backend selection. Expanded related tests and updated documentation. --- README.md | 27 +++- lf_toolkit/evaluation/image_upload.py | 116 ++++++++++++----- pyproject.toml | 7 ++ tests/evaluation/image_upload_test.py | 172 +++++++++++++++++++++++--- 4 files changed, 271 insertions(+), 51 deletions(-) diff --git a/README.md b/README.md index 7abb764..63c8fee 100644 --- a/README.md +++ b/README.md @@ -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 @@ -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 | |---|---| @@ -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): diff --git a/lf_toolkit/evaluation/image_upload.py b/lf_toolkit/evaluation/image_upload.py index d0d69f4..185dcc7 100644 --- a/lf_toolkit/evaluation/image_upload.py +++ b/lf_toolkit/evaluation/image_upload.py @@ -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]] = { @@ -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'], @@ -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)}") diff --git a/pyproject.toml b/pyproject.toml index f544eb5..00039ef 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -38,6 +38,12 @@ fastapi = { version = "^0.115.0", optional = true } pywin32 = { version = "^306", platform = "win32", optional = true } +########################## +# gcs image-upload dependencies +########################## + +google-cloud-storage = { version = "^2.18", optional = true } + ########################## # dev dependencies ########################## @@ -66,6 +72,7 @@ datamodel-code-generator = "^0.55.0" [tool.poetry.extras] parsing = ["antlr4-python3-runtime", "lark", "latex2sympy"] ipc = ["pywin32"] +gcs = ["google-cloud-storage"] http = ["fastapi"] [tool.isort] diff --git a/tests/evaluation/image_upload_test.py b/tests/evaluation/image_upload_test.py index b4a0125..aec7f28 100644 --- a/tests/evaluation/image_upload_test.py +++ b/tests/evaluation/image_upload_test.py @@ -8,6 +8,7 @@ from lf_toolkit.evaluation.image_upload import ( generate_file_name, get_s3_bucket_uri, + resolve_upload_backend, upload_image, ImageUploadError, InvalidMimeTypeError, @@ -15,6 +16,13 @@ ) +def fake_getenv(mapping): + """Build an os.getenv replacement backed by ``mapping`` (respects default).""" + def _getenv(key, default=None): + return mapping.get(key, default) + return _getenv + + class TestGenerateFileName: """Test suite for generate_file_name function""" @@ -94,8 +102,38 @@ def test_get_s3_bucket_uri_empty_string(self): get_s3_bucket_uri() -class TestUploadImage: - """Test suite for upload_image function""" +class TestResolveUploadBackend: + """Test suite for resolve_upload_backend function""" + + @patch('lf_toolkit.evaluation.image_upload.os.getenv') + def test_defaults_to_s3(self, mock_getenv): + mock_getenv.side_effect = fake_getenv({}) + assert resolve_upload_backend() == 's3' + + @patch('lf_toolkit.evaluation.image_upload.os.getenv') + def test_gcs_bucket_selects_gcs(self, mock_getenv): + mock_getenv.side_effect = fake_getenv({'GCS_BUCKET': 'my-bucket'}) + assert resolve_upload_backend() == 'gcs' + + @patch('lf_toolkit.evaluation.image_upload.os.getenv') + def test_explicit_backend_wins_over_autodetect(self, mock_getenv): + mock_getenv.side_effect = fake_getenv({ + 'IMAGE_UPLOAD_BACKEND': 'S3', + 'GCS_BUCKET': 'my-bucket', + }) + assert resolve_upload_backend() == 's3' + + @patch('lf_toolkit.evaluation.image_upload.os.getenv') + def test_unknown_backend_value_falls_back_to_autodetect(self, mock_getenv): + mock_getenv.side_effect = fake_getenv({ + 'IMAGE_UPLOAD_BACKEND': 'azure', + 'S3_BUCKET_URI': 'https://s3.amazonaws.com/bucket', + }) + assert resolve_upload_backend() == 's3' + + +class TestUploadImageS3: + """Test suite for upload_image function (S3 backend)""" @patch('lf_toolkit.evaluation.image_upload.requests.request') @patch('lf_toolkit.evaluation.image_upload.get_aws_signed_request') @@ -103,11 +141,9 @@ class TestUploadImage: @patch('lf_toolkit.evaluation.image_upload.uuid.uuid4') def test_successful_upload(self, mock_uuid, mock_getenv, mock_get_aws_signed_request, mock_request): """Test successful image upload with UUID-based filename""" - # Setup mocks mock_uuid.return_value = uuid.UUID('12345678-1234-5678-1234-567812345678') - mock_getenv.return_value = 'https://s3.amazonaws.com/eduvision' + mock_getenv.side_effect = fake_getenv({'S3_BUCKET_URI': 'https://s3.amazonaws.com/eduvision'}) - # Mock the AWS signed request mock_prepared_request = Mock() mock_prepared_request.method = 'PUT' mock_prepared_request.url = 'https://s3.amazonaws.com/eduvision/eduvision/12345678-1234-5678-1234-567812345678.jpeg' @@ -122,14 +158,11 @@ def test_successful_upload(self, mock_uuid, mock_getenv, mock_get_aws_signed_req mock_response.status_code = 200 mock_request.return_value = mock_response - # Create a real PIL image for testing img = Image.new('RGB', (100, 100), color='red') img.format = 'JPEG' - # Execute result = upload_image(img, "eduvision") - # Verify response assert result == 'https://s3.amazonaws.com/eduvision/eduvision/12345678-1234-5678-1234-567812345678.jpeg' assert mock_request.called assert mock_request.call_args[1]['timeout'] == 30 @@ -140,11 +173,9 @@ def test_successful_upload(self, mock_uuid, mock_getenv, mock_get_aws_signed_req @patch('lf_toolkit.evaluation.image_upload.uuid.uuid4') def test_upload_with_png(self, mock_uuid, mock_getenv, mock_get_aws_signed_request, mock_request): """Test uploading PNG image with UUID-based filename""" - # Setup mocks mock_uuid.return_value = uuid.UUID('12345678-1234-5678-1234-567812345678') - mock_getenv.return_value = 'https://s3.amazonaws.com/eduvision' + mock_getenv.side_effect = fake_getenv({'S3_BUCKET_URI': 'https://s3.amazonaws.com/eduvision'}) - # Mock the AWS signed request mock_prepared_request = Mock() mock_prepared_request.method = 'PUT' mock_prepared_request.url = 'https://s3.amazonaws.com/eduvision/eduvision/12345678-1234-5678-1234-567812345678.png' @@ -169,7 +200,7 @@ def test_upload_with_png(self, mock_uuid, mock_getenv, mock_get_aws_signed_reque @patch('lf_toolkit.evaluation.image_upload.os.getenv') def test_upload_missing_s3_uri(self, mock_getenv): """Test upload fails when S3_BUCKET_URI is missing""" - mock_getenv.return_value = None + mock_getenv.side_effect = fake_getenv({}) img = Image.new('RGB', (100, 100)) img.format = 'JPEG' @@ -184,9 +215,8 @@ def test_upload_missing_s3_uri(self, mock_getenv): def test_upload_server_error(self, mock_uuid, mock_getenv, mock_get_aws_signed_request, mock_request): """Test upload fails when server returns error""" mock_uuid.return_value = uuid.UUID('12345678-1234-5678-1234-567812345678') - mock_getenv.return_value = 'https://s3.amazonaws.com/bucket' + mock_getenv.side_effect = fake_getenv({'S3_BUCKET_URI': 'https://s3.amazonaws.com/bucket'}) - # Mock the AWS signed request mock_prepared_request = Mock() mock_prepared_request.method = 'PUT' mock_prepared_request.url = 'https://s3.amazonaws.com/bucket/eduvision/12345678-1234-5678-1234-567812345678.jpeg' @@ -217,9 +247,8 @@ def test_upload_server_error(self, mock_uuid, mock_getenv, mock_get_aws_signed_r def test_upload_network_error(self, mock_uuid, mock_getenv, mock_get_aws_signed_request, mock_request): """Test upload fails on network error""" mock_uuid.return_value = uuid.UUID('12345678-1234-5678-1234-567812345678') - mock_getenv.return_value = 'https://s3.amazonaws.com/bucket' + mock_getenv.side_effect = fake_getenv({'S3_BUCKET_URI': 'https://s3.amazonaws.com/bucket'}) - # Mock the AWS signed request mock_prepared_request = Mock() mock_prepared_request.method = 'PUT' mock_prepared_request.url = 'https://s3.amazonaws.com/bucket/eduvision/12345678-1234-5678-1234-567812345678.jpeg' @@ -247,9 +276,8 @@ def test_upload_network_error(self, mock_uuid, mock_getenv, mock_get_aws_signed_ def test_upload_timeout_error(self, mock_uuid, mock_getenv, mock_get_aws_signed_request, mock_request): """Test upload fails on timeout""" mock_uuid.return_value = uuid.UUID('12345678-1234-5678-1234-567812345678') - mock_getenv.return_value = 'https://s3.amazonaws.com/bucket' + mock_getenv.side_effect = fake_getenv({'S3_BUCKET_URI': 'https://s3.amazonaws.com/bucket'}) - # Mock the AWS signed request mock_prepared_request = Mock() mock_prepared_request.method = 'PUT' mock_prepared_request.url = 'https://s3.amazonaws.com/bucket/eduvision/12345678-1234-5678-1234-567812345678.jpeg' @@ -277,9 +305,8 @@ def test_upload_timeout_error(self, mock_uuid, mock_getenv, mock_get_aws_signed_ def test_upload_image_no_format(self, mock_uuid, mock_getenv, mock_get_aws_signed_request, mock_request): """Test upload with image that has no format (defaults to PNG) uses UUID filename""" mock_uuid.return_value = uuid.UUID('12345678-1234-5678-1234-567812345678') - mock_getenv.return_value = 'https://s3.amazonaws.com/bucket/' + mock_getenv.side_effect = fake_getenv({'S3_BUCKET_URI': 'https://s3.amazonaws.com/bucket/'}) - # Mock the AWS signed request mock_prepared_request = Mock() mock_prepared_request.method = 'PUT' mock_prepared_request.url = 'https://s3.amazonaws.com/bucket/eduvision/12345678-1234-5678-1234-567812345678.png' @@ -302,6 +329,111 @@ def test_upload_image_no_format(self, mock_uuid, mock_getenv, mock_get_aws_signe assert result == 'https://s3.amazonaws.com/bucket/eduvision/12345678-1234-5678-1234-567812345678.png' +class TestUploadImageGCS: + """Test suite for upload_image function (GCS backend)""" + + def _mock_storage(self): + mock_blob = Mock() + mock_bucket = Mock() + mock_bucket.blob.return_value = mock_blob + mock_client = Mock() + mock_client.bucket.return_value = mock_bucket + mock_storage = Mock() + mock_storage.Client.return_value = mock_client + return mock_storage, mock_client, mock_bucket, mock_blob + + @patch('lf_toolkit.evaluation.image_upload.os.getenv') + @patch('lf_toolkit.evaluation.image_upload.uuid.uuid4') + def test_successful_gcs_upload(self, mock_uuid, mock_getenv): + mock_uuid.return_value = uuid.UUID('12345678-1234-5678-1234-567812345678') + mock_getenv.side_effect = fake_getenv({'GCS_BUCKET': 'plots-bucket'}) + mock_storage, mock_client, mock_bucket, mock_blob = self._mock_storage() + + img = Image.new('RGB', (100, 100), color='blue') + img.format = 'PNG' + + with patch('lf_toolkit.evaluation.image_upload._gcs_storage', mock_storage): + result = upload_image(img, "evaluatePython") + + assert result == ( + 'https://storage.googleapis.com/plots-bucket/evaluatePython/' + '12345678-1234-5678-1234-567812345678.png' + ) + mock_client.bucket.assert_called_once_with('plots-bucket') + mock_bucket.blob.assert_called_once_with( + 'evaluatePython/12345678-1234-5678-1234-567812345678.png' + ) + _, kwargs = mock_blob.upload_from_string.call_args + assert kwargs['content_type'] == 'image/png' + + @patch('lf_toolkit.evaluation.image_upload.os.getenv') + @patch('lf_toolkit.evaluation.image_upload.uuid.uuid4') + def test_gcs_public_base_url_override(self, mock_uuid, mock_getenv): + mock_uuid.return_value = uuid.UUID('12345678-1234-5678-1234-567812345678') + mock_getenv.side_effect = fake_getenv({ + 'IMAGE_UPLOAD_BACKEND': 'gcs', + 'GCS_BUCKET': 'plots-bucket', + 'GCS_PUBLIC_BASE_URL': 'https://cdn.example.com/', + }) + mock_storage, *_ = self._mock_storage() + + img = Image.new('RGB', (10, 10)) + img.format = 'PNG' + + with patch('lf_toolkit.evaluation.image_upload._gcs_storage', mock_storage): + result = upload_image(img, "evaluatePython") + + assert result == ( + 'https://cdn.example.com/plots-bucket/evaluatePython/' + '12345678-1234-5678-1234-567812345678.png' + ) + + @patch('lf_toolkit.evaluation.image_upload.os.getenv') + def test_gcs_missing_bucket(self, mock_getenv): + mock_getenv.side_effect = fake_getenv({'IMAGE_UPLOAD_BACKEND': 'gcs'}) + mock_storage, *_ = self._mock_storage() + + img = Image.new('RGB', (10, 10)) + img.format = 'PNG' + + with patch('lf_toolkit.evaluation.image_upload._gcs_storage', mock_storage): + with pytest.raises(MissingEnvironmentVariableError): + upload_image(img, "evaluatePython") + + @patch('lf_toolkit.evaluation.image_upload.os.getenv') + def test_gcs_backend_without_dependency(self, mock_getenv): + mock_getenv.side_effect = fake_getenv({ + 'IMAGE_UPLOAD_BACKEND': 'gcs', + 'GCS_BUCKET': 'plots-bucket', + }) + + img = Image.new('RGB', (10, 10)) + img.format = 'PNG' + + with patch('lf_toolkit.evaluation.image_upload._gcs_storage', None): + with pytest.raises(ImageUploadError) as exc_info: + upload_image(img, "evaluatePython") + + assert "google-cloud-storage is not installed" in str(exc_info.value) + + @patch('lf_toolkit.evaluation.image_upload.os.getenv') + @patch('lf_toolkit.evaluation.image_upload.uuid.uuid4') + def test_gcs_upload_error_wrapped(self, mock_uuid, mock_getenv): + mock_uuid.return_value = uuid.UUID('12345678-1234-5678-1234-567812345678') + mock_getenv.side_effect = fake_getenv({'GCS_BUCKET': 'plots-bucket'}) + mock_storage, _, _, mock_blob = self._mock_storage() + mock_blob.upload_from_string.side_effect = RuntimeError('boom') + + img = Image.new('RGB', (10, 10)) + img.format = 'PNG' + + with patch('lf_toolkit.evaluation.image_upload._gcs_storage', mock_storage): + with pytest.raises(ImageUploadError) as exc_info: + upload_image(img, "evaluatePython") + + assert "boom" in str(exc_info.value) + + class TestExceptionHierarchy: """Test suite for custom exception classes"""