fix: updates image generation request tracing to match latest OpenInference (#1520)

**Description**

updates image generation request tracing to match latest OpenInference
- Span name changed to ImagesResponse
- Add output.mime_type and output.value attributes
- Remove llm.model_name from request attributes (now in
invocation_parameters)
- Fix missing RecordResponse call in translator

**Related Issues/PRs (if applicable)**

#1519

Signed-off-by: Adrian Cole <[email protected]>
Co-authored-by: Ignasi Barrera <[email protected]>

Author: codefromthecrypt

internal/extproc/imagegeneration_processor.go

@@ -416,6 +416,7 @@ func (i *imageGenerationProcessorUpstreamFilter) SetBackend(ctx context.Context,
 	i.metrics.SetBackend(b)
 	i.modelNameOverride = b.ModelNameOverride
 	i.backendName = b.Name
+	i.span = rp.span
 	if err = i.selectTranslator(b.Schema); err != nil {
 		return fmt.Errorf("failed to select translator: %w", err)
 	}
@@ -442,7 +443,6 @@ func (i *imageGenerationProcessorUpstreamFilter) SetBackend(ctx context.Context,
 		}
 	}
 	rp.upstreamFilter = i
-	i.span = rp.span
 	return
 }
 

internal/extproc/imagegeneration_processor_test.go

@@ -20,6 +20,7 @@ import (
 	openaisdk "github.com/openai/openai-go/v2"
 	"github.com/stretchr/testify/require"
 
+	"github.com/envoyproxy/ai-gateway/internal/apischema/openai"
 	"github.com/envoyproxy/ai-gateway/internal/extproc/translator"
 	"github.com/envoyproxy/ai-gateway/internal/filterapi"
 	"github.com/envoyproxy/ai-gateway/internal/internalapi"
@@ -449,9 +450,9 @@ func TestImageGeneration_ParseBody(t *testing.T) {
 		jsonBody := `{"model":"gpt-image-1-mini","prompt":"a cat","size":"1024x1024","quality":"low"}`
 		modelName, rb, err := parseOpenAIImageGenerationBody(&extprocv3.HttpBody{Body: []byte(jsonBody)})
 		require.NoError(t, err)
-		require.Equal(t, "gpt-image-1-mini", modelName)
+		require.Equal(t, openai.ModelGPTImage1Mini, modelName)
 		require.NotNil(t, rb)
-		require.Equal(t, "gpt-image-1-mini", rb.Model)
+		require.Equal(t, openai.ModelGPTImage1Mini, rb.Model)
 		require.Equal(t, "a cat", rb.Prompt)
 	})
 	t.Run("error", func(t *testing.T) {

internal/extproc/translator/imagegeneration_openai_openai.go

@@ -145,8 +145,13 @@ func (o *openAIToOpenAIImageGenerationTranslator) ResponseBody(_ map[string]stri
 		tokenUsage.TotalTokens = uint32(resp.Usage.TotalTokens)   //nolint:gosec
 	}
 
-	// Provide response model for metrics
+	// There is no response model field, so use the request one.
 	responseModel = o.requestModel
 
+	// Record the response in the span if tracing is enabled.
+	if o.span != nil {
+		o.span.RecordResponse(resp)
+	}
+
 	return
 }

internal/extproc/translator/imagegeneration_openai_openai_test.go

@@ -18,7 +18,7 @@ import (
 )
 
 func TestOpenAIToOpenAIImageTranslator_RequestBody_ModelOverrideAndPath(t *testing.T) {
-	tr := NewImageGenerationOpenAIToOpenAITranslator("v1", "gpt-image-1", nil)
+	tr := NewImageGenerationOpenAIToOpenAITranslator("v1", openai.ModelGPTImage1Mini, nil)
 	req := &openaisdk.ImageGenerateParams{Model: openaisdk.ImageModelDallE3, Prompt: "a cat"}
 	original, _ := json.Marshal(req)
 
@@ -32,9 +32,9 @@ func TestOpenAIToOpenAIImageTranslator_RequestBody_ModelOverrideAndPath(t *testi
 
 	require.NotNil(t, bm)
 	mutated := bm.GetBody()
-	var got openaisdk.ImageGenerateParams
-	require.NoError(t, json.Unmarshal(mutated, &got))
-	require.Equal(t, "gpt-image-1", got.Model)
+	var actual openaisdk.ImageGenerateParams
+	require.NoError(t, json.Unmarshal(mutated, &actual))
+	require.Equal(t, openai.ModelGPTImage1Mini, actual.Model)
 }
 
 func TestOpenAIToOpenAIImageTranslator_RequestBody_ForceMutation(t *testing.T) {
@@ -67,11 +67,11 @@ func TestOpenAIToOpenAIImageTranslator_ResponseError_NonJSON(t *testing.T) {
 	require.NotNil(t, bm)
 
 	// Body should be OpenAI error JSON
-	var got struct {
+	var actual struct {
 		Error openai.ErrorType `json:"error"`
 	}
-	require.NoError(t, json.Unmarshal(bm.GetBody(), &got))
-	require.Equal(t, openAIBackendError, got.Error.Type)
+	require.NoError(t, json.Unmarshal(bm.GetBody(), &actual))
+	require.Equal(t, openAIBackendError, actual.Error.Type)
 }
 
 func TestOpenAIToOpenAIImageTranslator_ResponseBody_OK(t *testing.T) {
@@ -127,7 +127,7 @@ func TestOpenAIToOpenAIImageTranslator_ResponseError_ReadError(t *testing.T) {
 
 func TestOpenAIToOpenAIImageTranslator_ResponseBody_ModelPropagatesFromRequest(t *testing.T) {
 	// Use override so effective model differs from original
-	tr := NewImageGenerationOpenAIToOpenAITranslator("v1", "gpt-image-1", nil)
+	tr := NewImageGenerationOpenAIToOpenAITranslator("v1", openai.ModelGPTImage1Mini, nil)
 	req := &openaisdk.ImageGenerateParams{Model: openaisdk.ImageModelDallE3, Prompt: "a cat"}
 	original, _ := json.Marshal(req)
 	// Call RequestBody first to set requestModel inside translator
@@ -142,7 +142,7 @@ func TestOpenAIToOpenAIImageTranslator_ResponseBody_ModelPropagatesFromRequest(t
 	buf, _ := json.Marshal(resp)
 	_, _, _, respModel, err := tr.ResponseBody(map[string]string{}, bytes.NewReader(buf), true)
 	require.NoError(t, err)
-	require.Equal(t, "gpt-image-1", respModel)
+	require.Equal(t, openai.ModelGPTImage1Mini, respModel)
 }
 
 func TestOpenAIToOpenAIImageTranslator_ResponseHeaders_NoOp(t *testing.T) {
@@ -158,3 +158,28 @@ func TestOpenAIToOpenAIImageTranslator_ResponseBody_DecodeError(t *testing.T) {
 	require.Error(t, err)
 	require.Contains(t, err.Error(), "failed to decode response body")
 }
+
+func TestOpenAIToOpenAIImageTranslator_ResponseBody_RecordsSpan(t *testing.T) {
+	mockSpan := &mockImageGenerationSpan{}
+	tr := NewImageGenerationOpenAIToOpenAITranslator("v1", "", mockSpan)
+
+	resp := &openaisdk.ImagesResponse{
+		Data: []openaisdk.Image{{URL: "https://example.com/img.png"}},
+		Size: openaisdk.ImagesResponseSize1024x1024,
+	}
+	buf, _ := json.Marshal(resp)
+	_, _, _, _, err := tr.ResponseBody(map[string]string{}, bytes.NewReader(buf), true)
+	require.NoError(t, err)
+	require.NotNil(t, mockSpan.recordedResponse)
+}
+
+type mockImageGenerationSpan struct {
+	recordedResponse *openaisdk.ImagesResponse
+}
+
+func (m *mockImageGenerationSpan) RecordResponse(resp *openaisdk.ImagesResponse) {
+	m.recordedResponse = resp
+}
+
+func (m *mockImageGenerationSpan) EndSpanOnError(int, []byte) {}
+func (m *mockImageGenerationSpan) EndSpan()                   {}

internal/tracing/image_generation_span_test.go

@@ -23,7 +23,7 @@ import (
 type testImageGenerationRecorder struct{}
 
 func (r testImageGenerationRecorder) StartParams(_ *openaisdk.ImageGenerateParams, _ []byte) (string, []oteltrace.SpanStartOption) {
-	return "ImageGeneration", nil
+	return "ImagesResponse", nil
 }
 
 func (r testImageGenerationRecorder) RecordRequest(span oteltrace.Span, req *openaisdk.ImageGenerateParams, _ []byte) {

internal/tracing/image_generation_tracer_test.go

@@ -71,7 +71,7 @@ func TestImageGenerationTracer_StartSpanAndInjectHeaders(t *testing.T) {
 			name:             "basic image generation request",
 			req:              imageGenReq,
 			existingHeaders:  map[string]string{},
-			expectedSpanName: "ImageGeneration",
+			expectedSpanName: "ImagesResponse",
 			expectedAttrs: []attribute.KeyValue{
 				attribute.String("model", imageGenReq.Model),
 				attribute.String("prompt", imageGenReq.Prompt),
@@ -90,7 +90,7 @@ func TestImageGenerationTracer_StartSpanAndInjectHeaders(t *testing.T) {
 			existingHeaders: map[string]string{
 				"traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01",
 			},
-			expectedSpanName: "ImageGeneration",
+			expectedSpanName: "ImagesResponse",
 			expectedAttrs: []attribute.KeyValue{
 				attribute.String("model", imageGenReq.Model),
 				attribute.String("prompt", imageGenReq.Prompt),
@@ -115,7 +115,7 @@ func TestImageGenerationTracer_StartSpanAndInjectHeaders(t *testing.T) {
 				N:              openaiparam.NewOpt[int64](2),
 			},
 			existingHeaders:  map[string]string{},
-			expectedSpanName: "ImageGeneration",
+			expectedSpanName: "ImagesResponse",
 			expectedAttrs: []attribute.KeyValue{
 				attribute.String("model", openaisdk.ImageModelGPTImage1),
 				attribute.String("prompt", "a cat and a dog"),
@@ -366,7 +366,7 @@ var _ tracing.ImageGenerationRecorder = testImageGenTracerRecorder{}
 type testImageGenTracerRecorder struct{}
 
 func (r testImageGenTracerRecorder) StartParams(_ *openaisdk.ImageGenerateParams, _ []byte) (spanName string, opts []oteltrace.SpanStartOption) {
-	return "ImageGeneration", imageGenStartOpts
+	return "ImagesResponse", imageGenStartOpts
 }
 
 func (r testImageGenTracerRecorder) RecordRequest(span oteltrace.Span, req *openaisdk.ImageGenerateParams, body []byte) {

internal/tracing/openinference/openai/image_generation.go

@@ -52,7 +52,7 @@ var imageGenStartOpts = []trace.SpanStartOption{trace.WithSpanKind(trace.SpanKin
 
 // StartParams implements the same method as defined in tracing.ImageGenerationRecorder.
 func (r *ImageGenerationRecorder) StartParams(*openaisdk.ImageGenerateParams, []byte) (spanName string, opts []trace.SpanStartOption) {
-	return "ImageGeneration", imageGenStartOpts
+	return "ImagesResponse", imageGenStartOpts
 }
 
 // RecordRequest implements the same method as defined in tracing.ImageGenerationRecorder.
@@ -86,11 +86,10 @@ func (r *ImageGenerationRecorder) RecordResponseOnError(span trace.Span, statusC
 }
 
 // buildImageGenerationRequestAttributes builds OpenInference attributes from the image generation request.
-func buildImageGenerationRequestAttributes(req *openaisdk.ImageGenerateParams, body string, config *openinference.TraceConfig) []attribute.KeyValue {
+func buildImageGenerationRequestAttributes(_ *openaisdk.ImageGenerateParams, body string, config *openinference.TraceConfig) []attribute.KeyValue {
 	attrs := []attribute.KeyValue{
 		attribute.String(openinference.SpanKind, openinference.SpanKindLLM),
 		attribute.String(openinference.LLMSystem, openinference.LLMSystemOpenAI),
-		attribute.String(openinference.LLMModelName, req.Model),
 	}
 
 	if config.HideInputs {
@@ -100,6 +99,10 @@ func buildImageGenerationRequestAttributes(req *openaisdk.ImageGenerateParams, b
 		attrs = append(attrs, attribute.String(openinference.InputMimeType, openinference.MimeTypeJSON))
 	}
 
+	if !config.HideLLMInvocationParameters {
+		attrs = append(attrs, attribute.String(openinference.LLMInvocationParameters, body))
+	}
+
 	return attrs
 }