diff --git a/backend/internal/service/openai_gateway_response_handling.go b/backend/internal/service/openai_gateway_response_handling.go index 2acf0a9307..7e59f778e6 100644 --- a/backend/internal/service/openai_gateway_response_handling.go +++ b/backend/internal/service/openai_gateway_response_handling.go @@ -967,9 +967,30 @@ func extractOpenAIUsageFromJSONBytes(body []byte) (OpenAIUsage, bool) { return OpenAIUsage{}, false } if usage, ok := openAIUsageFromGJSON(gjson.GetBytes(body, "usage")); ok { + mergeHostedImageGenToolUsage(gjson.GetBytes(body, "tool_usage.image_gen"), &usage) return usage, true } - return openAIUsageFromGJSON(gjson.GetBytes(body, "response.usage")) + if usage, ok := openAIUsageFromGJSON(gjson.GetBytes(body, "response.usage")); ok { + mergeHostedImageGenToolUsage(gjson.GetBytes(body, "response.tool_usage.image_gen"), &usage) + return usage, true + } + return OpenAIUsage{}, false +} + +func mergeHostedImageGenToolUsage(imageGen gjson.Result, usage *OpenAIUsage) { + if !imageGen.Exists() || !imageGen.IsObject() { + return + } + if usage.ImageOutputTokens == 0 { + if v := imageGen.Get("output_tokens_details.image_tokens").Int(); v > 0 { + usage.ImageOutputTokens = int(v) + } + } + if usage.ImageInputTokens == 0 { + if v := imageGen.Get("input_tokens_details.image_tokens").Int(); v > 0 { + usage.ImageInputTokens = int(v) + } + } } func extractOpenAIResponseIDFromJSONBytes(body []byte) string { diff --git a/backend/internal/service/openai_gateway_response_handling_image_usage_test.go b/backend/internal/service/openai_gateway_response_handling_image_usage_test.go new file mode 100644 index 0000000000..fad1e1ac62 --- /dev/null +++ b/backend/internal/service/openai_gateway_response_handling_image_usage_test.go @@ -0,0 +1,157 @@ +package service + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/tidwall/gjson" +) + +func TestExtractOpenAIUsageFromJSONBytes_MergesHostedImageGenToolUsage(t *testing.T) { + // SSE response.completed event with response.usage + response.tool_usage.image_gen + body := []byte(`{ + "type": "response.completed", + "response": { + "usage": { + "input_tokens": 43792, + "output_tokens": 1005, + "total_tokens": 44797 + }, + "tool_usage": { + "image_gen": { + "input_tokens": 7918, + "input_tokens_details": {"image_tokens": 7620, "text_tokens": 298}, + "output_tokens": 186, + "output_tokens_details": {"image_tokens": 186, "text_tokens": 0}, + "total_tokens": 8104 + } + } + } + }`) + + usage, ok := extractOpenAIUsageFromJSONBytes(body) + assert.True(t, ok) + assert.Equal(t, 43792, usage.InputTokens, "input_tokens from response.usage") + assert.Equal(t, 1005, usage.OutputTokens, "output_tokens from response.usage") + assert.Equal(t, 186, usage.ImageOutputTokens, "image output tokens merged from tool_usage.image_gen") + assert.Equal(t, 7620, usage.ImageInputTokens, "image input tokens merged from tool_usage.image_gen") +} + +func TestExtractOpenAIUsageFromJSONBytes_NonStreamingMergesImageGen(t *testing.T) { + // Non-streaming response with top-level usage + tool_usage + body := []byte(`{ + "id": "resp_abc123", + "object": "response", + "usage": { + "input_tokens": 5000, + "output_tokens": 200 + }, + "tool_usage": { + "image_gen": { + "input_tokens": 3000, + "input_tokens_details": {"image_tokens": 2800, "text_tokens": 200}, + "output_tokens": 150, + "output_tokens_details": {"image_tokens": 150, "text_tokens": 0}, + "total_tokens": 3150 + } + } + }`) + + usage, ok := extractOpenAIUsageFromJSONBytes(body) + assert.True(t, ok) + assert.Equal(t, 5000, usage.InputTokens) + assert.Equal(t, 200, usage.OutputTokens) + assert.Equal(t, 150, usage.ImageOutputTokens, "image output tokens from tool_usage.image_gen") + assert.Equal(t, 2800, usage.ImageInputTokens, "image input tokens from tool_usage.image_gen") +} + +func TestExtractOpenAIUsageFromJSONBytes_NoToolUsageUnchanged(t *testing.T) { + // Standard response without tool_usage — behavior should be unchanged + body := []byte(`{ + "usage": { + "input_tokens": 100, + "output_tokens": 50 + } + }`) + + usage, ok := extractOpenAIUsageFromJSONBytes(body) + assert.True(t, ok) + assert.Equal(t, 100, usage.InputTokens) + assert.Equal(t, 50, usage.OutputTokens) + assert.Equal(t, 0, usage.ImageOutputTokens, "no image tokens without tool_usage") + assert.Equal(t, 0, usage.ImageInputTokens, "no image tokens without tool_usage") +} + +func TestExtractOpenAIUsageFromJSONBytes_BaseUsageHasImageTokensNoOverride(t *testing.T) { + // If base usage already has image tokens (e.g. from output_tokens_details), + // tool_usage should NOT override them. + body := []byte(`{ + "usage": { + "input_tokens": 100, + "output_tokens": 50, + "output_tokens_details": {"image_tokens": 30} + }, + "tool_usage": { + "image_gen": { + "input_tokens": 200, + "output_tokens": 100, + "output_tokens_details": {"image_tokens": 100, "text_tokens": 0}, + "total_tokens": 300 + } + } + }`) + + usage, ok := extractOpenAIUsageFromJSONBytes(body) + assert.True(t, ok) + assert.Equal(t, 30, usage.ImageOutputTokens, "base usage image tokens preserved, not overridden") +} + +func TestMergeHostedImageGenToolUsage_EmptyImageGen(t *testing.T) { + tests := []struct { + name string + json string + }{ + {"missing", `{}`}, + {"null", `{"image_gen": null}`}, + {"not object", `{"image_gen": 42}`}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + usage := OpenAIUsage{InputTokens: 100, OutputTokens: 50} + original := usage + body := []byte(tt.json) + mergeHostedImageGenToolUsage(gjson.GetBytes(body, "image_gen"), &usage) + assert.Equal(t, original, usage, "usage should be unchanged") + }) + } +} + +func TestParseSSEUsageBytes_ResponseCompletedWithImageGen(t *testing.T) { + svc := &OpenAIGatewayService{} + data := []byte(`{ + "type": "response.completed", + "response": { + "usage": { + "input_tokens": 10000, + "output_tokens": 500 + }, + "tool_usage": { + "image_gen": { + "input_tokens": 4000, + "input_tokens_details": {"image_tokens": 3800, "text_tokens": 200}, + "output_tokens": 186, + "output_tokens_details": {"image_tokens": 186, "text_tokens": 0}, + "total_tokens": 4186 + } + } + } + }`) + + usage := &OpenAIUsage{} + svc.parseSSEUsageBytes(data, usage) + + assert.Equal(t, 10000, usage.InputTokens) + assert.Equal(t, 500, usage.OutputTokens) + assert.Equal(t, 186, usage.ImageOutputTokens, "image output tokens from SSE tool_usage") + assert.Equal(t, 3800, usage.ImageInputTokens, "image input tokens from SSE tool_usage") +}