feat(translator): add OpenAI image generation translation layer

- Implement OpenAI-to-OpenAI image generation request/response translation
- Add comprehensive error handling for image generation failures
- Update translator interface to support image generation endpoints
- Include test coverage for translation functionality

Signed-off-by: Hrushikesh Patil <[email protected]>

Author: nutanix-Hrushikesh

internal/extproc/translator/imagegeneration_error.go

@@ -0,0 +1,17 @@
+// Copyright Envoy AI Gateway Authors
+// SPDX-License-Identifier: Apache-2.0
+// The full text of the Apache license is available in the LICENSE file at
+// the root of the repo.
+
+package translator
+
+// ImageGenerationError represents an error response from the OpenAI Images API.
+// This schema matches OpenAI's documented error wire format.
+type ImageGenerationError struct {
+	Error struct {
+		Type    string  `json:"type"`
+		Message string  `json:"message"`
+		Code    *string `json:"code,omitempty"`
+		Param   *string `json:"param,omitempty"`
+	} `json:"error"`
+}

internal/extproc/translator/imagegeneration_openai_openai.go

@@ -10,14 +10,13 @@ import (
 	"fmt"
 	"io"
 	"path"
-	"strconv"
 
 	corev3 "github.com/envoyproxy/go-control-plane/envoy/config/core/v3"
 	extprocv3 "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3"
 	"github.com/tidwall/sjson"
 
-	"github.com/envoyproxy/ai-gateway/internal/apischema/openai"
 	tracing "github.com/envoyproxy/ai-gateway/internal/tracing/api"
+	openaisdk "github.com/openai/openai-go/v2"
 )
 
 // NewImageGenerationOpenAIToOpenAITranslator implements [Factory] for OpenAI to OpenAI image generation translation.
@@ -33,7 +32,7 @@ type openAIToOpenAIImageGenerationTranslator struct {
 }
 
 // RequestBody implements [ImageGenerationTranslator.RequestBody].
-func (o *openAIToOpenAIImageGenerationTranslator) RequestBody(original []byte, req *openai.ImageGenerationRequest, forceBodyMutation bool) (
+func (o *openAIToOpenAIImageGenerationTranslator) RequestBody(original []byte, req *openaisdk.ImageGenerateParams, forceBodyMutation bool) (
 	headerMutation *extprocv3.HeaderMutation, bodyMutation *extprocv3.BodyMutation, err error,
 ) {
 	var newBody []byte
@@ -63,10 +62,8 @@ func (o *openAIToOpenAIImageGenerationTranslator) RequestBody(original []byte, r
 		bodyMutation = &extprocv3.BodyMutation{
 			Mutation: &extprocv3.BodyMutation_Body{Body: newBody},
 		}
-		headerMutation.SetHeaders = append(headerMutation.SetHeaders, &corev3.HeaderValueOption{Header: &corev3.HeaderValue{
-			Key:      "content-length",
-			RawValue: []byte(strconv.Itoa(len(newBody))),
-		}})
+		// Note: content-length header is set via dynamic metadata in the processor
+		// to avoid conflicts with Envoy's REPLACE_AND_CONTINUE processing mode
 	}
 	return
 }
@@ -79,12 +76,12 @@ func (o *openAIToOpenAIImageGenerationTranslator) ResponseError(respHeaders map[
 ) {
 	statusCode := respHeaders[statusHeaderName]
 	if v, ok := respHeaders[contentTypeHeaderName]; ok && v != jsonContentType {
-		var openaiError openai.ImageGenerationError
+		var openaiError ImageGenerationError
 		buf, err := io.ReadAll(body)
 		if err != nil {
 			return nil, nil, fmt.Errorf("failed to read error body: %w", err)
 		}
-		openaiError = openai.ImageGenerationError{
+		openaiError = ImageGenerationError{
 			Error: struct {
 				Type    string  `json:"type"`
 				Message string  `json:"message"`
@@ -102,6 +99,11 @@ func (o *openAIToOpenAIImageGenerationTranslator) ResponseError(respHeaders map[
 			return nil, nil, fmt.Errorf("failed to marshal error body: %w", err)
 		}
 		headerMutation = &extprocv3.HeaderMutation{}
+		// Ensure downstream sees a JSON error payload
+		headerMutation.SetHeaders = append(headerMutation.SetHeaders, &corev3.HeaderValueOption{Header: &corev3.HeaderValue{
+			Key:      contentTypeHeaderName,
+			RawValue: []byte(jsonContentType),
+		}})
 		setContentLength(headerMutation, mut.Body)
 		return headerMutation, &extprocv3.BodyMutation{Mutation: mut}, nil
 	}
@@ -117,22 +119,46 @@ func (o *openAIToOpenAIImageGenerationTranslator) ResponseHeaders(map[string]str
 func (o *openAIToOpenAIImageGenerationTranslator) ResponseBody(_ map[string]string, body io.Reader, _ bool) (
 	headerMutation *extprocv3.HeaderMutation, bodyMutation *extprocv3.BodyMutation, tokenUsage LLMTokenUsage, imageMetadata ImageGenerationMetadata, err error,
 ) {
-	resp := &openai.ImageGenerationResponse{}
-	if err := json.NewDecoder(body).Decode(&resp); err != nil {
+	// Read the entire response body first to debug any issues
+	bodyBytes, err := io.ReadAll(body)
+	if err != nil {
+		return nil, nil, tokenUsage, imageMetadata, fmt.Errorf("failed to read response body: %w", err)
+	}
+
+	// Debug logging for response body content
+	bodyPreview := string(bodyBytes)
+	if len(bodyPreview) > 200 {
+		bodyPreview = bodyPreview[:200] + "..."
+	}
+	fmt.Printf("DEBUG: Image generation translator received body - Length: %d, Preview: %s\n", len(bodyBytes), bodyPreview)
+
+	// Check if body looks like JSON
+	if len(bodyBytes) > 0 && bodyBytes[0] != '{' && bodyBytes[0] != '[' {
+		previewLen := 10
+		if len(bodyBytes) < previewLen {
+			previewLen = len(bodyBytes)
+		}
+		fmt.Printf("DEBUG: Body does not start with JSON character. First %d bytes: %v\n", previewLen, bodyBytes[:previewLen])
+	}
+
+	// Decode using OpenAI SDK v2 schema to avoid drift.
+	resp := &openaisdk.ImagesResponse{}
+	if err := json.Unmarshal(bodyBytes, &resp); err != nil {
+		fmt.Printf("DEBUG: JSON unmarshal failed - Error: %v, Body preview: %s\n", err, bodyPreview)
 		return nil, nil, tokenUsage, imageMetadata, fmt.Errorf("failed to unmarshal body: %w", err)
 	}
 
 	// Populate token usage if provided (GPT-Image-1); otherwise remain zero.
-	if resp.Usage != nil {
+	if resp.JSON.Usage.Valid() {
 		tokenUsage.InputTokens = uint32(resp.Usage.InputTokens)   //nolint:gosec
 		tokenUsage.OutputTokens = uint32(resp.Usage.OutputTokens) //nolint:gosec
 		tokenUsage.TotalTokens = uint32(resp.Usage.TotalTokens)   //nolint:gosec
 	}
 
-	// Extract image generation metadata for metrics
+	// Extract image generation metadata for metrics (model may be absent in SDK response)
 	imageMetadata.ImageCount = len(resp.Data)
-	imageMetadata.Model = resp.Model
-	imageMetadata.Size = resp.Size
+	imageMetadata.Model = ""
+	imageMetadata.Size = string(resp.Size)
 
 	return
 }

internal/extproc/translator/imagegeneration_openai_openai_test.go

@@ -12,12 +12,12 @@ import (
 
 	"github.com/stretchr/testify/require"
 
-	"github.com/envoyproxy/ai-gateway/internal/apischema/openai"
+	openaisdk "github.com/openai/openai-go/v2"
 )
 
 func TestOpenAIToOpenAIImageTranslator_RequestBody_ModelOverrideAndPath(t *testing.T) {
 	tr := NewImageGenerationOpenAIToOpenAITranslator("v1", "gpt-image-1")
-	req := &openai.ImageGenerationRequest{Model: openai.ModelDalle3, Prompt: "a cat"}
+	req := &openaisdk.ImageGenerateParams{Model: openaisdk.ImageModelDallE3, Prompt: "a cat"}
 	original, _ := json.Marshal(req)
 
 	hm, bm, err := tr.RequestBody(original, req, false)
@@ -30,14 +30,14 @@ func TestOpenAIToOpenAIImageTranslator_RequestBody_ModelOverrideAndPath(t *testi
 
 	require.NotNil(t, bm)
 	mutated := bm.GetBody()
-	var got openai.ImageGenerationRequest
+	var got openaisdk.ImageGenerateParams
 	require.NoError(t, json.Unmarshal(mutated, &got))
-	require.Equal(t, "gpt-image-1", got.Model)
+	require.Equal(t, "gpt-image-1", string(got.Model))
 }
 
 func TestOpenAIToOpenAIImageTranslator_RequestBody_ForceMutation(t *testing.T) {
 	tr := NewImageGenerationOpenAIToOpenAITranslator("v1", "")
-	req := &openai.ImageGenerationRequest{Model: openai.ModelDalle2, Prompt: "a cat"}
+	req := &openaisdk.ImageGenerateParams{Model: openaisdk.ImageModelDallE2, Prompt: "a cat"}
 	original, _ := json.Marshal(req)
 
 	hm, bm, err := tr.RequestBody(original, req, true)
@@ -65,14 +65,14 @@ func TestOpenAIToOpenAIImageTranslator_ResponseError_NonJSON(t *testing.T) {
 	require.NotNil(t, bm)
 
 	// Body should be OpenAI error JSON
-	var got openai.ImageGenerationError
+	var got ImageGenerationError
 	require.NoError(t, json.Unmarshal(bm.GetBody(), &got))
 	require.Equal(t, openAIBackendError, got.Error.Type)
 }
 
 func TestOpenAIToOpenAIImageTranslator_ResponseBody_OK(t *testing.T) {
 	tr := NewImageGenerationOpenAIToOpenAITranslator("v1", "")
-	resp := &openai.ImageGenerationResponse{Model: openai.ModelDalle3}
+	resp := &openaisdk.ImagesResponse{Size: openaisdk.ImagesResponseSize1024x1024}
 	buf, _ := json.Marshal(resp)
 	hm, bm, usage, metadata, err := tr.ResponseBody(map[string]string{}, bytes.NewReader(buf), true)
 	require.NoError(t, err)
@@ -81,9 +81,6 @@ func TestOpenAIToOpenAIImageTranslator_ResponseBody_OK(t *testing.T) {
 	require.Equal(t, uint32(0), usage.InputTokens)
 	require.Equal(t, uint32(0), usage.TotalTokens)
 	require.Equal(t, 0, metadata.ImageCount)
-	require.Equal(t, openai.ModelDalle3, metadata.Model)
-	require.Equal(t, "", metadata.Size)
+	require.Equal(t, "", metadata.Model)
+	require.Equal(t, string(openaisdk.ImagesResponseSize1024x1024), metadata.Size)
 }
-
-// Ensure the helper compiles in this package (not used directly but keeps imports clean in case of future changes).
-// (no-op local stub removed; production helper exists in package translator)

internal/extproc/translator/translator.go

@@ -17,6 +17,7 @@ import (
 	"github.com/envoyproxy/ai-gateway/internal/apischema/openai"
 	"github.com/envoyproxy/ai-gateway/internal/internalapi"
 	tracing "github.com/envoyproxy/ai-gateway/internal/tracing/api"
+	openaisdk "github.com/openai/openai-go/v2"
 )
 
 const (
@@ -240,10 +241,10 @@ var SJSONOptions = &sjson.Options{
 type ImageGenerationTranslator interface {
 	// RequestBody translates the request body.
 	// 	- raw is the raw request body.
-	// 	- body is the request body parsed into the [openai.ImageGenerationRequest].
+	// 	- body is the request body parsed into the OpenAI SDK [openaisdk.ImageGenerateParams].
 	//	- forceBodyMutation is true if the translator should always mutate the body, even if no changes are made.
 	//	- This returns headerMutation and bodyMutation that can be nil to indicate no mutation.
-	RequestBody(raw []byte, body *openai.ImageGenerationRequest, forceBodyMutation bool) (
+	RequestBody(raw []byte, body *openaisdk.ImageGenerateParams, forceBodyMutation bool) (
 		headerMutation *extprocv3.HeaderMutation,
 		bodyMutation *extprocv3.BodyMutation,
 		err error,