chore(lint): fix revive, unconvert, and testifylint in image generation and tracing

nutanix-Hrushikesh · nutanix-Hrushikesh · commit 46012bbad9ef · 2025-10-06T22:29:31.000+05:30
- Ignore unused parameters in tests and translator (revive)\n- Remove unnecessary string conversions now that model types are strings (unconvert)\n- Use require.Empty for empty string check (testifylint)\n- Remove unused test variable in span tests (unused)\n- Keep extractUsageFromBufferEvent with lint ignore for future parity
diff --git a/.gitignore b/.gitignore
@@ -49,3 +49,5 @@ inference-extension-conformance-test-report.yaml
 
 .goose
 tests/e2e-inference-extension/logs/
+
+/aigw
diff --git a/cmd/aigw/README.md b/cmd/aigw/README.md
@@ -51,6 +51,21 @@ Here are values we use for Ollama:
      docker compose run --rm embeddings
      ```
 
+   - Image generation:
+     - Using service:
+       ```bash
+       docker compose run --rm image-generation
+       ```
+     - Using curl (save to file):
+       ```bash
+       curl -s \
+         -X POST http://localhost:1975/v1/images/generations \
+         -H "Authorization: Bearer unused" \
+         -H "Content-Type: application/json" \
+         -d '{"model":"gpt-image-1","prompt":"A watercolor painting of a red fox in a birch forest","size":"512x512"}' \
+         | jq -r '.data[0].b64_json' | base64 -d > image.png
+       ```
+
 4. **Create embeddings**:
 
    The `create-embeddings` service uses `curl` to send an embeddings request
@@ -144,6 +159,8 @@ This configures the OTLP endpoint to otel-tui on port 4318.
    ```bash
    COMPOSE_PROFILES=<profile> docker compose -f docker-compose-otel.yaml run --build --rm chat-completion
    COMPOSE_PROFILES=<profile> docker compose -f docker-compose-otel.yaml run --build --rm create-embeddings
+      # Image generation
+      COMPOSE_PROFILES=<profile> docker compose -f docker-compose-otel.yaml run --build --rm image-generation
    ```
 
 3. **Check telemetry output**:
diff --git a/cmd/aigw/docker-compose.yaml b/cmd/aigw/docker-compose.yaml
@@ -99,3 +99,23 @@ services:
           -d "{\"model\":\"$$EMBEDDINGS_MODEL\",\"input\":\"How do I reset my password?\"}"
     extra_hosts:  # localhost:host-gateway trick doesn't work with aigw
       - "host.docker.internal:host-gateway"
+
+  # image-generation is a simple curl-based test client for sending image
+  # generation requests to aigw.
+  image-generation:
+    image: golang:1.25
+    container_name: image-generation
+    profiles: ["test"]
+    env_file:
+      - ../../.env.ollama
+    command:
+      - sh
+      - -c
+      - |
+        curl -s -w %{http_code} \
+          -X POST http://aigw:1975/v1/images/generations \
+          -H "Authorization: Bearer unused" \
+          -H "Content-Type: application/json" \
+          -d "{\"model\":\"$$CHAT_MODEL\",\"prompt\":\"A watercolor painting of a red fox in a birch forest\",\"size\":\"256x256\"}"
+    extra_hosts:  # localhost:host-gateway trick doesn't work with aigw
+      - "host.docker.internal:host-gateway"
diff --git a/go.mod b/go.mod
@@ -10,6 +10,7 @@ require (
 	github.com/Azure/azure-sdk-for-go/sdk/azidentity v1.12.0
 	github.com/a8m/envsubst v1.4.3
 	github.com/alecthomas/kong v1.12.1
+	github.com/andybalholm/brotli v1.2.0
 	github.com/anthropics/anthropic-sdk-go v1.12.0
 	github.com/aws/aws-sdk-go-v2 v1.39.1
 	github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.7.1
@@ -83,7 +84,6 @@ require (
 	github.com/AzureAD/microsoft-authentication-library-for-go v1.5.0 // indirect
 	github.com/Microsoft/go-winio v0.6.2 // indirect
 	github.com/NYTimes/gziphandler v1.1.1 // indirect
-	github.com/andybalholm/brotli v1.2.0 // indirect
 	github.com/antlr4-go/antlr/v4 v4.13.1 // indirect
 	github.com/asaskevich/govalidator v0.0.0-20230301143203-a9d515a09cc2 // indirect
 	github.com/avast/retry-go v3.0.0+incompatible // indirect
diff --git a/go.sum b/go.sum
@@ -425,6 +425,8 @@ github.com/x448/float16 v0.8.4 h1:qLwI1I70+NjRFUR3zs1JPUCgaCXSh3SW62uAKT1mSBM=
 github.com/x448/float16 v0.8.4/go.mod h1:14CWIYCyZA/cWjXOioeEpHeN/83MdbZDRQHoFcYsOfg=
 github.com/xiang90/probing v0.0.0-20221125231312-a49e3df8f510 h1:S2dVYn90KE98chqDkyE9Z4N61UnQd+KOfgp5Iu53llk=
 github.com/xiang90/probing v0.0.0-20221125231312-a49e3df8f510/go.mod h1:UETIi67q53MR2AWcXfiuqkDkRtnGDLqkBTpCHuJHxtU=
+github.com/xyproto/randomstring v1.0.5 h1:YtlWPoRdgMu3NZtP45drfy1GKoojuR7hmRcnhZqKjWU=
+github.com/xyproto/randomstring v1.0.5/go.mod h1:rgmS5DeNXLivK7YprL0pY+lTuhNQW3iGxZ18UQApw/E=
 github.com/yosida95/uritemplate/v3 v3.0.2 h1:Ed3Oyj9yrmi9087+NczuL5BwkIc4wvTb5zIM+UJPGz4=
 github.com/yosida95/uritemplate/v3 v3.0.2/go.mod h1:ILOh0sOhIJR3+L/8afwt/kE++YT040gmv5BQTMR2HP4=
 github.com/yuin/goldmark v1.1.27/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
diff --git a/internal/extproc/imagegeneration_processor.go b/internal/extproc/imagegeneration_processor.go
@@ -18,6 +18,7 @@ import (
 	corev3 "github.com/envoyproxy/go-control-plane/envoy/config/core/v3"
 	extprocv3http "github.com/envoyproxy/go-control-plane/envoy/extensions/filters/http/ext_proc/v3"
 	extprocv3 "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3"
+	openaisdk "github.com/openai/openai-go/v2"
 	"google.golang.org/protobuf/types/known/structpb"
 
 	"github.com/envoyproxy/ai-gateway/internal/extproc/backendauth"
@@ -27,7 +28,6 @@ import (
 	"github.com/envoyproxy/ai-gateway/internal/internalapi"
 	"github.com/envoyproxy/ai-gateway/internal/metrics"
 	tracing "github.com/envoyproxy/ai-gateway/internal/tracing/api"
-	openaisdk "github.com/openai/openai-go/v2"
 )
 
 // ImageGenerationProcessorFactory returns a factory method to instantiate the image generation processor.
@@ -215,11 +215,11 @@ func (i *imageGenerationProcessorUpstreamFilter) ProcessRequestHeaders(ctx conte
 	// Start tracking metrics for this request.
 	i.metrics.StartRequest(i.requestHeaders)
 	// Set the original model from the request body before any overrides
-	i.metrics.SetOriginalModel(internalapi.OriginalModel(i.originalRequestBody.Model))
+	i.metrics.SetOriginalModel(i.originalRequestBody.Model)
 	// Set the request model for metrics from the original model or override if applied.
-	reqModel := cmp.Or(i.requestHeaders[i.config.modelNameHeaderKey], string(i.originalRequestBody.Model))
-	i.metrics.SetRequestModel(internalapi.RequestModel(reqModel))
-	i.metrics.SetResponseModel(internalapi.ResponseModel(reqModel))
+	reqModel := cmp.Or(i.requestHeaders[i.config.modelNameHeaderKey], i.originalRequestBody.Model)
+	i.metrics.SetRequestModel(reqModel)
+	i.metrics.SetResponseModel(reqModel)
 
 	// We force the body mutation in the following cases:
 	// * The request is a retry request because the body mutation might have happened the previous iteration.
@@ -471,9 +471,9 @@ func (i *imageGenerationProcessorUpstreamFilter) SetBackend(ctx context.Context,
 	i.headerMutator = headermutator.NewHeaderMutator(b.HeaderMutation, rp.requestHeaders)
 	// Sync header with backend model so header-derived labels/CEL use the actual model.
 	if i.modelNameOverride != "" {
-		i.requestHeaders[i.config.modelNameHeaderKey] = string(i.modelNameOverride)
+		i.requestHeaders[i.config.modelNameHeaderKey] = i.modelNameOverride
 		// Update metrics with the overridden model
-		i.metrics.SetRequestModel(internalapi.RequestModel(i.modelNameOverride))
+		i.metrics.SetRequestModel(i.modelNameOverride)
 	}
 	i.originalRequestBody = rp.originalRequestBody
 	i.originalRequestBodyRaw = rp.originalRequestBodyRaw
@@ -498,5 +498,5 @@ func parseOpenAIImageGenerationBody(body *extprocv3.HttpBody) (modelName string,
 	if err := json.Unmarshal(body.Body, &openAIReq); err != nil {
 		return "", nil, fmt.Errorf("failed to unmarshal body: %w", err)
 	}
-	return string(openAIReq.Model), &openAIReq, nil
+	return openAIReq.Model, &openAIReq, nil
 }
diff --git a/internal/extproc/imagegeneration_processor_test.go b/internal/extproc/imagegeneration_processor_test.go
@@ -15,13 +15,13 @@ import (
 
 	corev3 "github.com/envoyproxy/go-control-plane/envoy/config/core/v3"
 	extprocv3 "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3"
+	openaisdk "github.com/openai/openai-go/v2"
 	"github.com/stretchr/testify/require"
 
 	"github.com/envoyproxy/ai-gateway/internal/extproc/translator"
 	"github.com/envoyproxy/ai-gateway/internal/filterapi"
 	"github.com/envoyproxy/ai-gateway/internal/llmcostcel"
 	tracing "github.com/envoyproxy/ai-gateway/internal/tracing/api"
-	openaisdk "github.com/openai/openai-go/v2"
 )
 
 func TestImageGeneration_Schema(t *testing.T) {
@@ -95,7 +95,7 @@ func (m *mockImageGenerationSpan) EndSpanOnError(status int, body []byte) {
 	m.errBody = string(body)
 }
 
-func (m *mockImageGenerationSpan) RecordResponse(resp *openaisdk.ImagesResponse) {
+func (m *mockImageGenerationSpan) RecordResponse(_ *openaisdk.ImagesResponse) {
 	// Mock implementation
 }
 
@@ -390,7 +390,7 @@ type mockImageGenerationTranslator struct {
 	retImageMetadata            translator.ImageGenerationMetadata
 }
 
-func (m *mockImageGenerationTranslator) RequestBody(original []byte, req *openaisdk.ImageGenerateParams, forceBodyMutation bool) (*extprocv3.HeaderMutation, *extprocv3.BodyMutation, error) {
+func (m *mockImageGenerationTranslator) RequestBody(_ []byte, req *openaisdk.ImageGenerateParams, forceBodyMutation bool) (*extprocv3.HeaderMutation, *extprocv3.BodyMutation, error) {
 	if m.expRequestBody != nil {
 		require.Equal(m.t, m.expRequestBody, req)
 	}
@@ -410,6 +410,7 @@ func (m *mockImageGenerationTranslator) ResponseHeaders(headers map[string]strin
 }
 
 func (m *mockImageGenerationTranslator) ResponseBody(headers map[string]string, body io.Reader, endOfStream bool) (*extprocv3.HeaderMutation, *extprocv3.BodyMutation, translator.LLMTokenUsage, translator.ImageGenerationMetadata, error) {
+	_ = headers
 	if m.expResponseBody != nil {
 		bodyBytes, _ := io.ReadAll(body)
 		require.Equal(m.t, m.expResponseBody.Body, bodyBytes)
@@ -419,13 +420,15 @@ func (m *mockImageGenerationTranslator) ResponseBody(headers map[string]string,
 }
 
 func (m *mockImageGenerationTranslator) ResponseError(headers map[string]string, body io.Reader) (*extprocv3.HeaderMutation, *extprocv3.BodyMutation, error) {
+	_ = headers
+	_ = body
 	return m.retHeaderMutation, m.retBodyMutation, m.retErr
 }
 
 // imageGenerationBodyFromModel returns a minimal valid image generation request for tests.
 func imageGenerationBodyFromModel(t *testing.T, model string) []byte {
 	t.Helper()
-	b, err := json.Marshal(&openaisdk.ImageGenerateParams{Model: openaisdk.ImageModel(model), Prompt: "a cat"})
+	b, err := json.Marshal(&openaisdk.ImageGenerateParams{Model: model, Prompt: "a cat"})
 	require.NoError(t, err)
 	return b
 }
diff --git a/internal/extproc/mocks_test.go b/internal/extproc/mocks_test.go
@@ -410,47 +410,57 @@ func (m *mockImageGenerationMetrics) StartRequest(map[string]string) {}
 func (m *mockImageGenerationMetrics) SetOriginalModel(originalModel string) {
 	m.model = originalModel
 }
+
 func (m *mockImageGenerationMetrics) SetRequestModel(requestModel string) {
 	m.model = requestModel
 }
+
 func (m *mockImageGenerationMetrics) SetResponseModel(responseModel string) {
 	m.model = responseModel
 }
-func (m *mockImageGenerationMetrics) SetModel(requestModel, responseModel string) {
+
+func (m *mockImageGenerationMetrics) SetModel(_ string, responseModel string) {
 	m.model = responseModel
 }
 func (m *mockImageGenerationMetrics) SetBackend(b *filterapi.Backend) { m.backend = b.Name }
 func (m *mockImageGenerationMetrics) RecordTokenUsage(_ context.Context, input, output uint32, _ map[string]string) {
 	m.tokenUsageCount += int(input + output)
 }
+
 func (m *mockImageGenerationMetrics) RecordRequestCompletion(_ context.Context, success bool, _ map[string]string) {
 	if success {
 		m.requestSuccessCount++
 	} else {
 		m.requestErrorCount++
 	}
 }
+
 func (m *mockImageGenerationMetrics) RecordImageGeneration(_ context.Context, _ int, _ string, _ string, _ map[string]string) {
 }
 
 func (m *mockImageGenerationMetrics) RequireRequestFailure(t *testing.T) {
 	require.Equal(t, 0, m.requestSuccessCount)
 	require.Equal(t, 1, m.requestErrorCount)
 }
+
 func (m *mockImageGenerationMetrics) RequireRequestNotCompleted(t *testing.T) {
 	require.Equal(t, 0, m.requestSuccessCount)
 	require.Equal(t, 0, m.requestErrorCount)
 }
+
 func (m *mockImageGenerationMetrics) RequireRequestSuccess(t *testing.T) {
 	require.Equal(t, 1, m.requestSuccessCount)
 	require.Equal(t, 0, m.requestErrorCount)
 }
+
 func (m *mockImageGenerationMetrics) RequireSelectedModel(t *testing.T, model string) {
 	require.Equal(t, model, m.model)
 }
+
 func (m *mockImageGenerationMetrics) RequireSelectedBackend(t *testing.T, backend string) {
 	require.Equal(t, backend, m.backend)
 }
+
 func (m *mockImageGenerationMetrics) RequireTokensRecorded(t *testing.T, count int) {
 	require.Equal(t, count, m.tokenUsageCount)
 }
diff --git a/internal/extproc/translator/imagegeneration_openai_openai.go b/internal/extproc/translator/imagegeneration_openai_openai.go
@@ -14,11 +14,11 @@ import (
 
 	corev3 "github.com/envoyproxy/go-control-plane/envoy/config/core/v3"
 	extprocv3 "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3"
+	openaisdk "github.com/openai/openai-go/v2"
 	"github.com/tidwall/sjson"
 
 	"github.com/envoyproxy/ai-gateway/internal/internalapi"
 	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.
@@ -36,7 +36,7 @@ type openAIToOpenAIImageGenerationTranslator struct {
 }
 
 // RequestBody implements [ImageGenerationTranslator.RequestBody].
-func (o *openAIToOpenAIImageGenerationTranslator) RequestBody(original []byte, req *openaisdk.ImageGenerateParams, forceBodyMutation bool) (
+func (o *openAIToOpenAIImageGenerationTranslator) RequestBody(original []byte, _ *openaisdk.ImageGenerateParams, forceBodyMutation bool) (
 	headerMutation *extprocv3.HeaderMutation, bodyMutation *extprocv3.BodyMutation, err error,
 ) {
 	var newBody []byte
@@ -172,7 +172,9 @@ func (o *openAIToOpenAIImageGenerationTranslator) ResponseBody(_ map[string]stri
 // extractUsageFromBufferEvent extracts token usage from buffered streaming events.
 // This is currently not applicable for image generation as it doesn't use streaming.
 // TODO: Implement if streaming support is added for image generation in the future.
-func (o *openAIToOpenAIImageGenerationTranslator) extractUsageFromBufferEvent(span tracing.ImageGenerationSpan) LLMTokenUsage {
-	// Image generation doesn't use streaming, so no token usage to extract
+// NOTE: image generation currently does not use streaming; keep for future parity with other translators.
+//
+//lint:ignore U1000 kept for parity and future use
+func (o *openAIToOpenAIImageGenerationTranslator) extractUsageFromBufferEvent(_ tracing.ImageGenerationSpan) LLMTokenUsage {
 	return LLMTokenUsage{}
 }
diff --git a/internal/extproc/translator/imagegeneration_openai_openai_test.go b/internal/extproc/translator/imagegeneration_openai_openai_test.go
@@ -10,9 +10,8 @@ import (
 	"encoding/json"
 	"testing"
 
-	"github.com/stretchr/testify/require"
-
 	openaisdk "github.com/openai/openai-go/v2"
+	"github.com/stretchr/testify/require"
 )
 
 func TestOpenAIToOpenAIImageTranslator_RequestBody_ModelOverrideAndPath(t *testing.T) {
@@ -32,7 +31,7 @@ func TestOpenAIToOpenAIImageTranslator_RequestBody_ModelOverrideAndPath(t *testi
 	mutated := bm.GetBody()
 	var got openaisdk.ImageGenerateParams
 	require.NoError(t, json.Unmarshal(mutated, &got))
-	require.Equal(t, "gpt-image-1", string(got.Model))
+	require.Equal(t, "gpt-image-1", got.Model)
 }
 
 func TestOpenAIToOpenAIImageTranslator_RequestBody_ForceMutation(t *testing.T) {
@@ -81,6 +80,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, "", metadata.Model)
+	require.Empty(t, metadata.Model)
 	require.Equal(t, string(openaisdk.ImagesResponseSize1024x1024), metadata.Size)
 }
diff --git a/internal/extproc/translator/translator.go b/internal/extproc/translator/translator.go
@@ -11,13 +11,13 @@ import (
 
 	corev3 "github.com/envoyproxy/go-control-plane/envoy/config/core/v3"
 	extprocv3 "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3"
+	openaisdk "github.com/openai/openai-go/v2"
 	"github.com/tidwall/sjson"
 
 	anthropicschema "github.com/envoyproxy/ai-gateway/internal/apischema/anthropic"
 	"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 (
diff --git a/internal/tracing/api/api.go b/internal/tracing/api/api.go
@@ -11,11 +11,11 @@ import (
 	"context"
 
 	extprocv3 "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3"
+	openaisdk "github.com/openai/openai-go/v2"
 	"go.opentelemetry.io/otel/propagation"
 	"go.opentelemetry.io/otel/trace"
 
 	"github.com/envoyproxy/ai-gateway/internal/apischema/openai"
-	openaisdk "github.com/openai/openai-go/v2"
 )
 
 var _ Tracing = NoopTracing{}
diff --git a/internal/tracing/image_generation_span.go b/internal/tracing/image_generation_span.go
@@ -6,10 +6,10 @@
 package tracing
 
 import (
+	openaisdk "github.com/openai/openai-go/v2"
 	"go.opentelemetry.io/otel/trace"
 
 	tracing "github.com/envoyproxy/ai-gateway/internal/tracing/api"
-	openaisdk "github.com/openai/openai-go/v2"
 )
 
 // Ensure imageGenerationSpan implements ImageGenerationSpan.
diff --git a/internal/tracing/image_generation_span_test.go b/internal/tracing/image_generation_span_test.go
diff --git a/internal/tracing/image_generation_tracer.go b/internal/tracing/image_generation_tracer.go
diff --git a/internal/tracing/image_generation_tracer_test.go b/internal/tracing/image_generation_tracer_test.go
diff --git a/internal/tracing/openinference/openai/image_generation.go b/internal/tracing/openinference/openai/image_generation.go
diff --git a/internal/tracing/openinference/openai/image_generation_test.go b/internal/tracing/openinference/openai/image_generation_test.go

-Original file line number
+Diff line change
 .goose
 tests/e2e-inference-extension/logs/
++
 +/aigw
Original file line number	Diff line number	Diff line change
`@@ -15,13 +15,13 @@ import (`
`15`	`15`
`16`	`16`	`corev3 "github.com/envoyproxy/go-control-plane/envoy/config/core/v3"`
`17`	`17`	`extprocv3 "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3"`
	`18`	`+ openaisdk "github.com/openai/openai-go/v2"`
`18`	`19`	`"github.com/stretchr/testify/require"`
`19`	`20`
`20`	`21`	`"github.com/envoyproxy/ai-gateway/internal/extproc/translator"`
`21`	`22`	`"github.com/envoyproxy/ai-gateway/internal/filterapi"`
`22`	`23`	`"github.com/envoyproxy/ai-gateway/internal/llmcostcel"`
`23`	`24`	`tracing "github.com/envoyproxy/ai-gateway/internal/tracing/api"`
`24`		`- openaisdk "github.com/openai/openai-go/v2"`
`25`	`25`	`)`
`26`	`26`
`27`	`27`	`func TestImageGeneration_Schema(t *testing.T) {`
`@@ -95,7 +95,7 @@ func (m *mockImageGenerationSpan) EndSpanOnError(status int, body []byte) {`
`95`	`95`	`m.errBody = string(body)`
`96`	`96`	`}`
`97`	`97`
`98`		`-func (m mockImageGenerationSpan) RecordResponse(resp openaisdk.ImagesResponse) {`
	`98`	`+func (m mockImageGenerationSpan) RecordResponse(_ openaisdk.ImagesResponse) {`
`99`	`99`	`// Mock implementation`
`100`	`100`	`}`
`101`	`101`
`@@ -390,7 +390,7 @@ type mockImageGenerationTranslator struct {`
`390`	`390`	`retImageMetadata translator.ImageGenerationMetadata`
`391`	`391`	`}`
`392`	`392`
`393`		`-func (m mockImageGenerationTranslator) RequestBody(original []byte, req openaisdk.ImageGenerateParams, forceBodyMutation bool) (extprocv3.HeaderMutation, extprocv3.BodyMutation, error) {`
	`393`	`+func (m mockImageGenerationTranslator) RequestBody(_ []byte, req openaisdk.ImageGenerateParams, forceBodyMutation bool) (extprocv3.HeaderMutation, extprocv3.BodyMutation, error) {`
`394`	`394`	`if m.expRequestBody != nil {`
`395`	`395`	`require.Equal(m.t, m.expRequestBody, req)`
`396`	`396`	`}`
`@@ -410,6 +410,7 @@ func (m *mockImageGenerationTranslator) ResponseHeaders(headers map[string]strin`
`410`	`410`	`}`
`411`	`411`
`412`	`412`	`func (m mockImageGenerationTranslator) ResponseBody(headers map[string]string, body io.Reader, endOfStream bool) (extprocv3.HeaderMutation, *extprocv3.BodyMutation, translator.LLMTokenUsage, translator.ImageGenerationMetadata, error) {`
	`413`	`+ _ = headers`
`413`	`414`	`if m.expResponseBody != nil {`
`414`	`415`	`bodyBytes, _ := io.ReadAll(body)`
`415`	`416`	`require.Equal(m.t, m.expResponseBody.Body, bodyBytes)`
`@@ -419,13 +420,15 @@ func (m *mockImageGenerationTranslator) ResponseBody(headers map[string]string,`
`419`	`420`	`}`
`420`	`421`
`421`	`422`	`func (m mockImageGenerationTranslator) ResponseError(headers map[string]string, body io.Reader) (extprocv3.HeaderMutation, *extprocv3.BodyMutation, error) {`
	`423`	`+ _ = headers`
	`424`	`+ _ = body`
`422`	`425`	`return m.retHeaderMutation, m.retBodyMutation, m.retErr`
`423`	`426`	`}`
`424`	`427`
`425`	`428`	`// imageGenerationBodyFromModel returns a minimal valid image generation request for tests.`
`426`	`429`	`func imageGenerationBodyFromModel(t *testing.T, model string) []byte {`
`427`	`430`	`t.Helper()`
`428`		`- b, err := json.Marshal(&openaisdk.ImageGenerateParams{Model: openaisdk.ImageModel(model), Prompt: "a cat"})`
	`431`	`+ b, err := json.Marshal(&openaisdk.ImageGenerateParams{Model: model, Prompt: "a cat"})`
`429`	`432`	`require.NoError(t, err)`
`430`	`433`	`return b`
`431`	`434`	`}`
Original file line number	Diff line number	Diff line change
`@@ -410,47 +410,57 @@ func (m *mockImageGenerationMetrics) StartRequest(map[string]string) {}`
`410`	`410`	`func (m *mockImageGenerationMetrics) SetOriginalModel(originalModel string) {`
`411`	`411`	`m.model = originalModel`
`412`	`412`	`}`
	`413`	`+`
`413`	`414`	`func (m *mockImageGenerationMetrics) SetRequestModel(requestModel string) {`
`414`	`415`	`m.model = requestModel`
`415`	`416`	`}`
	`417`	`+`
`416`	`418`	`func (m *mockImageGenerationMetrics) SetResponseModel(responseModel string) {`
`417`	`419`	`m.model = responseModel`
`418`	`420`	`}`
`419`		`-func (m *mockImageGenerationMetrics) SetModel(requestModel, responseModel string) {`
	`421`	`+`
	`422`	`+func (m *mockImageGenerationMetrics) SetModel(_ string, responseModel string) {`
`420`	`423`	`m.model = responseModel`
`421`	`424`	`}`
`422`	`425`	`func (m mockImageGenerationMetrics) SetBackend(b filterapi.Backend) { m.backend = b.Name }`
`423`	`426`	`func (m *mockImageGenerationMetrics) RecordTokenUsage(_ context.Context, input, output uint32, _ map[string]string) {`
`424`	`427`	`m.tokenUsageCount += int(input + output)`
`425`	`428`	`}`
	`429`	`+`
`426`	`430`	`func (m *mockImageGenerationMetrics) RecordRequestCompletion(_ context.Context, success bool, _ map[string]string) {`
`427`	`431`	`if success {`
`428`	`432`	`m.requestSuccessCount++`
`429`	`433`	`} else {`
`430`	`434`	`m.requestErrorCount++`
`431`	`435`	`}`
`432`	`436`	`}`
	`437`	`+`
`433`	`438`	`func (m *mockImageGenerationMetrics) RecordImageGeneration(_ context.Context, _ int, _ string, _ string, _ map[string]string) {`
`434`	`439`	`}`
`435`	`440`
`436`	`441`	`func (m mockImageGenerationMetrics) RequireRequestFailure(t testing.T) {`
`437`	`442`	`require.Equal(t, 0, m.requestSuccessCount)`
`438`	`443`	`require.Equal(t, 1, m.requestErrorCount)`
`439`	`444`	`}`
	`445`	`+`
`440`	`446`	`func (m mockImageGenerationMetrics) RequireRequestNotCompleted(t testing.T) {`
`441`	`447`	`require.Equal(t, 0, m.requestSuccessCount)`
`442`	`448`	`require.Equal(t, 0, m.requestErrorCount)`
`443`	`449`	`}`
	`450`	`+`
`444`	`451`	`func (m mockImageGenerationMetrics) RequireRequestSuccess(t testing.T) {`
`445`	`452`	`require.Equal(t, 1, m.requestSuccessCount)`
`446`	`453`	`require.Equal(t, 0, m.requestErrorCount)`
`447`	`454`	`}`
	`455`	`+`
`448`	`456`	`func (m mockImageGenerationMetrics) RequireSelectedModel(t testing.T, model string) {`
`449`	`457`	`require.Equal(t, model, m.model)`
`450`	`458`	`}`
	`459`	`+`
`451`	`460`	`func (m mockImageGenerationMetrics) RequireSelectedBackend(t testing.T, backend string) {`
`452`	`461`	`require.Equal(t, backend, m.backend)`
`453`	`462`	`}`
	`463`	`+`
`454`	`464`	`func (m mockImageGenerationMetrics) RequireTokensRecorded(t testing.T, count int) {`
`455`	`465`	`require.Equal(t, count, m.tokenUsageCount)`
`456`	`466`	`}`