feat(google_genai): add generate_images tracing with Attachment support

Add tracing for the Google GenAI generate_images API (both sync and
async). Generated image bytes are wrapped in Attachment objects so the
SDK uploads them to object storage and the Braintrust UI can render
them, following the same pattern used for input-side inline_data in
both the google_genai and adk integrations.

Key changes:
- ModelsGenerateImagesPatcher and AsyncModelsGenerateImagesPatcher
- _generate_images_wrapper / _async_generate_images_wrapper
- Output includes image metadata (size, mime type, safety attrs) plus
  an Attachment per generated image in {"image_url": {"url": attachment}}
- VCR before_record_response sanitizes image bytes in cassettes
- Sync and async test cases validate span structure, metrics, and
  Attachment objects

Author: AbhiPrasad

py/src/braintrust/integrations/google_genai/cassettes/test_generate_images.yaml

@@ -0,0 +1,57 @@
+interactions:
+- request:
+    body: '{"instances": [{"prompt": "A watercolor fox in a forest"}], "parameters":
+      {"sampleCount": 1, "aspectRatio": "1:1", "safetySetting": "BLOCK_LOW_AND_ABOVE",
+      "includeRaiReason": true}}'
+    headers:
+      Accept:
+      - '*/*'
+      Accept-Encoding:
+      - gzip, deflate, zstd
+      Connection:
+      - keep-alive
+      Content-Length:
+      - '181'
+      Content-Type:
+      - application/json
+      Host:
+      - generativelanguage.googleapis.com
+      user-agent:
+      - google-genai-sdk/1.66.0 gl-python/3.13.3
+      x-goog-api-client:
+      - google-genai-sdk/1.66.0 gl-python/3.13.3
+    method: POST
+    uri: https://generativelanguage.googleapis.com/v1beta/models/imagen-4.0-fast-generate-001:predict
+  response:
+    body:
+      string: '{"predictions": [{"bytesBase64Encoded": "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8DwHwAFBQIAX8jx0gAAAABJRU5ErkJggg==",
+        "mimeType": "image/png"}]}'
+    headers:
+      Alt-Svc:
+      - h3=":443"; ma=2592000,h3-29=":443"; ma=2592000
+      Content-Type:
+      - application/json; charset=UTF-8
+      Date:
+      - Fri, 27 Mar 2026 01:18:45 GMT
+      Server:
+      - scaffolding on HTTPServer2
+      Server-Timing:
+      - gfet4t7; dur=4159
+      Transfer-Encoding:
+      - chunked
+      Vary:
+      - Origin
+      - X-Origin
+      - Referer
+      X-Content-Type-Options:
+      - nosniff
+      X-Frame-Options:
+      - SAMEORIGIN
+      X-XSS-Protection:
+      - '0'
+      content-length:
+      - '2622509'
+    status:
+      code: 200
+      message: OK
+version: 1

py/src/braintrust/integrations/google_genai/cassettes/test_generate_images_async.yaml

@@ -0,0 +1,47 @@
+interactions:
+- request:
+    body: '{"instances": [{"prompt": "A watercolor fox in a forest"}], "parameters":
+      {"sampleCount": 1, "aspectRatio": "1:1", "safetySetting": "BLOCK_LOW_AND_ABOVE",
+      "includeRaiReason": true}}'
+    headers:
+      Content-Type:
+      - application/json
+      user-agent:
+      - google-genai-sdk/1.66.0 gl-python/3.13.3
+      x-goog-api-client:
+      - google-genai-sdk/1.66.0 gl-python/3.13.3
+    method: POST
+    uri: https://generativelanguage.googleapis.com/v1beta/models/imagen-4.0-fast-generate-001:predict
+  response:
+    body:
+      string: '{"predictions": [{"bytesBase64Encoded": "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8DwHwAFBQIAX8jx0gAAAABJRU5ErkJggg==",
+        "mimeType": "image/png"}]}'
+    headers:
+      Alt-Svc:
+      - h3=":443"; ma=2592000,h3-29=":443"; ma=2592000
+      Content-Type:
+      - application/json; charset=UTF-8
+      Date:
+      - Fri, 27 Mar 2026 01:18:49 GMT
+      Server:
+      - scaffolding on HTTPServer2
+      Server-Timing:
+      - gfet4t7; dur=3224
+      Transfer-Encoding:
+      - chunked
+      Vary:
+      - Origin
+      - X-Origin
+      - Referer
+      X-Content-Type-Options:
+      - nosniff
+      X-Frame-Options:
+      - SAMEORIGIN
+      X-XSS-Protection:
+      - '0'
+      content-length:
+      - '2750553'
+    status:
+      code: 200
+      message: OK
+version: 1

py/src/braintrust/integrations/google_genai/integration.py

@@ -8,9 +8,11 @@
     AsyncModelsEmbedContentPatcher,
     AsyncModelsGenerateContentPatcher,
     AsyncModelsGenerateContentStreamPatcher,
+    AsyncModelsGenerateImagesPatcher,
     ModelsEmbedContentPatcher,
     ModelsGenerateContentPatcher,
     ModelsGenerateContentStreamPatcher,
+    ModelsGenerateImagesPatcher,
 )
 
 
@@ -26,7 +28,9 @@ class GoogleGenAIIntegration(BaseIntegration):
         ModelsGenerateContentPatcher,
         ModelsGenerateContentStreamPatcher,
         ModelsEmbedContentPatcher,
+        ModelsGenerateImagesPatcher,
         AsyncModelsGenerateContentPatcher,
         AsyncModelsGenerateContentStreamPatcher,
         AsyncModelsEmbedContentPatcher,
+        AsyncModelsGenerateImagesPatcher,
     )

py/src/braintrust/integrations/google_genai/patchers.py

@@ -6,9 +6,11 @@
     _async_embed_content_wrapper,
     _async_generate_content_stream_wrapper,
     _async_generate_content_wrapper,
+    _async_generate_images_wrapper,
     _embed_content_wrapper,
     _generate_content_stream_wrapper,
     _generate_content_wrapper,
+    _generate_images_wrapper,
 )
 
 
@@ -44,6 +46,15 @@ class ModelsEmbedContentPatcher(FunctionWrapperPatcher):
     wrapper = _embed_content_wrapper
 
 
+class ModelsGenerateImagesPatcher(FunctionWrapperPatcher):
+    """Patch ``Models.generate_images`` for tracing."""
+
+    name = "google_genai.models.generate_images"
+    target_module = "google.genai.models"
+    target_path = "Models.generate_images"
+    wrapper = _generate_images_wrapper
+
+
 # ---------------------------------------------------------------------------
 # Async Models patchers
 # ---------------------------------------------------------------------------
@@ -74,3 +85,12 @@ class AsyncModelsEmbedContentPatcher(FunctionWrapperPatcher):
     target_module = "google.genai.models"
     target_path = "AsyncModels.embed_content"
     wrapper = _async_embed_content_wrapper
+
+
+class AsyncModelsGenerateImagesPatcher(FunctionWrapperPatcher):
+    """Patch ``AsyncModels.generate_images`` for tracing."""
+
+    name = "google_genai.async_models.generate_images"
+    target_module = "google.genai.models"
+    target_path = "AsyncModels.generate_images"
+    wrapper = _async_generate_images_wrapper

py/src/braintrust/integrations/google_genai/test_google_genai.py

@@ -1,10 +1,13 @@
+import gzip
+import json
 import os
 import time
 from pathlib import Path
 
 import pytest
 from braintrust import logger
 from braintrust.integrations.google_genai import setup_genai
+from braintrust.logger import Attachment
 from braintrust.test_helpers import init_test_logger
 from braintrust.wrappers.test_utils import verify_autoinstrument_script
 from google.genai import types
@@ -14,7 +17,59 @@
 PROJECT_NAME = "test-genai-app"
 MODEL = "gemini-2.0-flash-001"
 EMBEDDING_MODEL = "gemini-embedding-001"
+IMAGE_MODEL = "imagen-4.0-fast-generate-001"
 FIXTURES_DIR = Path(__file__).parent.parent.parent.parent.parent / "internal/golden/fixtures"
+TINY_PNG_BASE64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8DwHwAFBQIAX8jx0gAAAABJRU5ErkJggg=="
+
+
+def _sanitize_generate_images_body(value):
+    if isinstance(value, dict):
+        return {
+            key: (
+                TINY_PNG_BASE64
+                if key == "bytesBase64Encoded" and isinstance(val, str)
+                else _sanitize_generate_images_body(val)
+            )
+            for key, val in value.items()
+        }
+    if isinstance(value, list):
+        return [_sanitize_generate_images_body(item) for item in value]
+    return value
+
+
+def _sanitize_generate_images_response(response):
+    body = response.get("body", {})
+    payload = body.get("string")
+    if not payload:
+        return response
+
+    is_bytes = isinstance(payload, bytes)
+    is_gzipped = False
+
+    if is_bytes:
+        raw_payload = payload
+        if raw_payload[:2] == b"\x1f\x8b":
+            raw_payload = gzip.decompress(raw_payload)
+            is_gzipped = True
+        payload = raw_payload.decode("utf-8")
+
+    try:
+        parsed = json.loads(payload)
+    except Exception:
+        return response
+
+    sanitized = _sanitize_generate_images_body(parsed)
+    if sanitized == parsed:
+        return response
+
+    sanitized_payload = json.dumps(sanitized)
+    if is_bytes:
+        body["string"] = (
+            gzip.compress(sanitized_payload.encode("utf-8")) if is_gzipped else sanitized_payload.encode("utf-8")
+        )
+    else:
+        body["string"] = sanitized_payload
+    return response
 
 
 @pytest.fixture(scope="module")
@@ -27,14 +82,19 @@ def before_record_request(request):
         request.method = request.method.upper()
         return request
 
+    def before_record_response(response):
+        return _sanitize_generate_images_response(response)
+
     return {
         "record_mode": record_mode,
+        "decode_compressed_response": True,
         "filter_headers": [
             "authorization",
             "x-api-key",
             "x-goog-api-key",
         ],
         "before_record_request": before_record_request,
+        "before_record_response": before_record_response,
     }
 
 
@@ -669,6 +729,105 @@ def test_attachment_in_config(memory_logger):
     assert copied["temperature"] == 0.5
 
 
+@pytest.mark.vcr
+def test_generate_images(memory_logger):
+    assert not memory_logger.pop()
+
+    client = Client()
+    start = time.time()
+
+    response = client.models.generate_images(
+        model=IMAGE_MODEL,
+        prompt="A watercolor fox in a forest",
+        config=types.GenerateImagesConfig(
+            number_of_images=1,
+            aspect_ratio="1:1",
+            safety_filter_level="BLOCK_LOW_AND_ABOVE",
+            include_rai_reason=True,
+        ),
+    )
+    end = time.time()
+
+    assert len(response.generated_images) == 1
+    assert response.generated_images[0].image
+    assert response.generated_images[0].image.image_bytes
+
+    spans = memory_logger.pop()
+    assert len(spans) == 1
+    span = spans[0]
+    assert span["metadata"]["model"] == IMAGE_MODEL
+    assert span["input"]["prompt"] == "A watercolor fox in a forest"
+    assert span["input"]["config"]["number_of_images"] == 1
+    assert span["input"]["config"]["aspect_ratio"] == "1:1"
+    assert span["input"]["config"]["safety_filter_level"] == "BLOCK_LOW_AND_ABOVE"
+    assert span["input"]["config"]["include_rai_reason"] is True
+    assert span["output"]["generated_images_count"] == 1
+    generated_image = span["output"]["generated_images"][0]
+    assert generated_image["image_size_bytes"] > 0
+    assert generated_image["mime_type"] in {"image/png", "image/jpeg", "image/webp"}
+
+    # Verify the image bytes are stored as an Attachment for upload to object storage
+    assert "image_url" in generated_image
+    attachment = generated_image["image_url"]["url"]
+    assert isinstance(attachment, Attachment)
+    assert attachment.reference["type"] == "braintrust_attachment"
+    assert attachment.reference["content_type"] == generated_image["mime_type"]
+    assert attachment.reference["filename"].startswith("generated_image_")
+    assert attachment.reference["key"]
+
+    _assert_timing_metrics_are_valid(span["metrics"], start, end)
+
+
+@pytest.mark.vcr
+@pytest.mark.asyncio
+async def test_generate_images_async(memory_logger):
+    assert not memory_logger.pop()
+
+    client = Client()
+    start = time.time()
+
+    response = await client.aio.models.generate_images(
+        model=IMAGE_MODEL,
+        prompt="A watercolor fox in a forest",
+        config=types.GenerateImagesConfig(
+            number_of_images=1,
+            aspect_ratio="1:1",
+            safety_filter_level="BLOCK_LOW_AND_ABOVE",
+            include_rai_reason=True,
+        ),
+    )
+    end = time.time()
+
+    assert len(response.generated_images) == 1
+    assert response.generated_images[0].image
+    assert response.generated_images[0].image.image_bytes
+
+    spans = memory_logger.pop()
+    assert len(spans) == 1
+    span = spans[0]
+    assert span["metadata"]["model"] == IMAGE_MODEL
+    assert span["input"]["prompt"] == "A watercolor fox in a forest"
+    assert span["input"]["config"]["number_of_images"] == 1
+    assert span["input"]["config"]["aspect_ratio"] == "1:1"
+    assert span["input"]["config"]["safety_filter_level"] == "BLOCK_LOW_AND_ABOVE"
+    assert span["input"]["config"]["include_rai_reason"] is True
+    assert span["output"]["generated_images_count"] == 1
+    generated_image = span["output"]["generated_images"][0]
+    assert generated_image["image_size_bytes"] > 0
+    assert generated_image["mime_type"] in {"image/png", "image/jpeg", "image/webp"}
+
+    # Verify the image bytes are stored as an Attachment for upload to object storage
+    assert "image_url" in generated_image
+    attachment = generated_image["image_url"]["url"]
+    assert isinstance(attachment, Attachment)
+    assert attachment.reference["type"] == "braintrust_attachment"
+    assert attachment.reference["content_type"] == generated_image["mime_type"]
+    assert attachment.reference["filename"].startswith("generated_image_")
+    assert attachment.reference["key"]
+
+    _assert_timing_metrics_are_valid(span["metrics"], start, end)
+
+
 def test_nested_attachments_in_contents(memory_logger):
     """Test that nested attachments in contents are preserved."""
     from braintrust.bt_json import bt_safe_deep_copy