From d2f243a420e0f8c16127f37863719257e64adac5 Mon Sep 17 00:00:00 2001 From: jstar0 Date: Tue, 16 Jun 2026 18:06:16 +0800 Subject: [PATCH] perf: reduce chat responses fallback usage allocation --- relay/channel/openai/chat_via_responses.go | 14 +- .../channel/openai/chat_via_responses_test.go | 179 ++++++++++++++++++ service/stream_estimate.go | 167 ++++++++++++++++ service/stream_estimate_test.go | 148 +++++++++++++++ service/usage_count_bench_test.go | 146 ++++++++++++++ 5 files changed, 647 insertions(+), 7 deletions(-) create mode 100644 relay/channel/openai/chat_via_responses_test.go create mode 100644 service/stream_estimate.go create mode 100644 service/stream_estimate_test.go create mode 100644 service/usage_count_bench_test.go diff --git a/relay/channel/openai/chat_via_responses.go b/relay/channel/openai/chat_via_responses.go index 2c0752275daa..f3a955519903 100644 --- a/relay/channel/openai/chat_via_responses.go +++ b/relay/channel/openai/chat_via_responses.go @@ -104,7 +104,7 @@ func OaiResponsesToChatStreamHandler(c *gin.Context, info *relaycommon.RelayInfo var ( usage = &dto.Usage{} outputText strings.Builder - usageText strings.Builder + usageEst = service.NewStreamingEstimateByModel(info.UpstreamModelName) sentStart bool sentStop bool sawToolCall bool @@ -167,7 +167,7 @@ func OaiResponsesToChatStreamHandler(c *gin.Context, info *relaycommon.RelayInfo // return false // } // - // usageText.WriteString(delta) + // usageEst.WriteString(delta) // chunk := &dto.ChatCompletionsStreamResponse{ // Id: responseId, // Object: "chat.completion.chunk", @@ -208,7 +208,7 @@ func OaiResponsesToChatStreamHandler(c *gin.Context, info *relaycommon.RelayInfo return false } - usageText.WriteString(delta) + usageEst.WriteString(delta) chunk := &dto.ChatCompletionsStreamResponse{ Id: responseId, Object: "chat.completion.chunk", @@ -288,10 +288,10 @@ func OaiResponsesToChatStreamHandler(c *gin.Context, info *relaycommon.RelayInfo // Include tool call data in the local builder for fallback token estimation. if tool.Function.Name != "" { - usageText.WriteString(tool.Function.Name) + usageEst.WriteString(tool.Function.Name) } if argsDelta != "" { - usageText.WriteString(argsDelta) + usageEst.WriteString(argsDelta) } return true } @@ -365,7 +365,7 @@ func OaiResponsesToChatStreamHandler(c *gin.Context, info *relaycommon.RelayInfo if streamResp.Delta != "" { outputText.WriteString(streamResp.Delta) - usageText.WriteString(streamResp.Delta) + usageEst.WriteString(streamResp.Delta) delta := streamResp.Delta chunk := &dto.ChatCompletionsStreamResponse{ Id: responseId, @@ -516,7 +516,7 @@ func OaiResponsesToChatStreamHandler(c *gin.Context, info *relaycommon.RelayInfo } if usage.TotalTokens == 0 { - usage = service.ResponseText2Usage(c, usageText.String(), info.UpstreamModelName, info.GetEstimatePromptTokens()) + usage = service.StreamingEstimate2Usage(c, usageEst, info.GetEstimatePromptTokens()) } if !sentStart { diff --git a/relay/channel/openai/chat_via_responses_test.go b/relay/channel/openai/chat_via_responses_test.go new file mode 100644 index 000000000000..4951973ca759 --- /dev/null +++ b/relay/channel/openai/chat_via_responses_test.go @@ -0,0 +1,179 @@ +package openai + +import ( + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + relaycommon "github.com/QuantumNous/new-api/relay/common" + "github.com/QuantumNous/new-api/service" + "github.com/QuantumNous/new-api/types" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func newResponsesToChatStreamTestContext(t *testing.T, body string) (*gin.Context, *httptest.ResponseRecorder, *http.Response, *relaycommon.RelayInfo) { + t.Helper() + + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil) + c.Set(common.RequestIdKey, "responses-to-chat-test") + + resp := &http.Response{ + StatusCode: http.StatusOK, + Body: io.NopCloser(strings.NewReader(body)), + Header: http.Header{"Content-Type": []string{"text/event-stream"}}, + } + info := &relaycommon.RelayInfo{ + RelayFormat: types.RelayFormatOpenAI, + ShouldIncludeUsage: true, + ChannelMeta: &relaycommon.ChannelMeta{ + UpstreamModelName: "gpt-5.5", + }, + } + info.SetEstimatePromptTokens(37) + return c, recorder, resp, info +} + +func responsesSSE(lines ...string) string { + var b strings.Builder + for _, line := range lines { + b.WriteString("data: ") + b.WriteString(line) + b.WriteString("\n\n") + } + b.WriteString("data: [DONE]\n\n") + return b.String() +} + +func TestOaiResponsesToChatStreamFallbackUsageMatchesResponseText2Usage(t *testing.T) { + oldMode := gin.Mode() + gin.SetMode(gin.TestMode) + t.Cleanup(func() { gin.SetMode(oldMode) }) + + oldTimeout := constant.StreamingTimeout + constant.StreamingTimeout = 30 + t.Cleanup(func() { constant.StreamingTimeout = oldTimeout }) + + body := responsesSSE( + `{"type":"response.created","response":{"id":"resp_test","created_at":1710000000,"model":"gpt-5.5"}}`, + `{"type":"response.reasoning_summary_text.delta","delta":"Reasoning summary 123"}`, + `{"type":"response.reasoning_summary_text.done"}`, + `{"type":"response.reasoning_summary_text.delta","delta":"second paragraph"}`, + `{"type":"response.output_text.delta","delta":" Visible output 中文"}`, + `{"type":"response.output_item.added","item":{"id":"fc_1","type":"function_call","call_id":"call_1","name":"lookup","arguments":"{\"city\":\"Bei"}}`, + `{"type":"response.function_call_arguments.delta","item_id":"fc_1","delta":"jing\"}"}`, + `{"type":"response.completed","response":{"id":"resp_test","created_at":1710000001,"model":"gpt-5.5","usage":{"input_tokens":0,"output_tokens":0,"total_tokens":0}}}`, + ) + + c, recorder, resp, info := newResponsesToChatStreamTestContext(t, body) + + usage, err := OaiResponsesToChatStreamHandler(c, info, resp) + require.Nil(t, err) + require.NotNil(t, usage) + + expectedText := "Reasoning summary 123\n\nsecond paragraph Visible output 中文" + expected := service.ResponseText2Usage(gin.CreateTestContextOnly(httptest.NewRecorder(), gin.New()), expectedText, info.UpstreamModelName, info.GetEstimatePromptTokens()) + assert.Equal(t, expected.PromptTokens, usage.PromptTokens) + assert.Equal(t, expected.CompletionTokens, usage.CompletionTokens) + assert.Equal(t, expected.TotalTokens, usage.TotalTokens) + assert.True(t, common.GetContextKeyBool(c, constant.ContextKeyLocalCountTokens)) + assert.Contains(t, recorder.Body.String(), `"usage":{"prompt_tokens":37`) +} + +func TestOaiResponsesToChatStreamFallbackUsageCountsToolCalls(t *testing.T) { + oldMode := gin.Mode() + gin.SetMode(gin.TestMode) + t.Cleanup(func() { gin.SetMode(oldMode) }) + + oldTimeout := constant.StreamingTimeout + constant.StreamingTimeout = 30 + t.Cleanup(func() { constant.StreamingTimeout = oldTimeout }) + + body := responsesSSE( + `{"type":"response.created","response":{"id":"resp_test","created_at":1710000000,"model":"gpt-5.5"}}`, + `{"type":"response.output_item.added","item":{"id":"fc_1","type":"function_call","call_id":"call_1","name":"lookup","arguments":"{\"city\":\"Bei"}}`, + `{"type":"response.function_call_arguments.delta","item_id":"fc_1","delta":"jing\"}"}`, + `{"type":"response.completed","response":{"id":"resp_test","created_at":1710000001,"model":"gpt-5.5","usage":{"input_tokens":0,"output_tokens":0,"total_tokens":0}}}`, + ) + + c, recorder, resp, info := newResponsesToChatStreamTestContext(t, body) + + usage, err := OaiResponsesToChatStreamHandler(c, info, resp) + require.Nil(t, err) + require.NotNil(t, usage) + + expectedText := `lookup{"city":"Beijing"}` + expected := service.ResponseText2Usage(gin.CreateTestContextOnly(httptest.NewRecorder(), gin.New()), expectedText, info.UpstreamModelName, info.GetEstimatePromptTokens()) + assert.Equal(t, expected.PromptTokens, usage.PromptTokens) + assert.Equal(t, expected.CompletionTokens, usage.CompletionTokens) + assert.Equal(t, expected.TotalTokens, usage.TotalTokens) + assert.True(t, common.GetContextKeyBool(c, constant.ContextKeyLocalCountTokens)) + assert.Contains(t, recorder.Body.String(), `"finish_reason":"tool_calls"`) +} + +func TestOaiResponsesToChatStreamUsesUpstreamUsageWithoutLocalCountFlag(t *testing.T) { + oldMode := gin.Mode() + gin.SetMode(gin.TestMode) + t.Cleanup(func() { gin.SetMode(oldMode) }) + + oldTimeout := constant.StreamingTimeout + constant.StreamingTimeout = 30 + t.Cleanup(func() { constant.StreamingTimeout = oldTimeout }) + + body := responsesSSE( + `{"type":"response.created","response":{"id":"resp_test","created_at":1710000000,"model":"gpt-5.5"}}`, + `{"type":"response.output_text.delta","delta":"Visible output that must not be locally counted"}`, + `{"type":"response.completed","response":{"id":"resp_test","created_at":1710000001,"model":"gpt-5.5","usage":{"input_tokens":11,"output_tokens":13,"total_tokens":24,"input_tokens_details":{"cached_tokens":7},"completion_tokens_details":{"reasoning_tokens":5}}}}`, + ) + + c, recorder, resp, info := newResponsesToChatStreamTestContext(t, body) + + usage, err := OaiResponsesToChatStreamHandler(c, info, resp) + require.Nil(t, err) + require.NotNil(t, usage) + assert.Equal(t, 11, usage.PromptTokens) + assert.Equal(t, 13, usage.CompletionTokens) + assert.Equal(t, 24, usage.TotalTokens) + assert.Equal(t, 11, usage.InputTokens) + assert.Equal(t, 13, usage.OutputTokens) + assert.Equal(t, 7, usage.PromptTokensDetails.CachedTokens) + assert.Equal(t, 5, usage.CompletionTokenDetails.ReasoningTokens) + assert.False(t, common.GetContextKeyBool(c, constant.ContextKeyLocalCountTokens)) + assert.Contains(t, recorder.Body.String(), `"usage":{"prompt_tokens":11`) +} + +func TestOaiResponsesToChatStreamCompletesUpstreamUsageTotalWithoutLocalCountFlag(t *testing.T) { + oldMode := gin.Mode() + gin.SetMode(gin.TestMode) + t.Cleanup(func() { gin.SetMode(oldMode) }) + + oldTimeout := constant.StreamingTimeout + constant.StreamingTimeout = 30 + t.Cleanup(func() { constant.StreamingTimeout = oldTimeout }) + + body := responsesSSE( + `{"type":"response.created","response":{"id":"resp_test","created_at":1710000000,"model":"gpt-5.5"}}`, + `{"type":"response.output_text.delta","delta":"Visible output with upstream input and output tokens"}`, + `{"type":"response.completed","response":{"id":"resp_test","created_at":1710000001,"model":"gpt-5.5","usage":{"input_tokens":11,"output_tokens":13,"total_tokens":0}}}`, + ) + + c, recorder, resp, info := newResponsesToChatStreamTestContext(t, body) + + usage, err := OaiResponsesToChatStreamHandler(c, info, resp) + require.Nil(t, err) + require.NotNil(t, usage) + assert.Equal(t, 11, usage.PromptTokens) + assert.Equal(t, 13, usage.CompletionTokens) + assert.Equal(t, 24, usage.TotalTokens) + assert.Equal(t, 11, usage.InputTokens) + assert.Equal(t, 13, usage.OutputTokens) + assert.False(t, common.GetContextKeyBool(c, constant.ContextKeyLocalCountTokens)) + assert.Contains(t, recorder.Body.String(), `"usage":{"prompt_tokens":11`) +} diff --git a/service/stream_estimate.go b/service/stream_estimate.go new file mode 100644 index 000000000000..f4ef41cc8c85 --- /dev/null +++ b/service/stream_estimate.go @@ -0,0 +1,167 @@ +package service + +import ( + "math" + "strings" + "unicode" + "unicode/utf8" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/dto" + "github.com/gin-gonic/gin" +) + +type streamingEstimateWordType int + +const ( + streamingEstimateNone streamingEstimateWordType = iota + streamingEstimateLatin + streamingEstimateNumber +) + +type StreamingEstimateByModel struct { + m multipliers + count float64 + currentWordType streamingEstimateWordType + pending []byte + hasInput bool +} + +func NewStreamingEstimateByModel(model string) *StreamingEstimateByModel { + model = strings.ToLower(model) + provider := OpenAI + if strings.Contains(model, "gemini") { + provider = Gemini + } else if strings.Contains(model, "claude") { + provider = Claude + } + return &StreamingEstimateByModel{m: getMultipliers(provider)} +} + +func (e *StreamingEstimateByModel) WriteString(text string) { + if e == nil || text == "" { + return + } + e.hasInput = true + if len(e.pending) > 0 { + for len(text) > 0 { + e.pending = append(e.pending, text[0]) + text = text[1:] + if utf8.FullRune(e.pending) { + break + } + } + if !utf8.FullRune(e.pending) { + return + } + r, size := utf8.DecodeRune(e.pending) + e.writeRune(r) + if size < len(e.pending) { + text = string(e.pending[size:]) + text + } + e.pending = e.pending[:0] + } + prefix, pending := splitTrailingIncompleteUTF8(text) + for _, r := range prefix { + e.writeRune(r) + } + if pending != "" { + e.pending = append(e.pending, pending...) + } +} + +func (e *StreamingEstimateByModel) Tokens() int { + if e == nil { + return 0 + } + snapshot := *e + if len(snapshot.pending) > 0 { + for _, r := range string(snapshot.pending) { + snapshot.writeRune(r) + } + } + if !snapshot.hasInput { + return 0 + } + return int(math.Ceil(snapshot.count)) + snapshot.m.BasePad +} + +func StreamingEstimate2Usage(c *gin.Context, e *StreamingEstimateByModel, promptTokens int) *dto.Usage { + common.SetContextKey(c, constant.ContextKeyLocalCountTokens, true) + usage := &dto.Usage{ + PromptTokens: promptTokens, + CompletionTokens: e.Tokens(), + } + usage.TotalTokens = usage.PromptTokens + usage.CompletionTokens + return usage +} + +func (e *StreamingEstimateByModel) writeRune(r rune) { + if unicode.IsSpace(r) { + e.currentWordType = streamingEstimateNone + if r == '\n' || r == '\t' { + e.count += e.m.Newline + } else { + e.count += e.m.Space + } + return + } + + if isCJK(r) { + e.currentWordType = streamingEstimateNone + e.count += e.m.CJK + return + } + + if isEmoji(r) { + e.currentWordType = streamingEstimateNone + e.count += e.m.Emoji + return + } + + if isLatinOrNumber(r) { + newType := streamingEstimateLatin + if unicode.IsNumber(r) { + newType = streamingEstimateNumber + } + if e.currentWordType == streamingEstimateNone || e.currentWordType != newType { + if newType == streamingEstimateNumber { + e.count += e.m.Number + } else { + e.count += e.m.Word + } + e.currentWordType = newType + } + return + } + + e.currentWordType = streamingEstimateNone + if isMathSymbol(r) { + e.count += e.m.MathSymbol + } else if r == '@' { + e.count += e.m.AtSign + } else if isURLDelim(r) { + e.count += e.m.URLDelim + } else { + e.count += e.m.Symbol + } +} + +func splitTrailingIncompleteUTF8(text string) (string, string) { + if text == "" { + return "", "" + } + start := len(text) - 1 + for start > 0 && isUTF8Continuation(text[start]) { + start-- + } + if start < len(text) && !utf8.FullRuneInString(text[start:]) { + return text[:start], text[start:] + } + return text, "" +} + +func isUTF8Continuation(b byte) bool { + return b&0xc0 == 0x80 +} diff --git a/service/stream_estimate_test.go b/service/stream_estimate_test.go new file mode 100644 index 000000000000..8edea3770b3f --- /dev/null +++ b/service/stream_estimate_test.go @@ -0,0 +1,148 @@ +package service + +import ( + "net/http/httptest" + "strings" + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestStreamingEstimateByModelMatchesEstimateTokenByModel(t *testing.T) { + cases := []struct { + name string + model string + chunks []string + }{ + { + name: "openai ascii split word", + model: "gpt-5.5", + chunks: []string{ + "Deterministic account", + "ing validation with 123", + "456 and https://example.com/a?b=c", + }, + }, + { + name: "claude cjk and symbols", + model: "claude-opus-4-8", + chunks: []string{ + "中文混合 English ", + "∑∫√ and emoji ✅", + "\nnew line", + }, + }, + { + name: "gemini long mixed", + model: "gemini-2.5-pro", + chunks: []string{ + strings.Repeat("alpha123 中文 ", 100), + strings.Repeat(" /path?x=y&z=1\n", 100), + }, + }, + { + name: "empty text remains zero", + model: "gpt-5.5", + chunks: nil, + }, + { + name: "whitespace only", + model: "gpt-5.5", + chunks: []string{ + " ", + "\n\t", + " ", + }, + }, + { + name: "symbols and mixed word boundaries", + model: "claude-opus-4-8", + chunks: []string{ + "abc", + "123", + "xyz@example.com", + " ∑∫√∞ /:?&=;#%", + }, + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + joined := strings.Join(tc.chunks, "") + estimator := NewStreamingEstimateByModel(tc.model) + for _, chunk := range tc.chunks { + estimator.WriteString(chunk) + } + got := estimator.Tokens() + want := EstimateTokenByModel(tc.model, joined) + require.Equal(t, want, got) + }) + } +} + +func TestStreamingEstimateByModelMatchesSplitUTF8(t *testing.T) { + text := "emoji ✅ 中文 ∑ math and URL https://example.com/a?b=c" + data := []byte(text) + estimator := NewStreamingEstimateByModel("claude-opus-4-8") + for i := 0; i < len(data); i++ { + estimator.WriteString(string(data[i : i+1])) + } + got := estimator.Tokens() + want := EstimateTokenByModel("claude-opus-4-8", text) + require.Equal(t, want, got) +} + +func TestStreamingEstimateByModelMatchesTrailingInvalidUTF8(t *testing.T) { + chunks := []string{"valid 中文 ", string([]byte{0xe2, 0x82})} + text := strings.Join(chunks, "") + estimator := NewStreamingEstimateByModel("gpt-5.5") + for _, chunk := range chunks { + estimator.WriteString(chunk) + } + got := estimator.Tokens() + want := EstimateTokenByModel("gpt-5.5", text) + require.Equal(t, want, got) +} + +func TestStreamingEstimateByModelMatchesPendingPlusMoreText(t *testing.T) { + data := []byte("✅abc123中文") + chunks := []string{ + string(data[:1]), + string(data[1:5]), + string(data[5:]), + } + text := strings.Join(chunks, "") + estimator := NewStreamingEstimateByModel("gemini-2.5-pro") + for _, chunk := range chunks { + estimator.WriteString(chunk) + } + got := estimator.Tokens() + want := EstimateTokenByModel("gemini-2.5-pro", text) + require.Equal(t, want, got) +} + +func TestStreamingEstimateToUsageMatchesResponseText2Usage(t *testing.T) { + text := "Reasoning summary 123\n\nVisible output 中文 with https://example.com/a?b=c" + model := "gpt-5.5" + promptTokens := 37 + + streaming := NewStreamingEstimateByModel(model) + for _, chunk := range []string{"Reasoning summary 123", "\n\nVisible output 中文", " with https://example.com/a?b=c"} { + streaming.WriteString(chunk) + } + + gotContext, _ := gin.CreateTestContext(httptest.NewRecorder()) + got := StreamingEstimate2Usage(gotContext, streaming, promptTokens) + + wantContext, _ := gin.CreateTestContext(httptest.NewRecorder()) + want := ResponseText2Usage(wantContext, text, model, promptTokens) + + assert.Equal(t, want.PromptTokens, got.PromptTokens) + assert.Equal(t, want.CompletionTokens, got.CompletionTokens) + assert.Equal(t, want.TotalTokens, got.TotalTokens) + assert.True(t, common.GetContextKeyBool(gotContext, constant.ContextKeyLocalCountTokens)) +} diff --git a/service/usage_count_bench_test.go b/service/usage_count_bench_test.go new file mode 100644 index 000000000000..b178e0e5bd70 --- /dev/null +++ b/service/usage_count_bench_test.go @@ -0,0 +1,146 @@ +package service + +import ( + "net/http/httptest" + "strings" + "testing" + + "github.com/gin-gonic/gin" +) + +var usageCountBenchSink int + +func init() { + InitTokenEncoders() + gin.SetMode(gin.TestMode) +} + +func benchmarkText(repeat int) string { + parts := []string{ + "Deterministic accounting validation compares prompt tokens, completion tokens, quota, stream status, and request identifiers.", + "中文内容用于覆盖 CJK 计数路径，避免只测英文 ASCII 文本。", + "Numbers 1234567890, URLs https://example.com/a?b=c&d=e, symbols ∑∫√∞, and emoji ✅ are included.", + } + return strings.Repeat(strings.Join(parts, "\n"), repeat) +} + +func BenchmarkEstimateTokenByModel(b *testing.B) { + cases := []struct { + name string + model string + repeat int + }{ + {name: "openai_small", model: "gpt-5.5", repeat: 1}, + {name: "openai_large", model: "gpt-5.5", repeat: 400}, + {name: "claude_small", model: "claude-sonnet-4", repeat: 1}, + {name: "claude_large", model: "claude-sonnet-4", repeat: 400}, + } + + for _, tc := range cases { + text := benchmarkText(tc.repeat) + b.Run(tc.name, func(b *testing.B) { + b.ReportAllocs() + for i := 0; i < b.N; i++ { + usageCountBenchSink = EstimateTokenByModel(tc.model, text) + } + }) + } +} + +func BenchmarkCountTextToken(b *testing.B) { + cases := []struct { + name string + model string + repeat int + }{ + {name: "openai_small", model: "gpt-5.5", repeat: 1}, + {name: "openai_large", model: "gpt-5.5", repeat: 400}, + {name: "claude_small", model: "claude-sonnet-4", repeat: 1}, + {name: "claude_large", model: "claude-sonnet-4", repeat: 400}, + } + + for _, tc := range cases { + text := benchmarkText(tc.repeat) + b.Run(tc.name, func(b *testing.B) { + b.ReportAllocs() + for i := 0; i < b.N; i++ { + usageCountBenchSink = CountTextToken(text, tc.model) + } + }) + } +} + +func BenchmarkResponseText2Usage(b *testing.B) { + cases := []struct { + name string + model string + repeat int + }{ + {name: "openai_small", model: "gpt-5.5", repeat: 1}, + {name: "openai_large", model: "gpt-5.5", repeat: 400}, + {name: "claude_small", model: "claude-sonnet-4", repeat: 1}, + {name: "claude_large", model: "claude-sonnet-4", repeat: 400}, + } + + for _, tc := range cases { + text := benchmarkText(tc.repeat) + b.Run(tc.name, func(b *testing.B) { + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + b.ReportAllocs() + for i := 0; i < b.N; i++ { + usage := ResponseText2Usage(c, text, tc.model, 123) + usageCountBenchSink = usage.TotalTokens + } + }) + } +} + +func BenchmarkChatViaResponsesFallbackUsage(b *testing.B) { + chunks := []string{ + "Reasoning summary 123", + "\n\nsecond paragraph", + " Visible output 中文", + `lookup{"city":"Beijing"}`, + " with https://example.com/a?b=c&d=e and symbols ∑∫√∞ ✅", + } + largeChunks := make([]string, 0, len(chunks)*400) + for i := 0; i < 400; i++ { + largeChunks = append(largeChunks, chunks...) + } + + cases := []struct { + name string + model string + chunks []string + }{ + {name: "old_builder_openai_large", model: "gpt-5.5", chunks: largeChunks}, + {name: "new_streaming_openai_large", model: "gpt-5.5", chunks: largeChunks}, + {name: "old_builder_claude_large", model: "claude-sonnet-4", chunks: largeChunks}, + {name: "new_streaming_claude_large", model: "claude-sonnet-4", chunks: largeChunks}, + } + + for _, tc := range cases { + b.Run(tc.name, func(b *testing.B) { + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + b.ReportAllocs() + for i := 0; i < b.N; i++ { + if strings.HasPrefix(tc.name, "old_builder") { + var usageText strings.Builder + for _, chunk := range tc.chunks { + usageText.WriteString(chunk) + } + usage := ResponseText2Usage(c, usageText.String(), tc.model, 123) + usageCountBenchSink = usage.TotalTokens + continue + } + + estimator := NewStreamingEstimateByModel(tc.model) + for _, chunk := range tc.chunks { + estimator.WriteString(chunk) + } + usage := StreamingEstimate2Usage(c, estimator, 123) + usageCountBenchSink = usage.TotalTokens + } + }) + } +}