Skip to content

Commit d9da239

Browse files
feat: update image generation translator constructor for tracing support
- Update NewImageGenerationOpenAIToOpenAITranslator constructor to accept tracing.ImageGenerationSpan parameter - Change modelNameOverride parameter type from string to internalapi.ModelNameOverride for type consistency - Add span field to openAIToOpenAIImageGenerationTranslator struct for tracing support - Update all test calls to pass nil span parameter to maintain compatibility Signed-off-by: Hrushikesh Patil <[email protected]>
1 parent d96a16a commit d9da239

File tree

2 files changed

+15
-9
lines changed

2 files changed

+15
-9
lines changed

internal/extproc/translator/imagegeneration_openai_openai.go

Lines changed: 11 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -10,25 +10,29 @@ import (
1010
"fmt"
1111
"io"
1212
"path"
13+
"strconv"
1314

1415
corev3 "github.com/envoyproxy/go-control-plane/envoy/config/core/v3"
1516
extprocv3 "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3"
1617
"github.com/tidwall/sjson"
1718

19+
"github.com/envoyproxy/ai-gateway/internal/internalapi"
1820
tracing "github.com/envoyproxy/ai-gateway/internal/tracing/api"
1921
openaisdk "github.com/openai/openai-go/v2"
2022
)
2123

2224
// NewImageGenerationOpenAIToOpenAITranslator implements [Factory] for OpenAI to OpenAI image generation translation.
23-
func NewImageGenerationOpenAIToOpenAITranslator(apiVersion string, modelNameOverride string) ImageGenerationTranslator {
24-
return &openAIToOpenAIImageGenerationTranslator{modelNameOverride: modelNameOverride, path: path.Join("/", apiVersion, "images/generations")}
25+
func NewImageGenerationOpenAIToOpenAITranslator(apiVersion string, modelNameOverride internalapi.ModelNameOverride, span tracing.ImageGenerationSpan) ImageGenerationTranslator {
26+
return &openAIToOpenAIImageGenerationTranslator{modelNameOverride: modelNameOverride, path: path.Join("/", apiVersion, "images/generations"), span: span}
2527
}
2628

2729
// openAIToOpenAIImageGenerationTranslator implements [ImageGenerationTranslator] for /v1/images/generations.
2830
type openAIToOpenAIImageGenerationTranslator struct {
29-
modelNameOverride string
31+
modelNameOverride internalapi.ModelNameOverride
3032
// The path of the images generations endpoint to be used for the request. It is prefixed with the OpenAI path prefix.
3133
path string
34+
// span is the tracing span for this request, inherited from the router filter.
35+
span tracing.ImageGenerationSpan
3236
}
3337

3438
// RequestBody implements [ImageGenerationTranslator.RequestBody].
@@ -62,8 +66,10 @@ func (o *openAIToOpenAIImageGenerationTranslator) RequestBody(original []byte, r
6266
bodyMutation = &extprocv3.BodyMutation{
6367
Mutation: &extprocv3.BodyMutation_Body{Body: newBody},
6468
}
65-
// Note: content-length header is set via dynamic metadata in the processor
66-
// to avoid conflicts with Envoy's REPLACE_AND_CONTINUE processing mode
69+
headerMutation.SetHeaders = append(headerMutation.SetHeaders, &corev3.HeaderValueOption{Header: &corev3.HeaderValue{
70+
Key: "content-length",
71+
RawValue: []byte(strconv.Itoa(len(newBody))),
72+
}})
6773
}
6874
return
6975
}

internal/extproc/translator/imagegeneration_openai_openai_test.go

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@ import (
1616
)
1717

1818
func TestOpenAIToOpenAIImageTranslator_RequestBody_ModelOverrideAndPath(t *testing.T) {
19-
tr := NewImageGenerationOpenAIToOpenAITranslator("v1", "gpt-image-1")
19+
tr := NewImageGenerationOpenAIToOpenAITranslator("v1", "gpt-image-1", nil)
2020
req := &openaisdk.ImageGenerateParams{Model: openaisdk.ImageModelDallE3, Prompt: "a cat"}
2121
original, _ := json.Marshal(req)
2222

@@ -36,7 +36,7 @@ func TestOpenAIToOpenAIImageTranslator_RequestBody_ModelOverrideAndPath(t *testi
3636
}
3737

3838
func TestOpenAIToOpenAIImageTranslator_RequestBody_ForceMutation(t *testing.T) {
39-
tr := NewImageGenerationOpenAIToOpenAITranslator("v1", "")
39+
tr := NewImageGenerationOpenAIToOpenAITranslator("v1", "", nil)
4040
req := &openaisdk.ImageGenerateParams{Model: openaisdk.ImageModelDallE2, Prompt: "a cat"}
4141
original, _ := json.Marshal(req)
4242

@@ -57,7 +57,7 @@ func TestOpenAIToOpenAIImageTranslator_RequestBody_ForceMutation(t *testing.T) {
5757
}
5858

5959
func TestOpenAIToOpenAIImageTranslator_ResponseError_NonJSON(t *testing.T) {
60-
tr := NewImageGenerationOpenAIToOpenAITranslator("v1", "")
60+
tr := NewImageGenerationOpenAIToOpenAITranslator("v1", "", nil)
6161
headers := map[string]string{contentTypeHeaderName: "text/plain", statusHeaderName: "503"}
6262
hm, bm, err := tr.ResponseError(headers, bytes.NewReader([]byte("backend error")))
6363
require.NoError(t, err)
@@ -71,7 +71,7 @@ func TestOpenAIToOpenAIImageTranslator_ResponseError_NonJSON(t *testing.T) {
7171
}
7272

7373
func TestOpenAIToOpenAIImageTranslator_ResponseBody_OK(t *testing.T) {
74-
tr := NewImageGenerationOpenAIToOpenAITranslator("v1", "")
74+
tr := NewImageGenerationOpenAIToOpenAITranslator("v1", "", nil)
7575
resp := &openaisdk.ImagesResponse{Size: openaisdk.ImagesResponseSize1024x1024}
7676
buf, _ := json.Marshal(resp)
7777
hm, bm, usage, metadata, err := tr.ResponseBody(map[string]string{}, bytes.NewReader(buf), true)

0 commit comments

Comments
 (0)