feat(component,ai,gemini): add image generation support (#1122)

pinglin · web-flow · commit d986614184cf · 2025-09-24T19:28:50.000+01:00
Because

- The Gemini component lacked support for image generation using the
`gemini-2.5-flash-image-preview` model
- Generated images were being exposed as raw binary data arrays in the
`candidates` field instead of being properly extracted to the `images`
field
- Image processing was duplicated and inefficient in streaming
responses, processing the same images multiple times

This commit

- Adds image generation capability to the Gemini component by extracting
images from `genai.GenerateContentResponse`
- Extracts generated images from API responses and converts them to
`format.Image` objects in the `Images` output field
- Cleans up raw binary `InlineData` from candidates after image
extraction to prevent JSON exposure of binary data
- Optimizes streaming responses by deferring image processing to final
response only, eliminating duplicate processing
- Updates component configuration documentation to clarify the
relationship between `images` field and `candidates` field
diff --git a/pkg/component/ai/gemini/v0/config/tasks.yaml b/pkg/component/ai/gemini/v0/config/tasks.yaml
@@ -1556,19 +1556,22 @@ TASK_CHAT:
           ID of the model to use.
           The value is one of the following:
             `gemini-2.5-pro`: Optimized for enhanced thinking and reasoning, multimodal understanding, advanced coding, and more.
-            `gemini-2.5-flash`: Optimized for Adaptive thinking, cost efficiency.
-            `gemini-2.0-flash-lite`: Optimized for Most cost-efficient model supporting high throughput.
+            `gemini-2.5-flash`: Optimized for adaptive thinking, cost efficiency.
+            `gemini-2.5-flash-lite`: Optimized for most cost-efficient model supporting high throughput.
+            `gemini-2.5-flash-image-preview`: Optimized for precise, conversational image generation and editing.
         type: string
         enum:
           - gemini-2.5-pro
           - gemini-2.5-flash
-          - gemini-2.0-flash-lite
+          - gemini-2.5-flash-lite
+          - gemini-2.5-flash-image-preview
         default: gemini-2.5-flash
         instillCredentialMap:
           values:
             - gemini-2.5-pro
             - gemini-2.5-flash
-            - gemini-2.0-flash-lite
+            - gemini-2.5-flash-lite
+            - gemini-2.5-flash-image-preview
           targets:
             - setup.api-key
       stream:
@@ -1709,20 +1712,31 @@ TASK_CHAT:
         title: Texts
         description: >-
           Simplified text output extracted from candidates. Each string represents the concatenated text content  from the corresponding candidate's parts,
-          including thought processes when `include-thoughts` is enabled.  This field provides easy access to the generated text without needing to traverse
+          including thought processes when `include-thoughts` is enabled. This field provides easy access to the generated text without needing to traverse
           the candidate structure.  Updated in real-time during streaming.
         type: array
         items:
           type: string
-      usage:
+      images:
         uiOrder: 1
+        title: Images
+        description: >-
+          Images output extracted and converted from candidates. This field provides easy access to the generated images as base64-encoded strings.
+          The original binary data is removed from the candidates field to prevent raw binary exposure in JSON output.  This field is only available when
+          the model supports image generation.
+        type: array
+        items:
+          title: Image
+          type: image/webp
+      usage:
+        uiOrder: 2
         title: Usage
         description: >-
           Token usage statistics: prompt tokens, completion tokens, total tokens, etc.
         type: object
         additionalProperties: true
       candidates:
-        uiOrder: 2
+        uiOrder: 3
         title: Candidates
         description: >-
           Complete candidate objects from the model containing rich metadata and structured content.  Each candidate includes safety ratings, finish reason,
@@ -1732,18 +1746,18 @@ TASK_CHAT:
         items:
           $ref: "#/$defs/candidate"
       usage-metadata:
-        uiOrder: 3
+        uiOrder: 4
         $ref: "#/$defs/usage-metadata"
       prompt-feedback:
-        uiOrder: 4
+        uiOrder: 5
         $ref: "#/$defs/prompt-feedback"
       model-version:
-        uiOrder: 5
+        uiOrder: 6
         title: Model Version
         description: The model version used to generate the response.
         type: string
       response-id:
-        uiOrder: 6
+        uiOrder: 7
         title: Response ID
         description: Identifier for this response.
         type: string
diff --git a/pkg/component/ai/gemini/v0/io.go b/pkg/component/ai/gemini/v0/io.go
@@ -59,8 +59,9 @@ func (t TaskChatInput) GetSystemInstruction() *genai.Content { return t.SystemIn
 // TaskChatOutput is the output for the TASK_CHAT task.
 type TaskChatOutput struct {
 	// Flattened chat output properties
-	Texts []string       `instill:"texts"`
-	Usage map[string]any `instill:"usage"`
+	Texts  []string       `instill:"texts"`
+	Images []format.Image `instill:"images"`
+	Usage  map[string]any `instill:"usage"`
 
 	// Use genai types directly with instill tags
 	Candidates     []*genai.Candidate                           `instill:"candidates"`
diff --git a/pkg/component/ai/gemini/v0/task_chat.go b/pkg/component/ai/gemini/v0/task_chat.go
@@ -3,10 +3,13 @@ package gemini
 import (
 	"context"
 	"fmt"
+	"strings"
 
 	"google.golang.org/genai"
 
 	"github.com/instill-ai/pipeline-backend/pkg/component/base"
+	"github.com/instill-ai/pipeline-backend/pkg/data"
+	"github.com/instill-ai/pipeline-backend/pkg/data/format"
 )
 
 func (e *execution) chat(ctx context.Context, job *base.Job) error {
@@ -260,6 +263,7 @@ func (e *execution) mergeResponseChunk(r *genai.GenerateContentResponse, finalRe
 func (e *execution) buildStreamOutput(texts []string, finalResp *genai.GenerateContentResponse) TaskChatOutput {
 	streamOutput := TaskChatOutput{
 		Texts:          texts,
+		Images:         []format.Image{},
 		Usage:          map[string]any{},
 		Candidates:     []*genai.Candidate{},
 		UsageMetadata:  nil,
@@ -281,6 +285,10 @@ func (e *execution) buildStreamOutput(texts []string, finalResp *genai.GenerateC
 			streamOutput.ResponseID = &ri
 		}
 
+		// Note: Image extraction and InlineData cleanup is deferred to renderFinal()
+		// to avoid processing the same images multiple times during streaming.
+		// Streaming responses will have empty Images array until the final response.
+
 		// Build usage map from UsageMetadata if available
 		if finalResp.UsageMetadata != nil {
 			streamOutput.Usage = buildUsageMap(finalResp.UsageMetadata)
@@ -348,6 +356,7 @@ func buildUsageMap(metadata *genai.GenerateContentResponseUsageMetadata) map[str
 func renderFinal(resp *genai.GenerateContentResponse, texts []string) TaskChatOutput {
 	out := TaskChatOutput{
 		Texts:          []string{},
+		Images:         []format.Image{},
 		Usage:          map[string]any{},
 		Candidates:     []*genai.Candidate{},
 		UsageMetadata:  nil,
@@ -385,6 +394,31 @@ func renderFinal(resp *genai.GenerateContentResponse, texts []string) TaskChatOu
 		}
 		out.Texts = acc
 	}
+
+	// Extract generated images from candidates and clean up InlineData to prevent raw binary exposure
+	if len(resp.Candidates) > 0 {
+		images := make([]format.Image, 0)
+		for _, c := range resp.Candidates {
+			if c.Content != nil {
+				for _, p := range c.Content.Parts {
+					if p != nil && p.InlineData != nil && strings.Contains(strings.ToLower(p.InlineData.MIMEType), "image") {
+						// Convert blob data to format.Image using the standard data package approach
+						// Normalize MIME type and use the existing NewImageFromBytes function
+						normalizedMimeType := strings.ToLower(strings.TrimSpace(strings.Split(p.InlineData.MIMEType, ";")[0]))
+						img, err := data.NewImageFromBytes(p.InlineData.Data, normalizedMimeType, "", true)
+						if err == nil {
+							images = append(images, img)
+						}
+						// Clean up InlineData to prevent raw binary data from being exposed in JSON output
+						// The binary data is already extracted and converted to format.Image above
+						p.InlineData = nil
+					}
+				}
+			}
+		}
+		out.Images = images
+	}
+
 	if resp.UsageMetadata != nil {
 		out.Usage = buildUsageMap(resp.UsageMetadata)
 	}
diff --git a/pkg/component/ai/gemini/v0/task_chat_test.go b/pkg/component/ai/gemini/v0/task_chat_test.go
@@ -5,6 +5,7 @@ import (
 	"encoding/base64"
 	"fmt"
 	"os"
+	"strings"
 	"testing"
 	"time"
 
@@ -1339,3 +1340,154 @@ func TestChatPerformanceOptimizations(t *testing.T) {
 		c.Check(videoTimeout > documentTimeout, qt.IsTrue)
 	})
 }
+
+func TestImageGeneration(t *testing.T) {
+	t.Parallel()
+
+	t.Run("image MIME type detection", func(t *testing.T) {
+		c := qt.New(t)
+
+		// Test the standard approach used in the codebase
+		// Test valid image MIME types
+		c.Check(strings.Contains(strings.ToLower("image/png"), "image"), qt.Equals, true)
+		c.Check(strings.Contains(strings.ToLower("image/jpeg"), "image"), qt.Equals, true)
+		c.Check(strings.Contains(strings.ToLower("image/gif"), "image"), qt.Equals, true)
+		c.Check(strings.Contains(strings.ToLower("image/webp"), "image"), qt.Equals, true)
+		c.Check(strings.Contains(strings.ToLower("IMAGE/PNG"), "image"), qt.Equals, true) // Case insensitive
+
+		// Test non-image MIME types
+		c.Check(strings.Contains(strings.ToLower("text/plain"), "image"), qt.Equals, false)
+		c.Check(strings.Contains(strings.ToLower("application/json"), "image"), qt.Equals, false)
+		c.Check(strings.Contains(strings.ToLower("video/mp4"), "image"), qt.Equals, false)
+		c.Check(strings.Contains(strings.ToLower("audio/wav"), "image"), qt.Equals, false)
+	})
+
+	t.Run("renderFinal with generated images", func(t *testing.T) {
+		c := qt.New(t)
+
+		// Create a mock response with generated images
+		// Use a simple 1x1 PNG image data
+		pngData := []byte{
+			0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A, // PNG signature
+			0x00, 0x00, 0x00, 0x0D, 0x49, 0x48, 0x44, 0x52, // IHDR chunk
+			0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x01, // 1x1 dimensions
+			0x08, 0x02, 0x00, 0x00, 0x00, 0x90, 0x77, 0x53, // bit depth, color type, etc.
+			0xDE, 0x00, 0x00, 0x00, 0x0C, 0x49, 0x44, 0x41, // IDAT chunk
+			0x54, 0x08, 0x99, 0x01, 0x01, 0x01, 0x00, 0x00, // pixel data
+			0xFE, 0xFF, 0x00, 0x00, 0x02, 0x00, 0x01, 0xE5, // checksum
+			0x27, 0xDE, 0xFC, 0x00, 0x00, 0x00, 0x00, 0x49, // IEND chunk
+			0x45, 0x4E, 0x44, 0xAE, 0x42, 0x60, 0x82,
+		}
+
+		resp := &genai.GenerateContentResponse{
+			Candidates: []*genai.Candidate{
+				{
+					Content: &genai.Content{
+						Parts: []*genai.Part{
+							{Text: "Here's your generated image:"},
+							{InlineData: &genai.Blob{MIMEType: "image/png", Data: pngData}},
+						},
+					},
+				},
+			},
+		}
+
+		result := renderFinal(resp, nil)
+
+		// Check that text was extracted
+		c.Check(result.Texts, qt.HasLen, 1)
+		c.Check(result.Texts[0], qt.Equals, "Here's your generated image:")
+
+		// Check that images were extracted
+		c.Check(result.Images, qt.HasLen, 1)
+		c.Check(result.Images[0].ContentType().String(), qt.Equals, "image/png")
+	})
+
+	t.Run("buildStreamOutput with generated images", func(t *testing.T) {
+		c := qt.New(t)
+
+		// Mock execution for the method receiver
+		e := &execution{}
+
+		// Use a simple 1x1 PNG image data
+		pngData := []byte{
+			0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A, // PNG signature
+			0x00, 0x00, 0x00, 0x0D, 0x49, 0x48, 0x44, 0x52, // IHDR chunk
+			0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x01, // 1x1 dimensions
+			0x08, 0x02, 0x00, 0x00, 0x00, 0x90, 0x77, 0x53, // bit depth, color type, etc.
+			0xDE, 0x00, 0x00, 0x00, 0x0C, 0x49, 0x44, 0x41, // IDAT chunk
+			0x54, 0x08, 0x99, 0x01, 0x01, 0x01, 0x00, 0x00, // pixel data
+			0xFE, 0xFF, 0x00, 0x00, 0x02, 0x00, 0x01, 0xE5, // checksum
+			0x27, 0xDE, 0xFC, 0x00, 0x00, 0x00, 0x00, 0x49, // IEND chunk
+			0x45, 0x4E, 0x44, 0xAE, 0x42, 0x60, 0x82,
+		}
+
+		texts := []string{"Generated image:"}
+		finalResp := &genai.GenerateContentResponse{
+			Candidates: []*genai.Candidate{
+				{
+					Content: &genai.Content{
+						Parts: []*genai.Part{
+							{Text: "Generated image:"},
+							{InlineData: &genai.Blob{MIMEType: "image/png", Data: pngData}},
+						},
+					},
+				},
+			},
+		}
+
+		result := e.buildStreamOutput(texts, finalResp)
+
+		// Check that texts are preserved
+		c.Check(result.Texts, qt.DeepEquals, texts)
+
+		// Check that images are NOT extracted during streaming (deferred to renderFinal)
+		c.Check(result.Images, qt.HasLen, 0)
+	})
+
+	t.Run("renderFinal with mixed content", func(t *testing.T) {
+		c := qt.New(t)
+
+		// Use a simple 1x1 PNG image data
+		pngData := []byte{
+			0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A, // PNG signature
+			0x00, 0x00, 0x00, 0x0D, 0x49, 0x48, 0x44, 0x52, // IHDR chunk
+			0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x01, // 1x1 dimensions
+			0x08, 0x02, 0x00, 0x00, 0x00, 0x90, 0x77, 0x53, // bit depth, color type, etc.
+			0xDE, 0x00, 0x00, 0x00, 0x0C, 0x49, 0x44, 0x41, // IDAT chunk
+			0x54, 0x08, 0x99, 0x01, 0x01, 0x01, 0x00, 0x00, // pixel data
+			0xFE, 0xFF, 0x00, 0x00, 0x02, 0x00, 0x01, 0xE5, // checksum
+			0x27, 0xDE, 0xFC, 0x00, 0x00, 0x00, 0x00, 0x49, // IEND chunk
+			0x45, 0x4E, 0x44, 0xAE, 0x42, 0x60, 0x82,
+		}
+
+		// Create a response with text, images, and non-image content
+		resp := &genai.GenerateContentResponse{
+			Candidates: []*genai.Candidate{
+				{
+					Content: &genai.Content{
+						Parts: []*genai.Part{
+							{Text: "Here are your images: "},
+							{InlineData: &genai.Blob{MIMEType: "image/png", Data: pngData}},
+							{Text: " and "},
+							{InlineData: &genai.Blob{MIMEType: "image/png", Data: pngData}}, // Use PNG data with PNG MIME for valid image
+							{Text: " Done!"},
+							{InlineData: &genai.Blob{MIMEType: "audio/wav", Data: []byte("audio-data")}}, // Non-image, should be ignored
+						},
+					},
+				},
+			},
+		}
+
+		result := renderFinal(resp, nil)
+
+		// Check that all text parts were concatenated
+		c.Check(result.Texts, qt.HasLen, 1)
+		c.Check(result.Texts[0], qt.Equals, "Here are your images:  and  Done!")
+
+		// Check that only image parts were extracted (audio should be ignored)
+		c.Check(result.Images, qt.HasLen, 2)
+		c.Check(result.Images[0].ContentType().String(), qt.Equals, "image/png")
+		c.Check(result.Images[1].ContentType().String(), qt.Equals, "image/png")
+	})
+}
