Added support for Google Cloud Storage in image uploads, including environment-based backend selection. Expanded related tests and updated documentation.

Author: m-messer

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):

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

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]