diff --git a/dto/channel_settings.go b/dto/channel_settings.go index b6a1ab9f7138..7b701434b78e 100644 --- a/dto/channel_settings.go +++ b/dto/channel_settings.go @@ -7,6 +7,13 @@ type ChannelSettings struct { PassThroughBodyEnabled bool `json:"pass_through_body_enabled,omitempty"` SystemPrompt string `json:"system_prompt,omitempty"` SystemPromptOverride bool `json:"system_prompt_override,omitempty"` + // TrustUpstreamUsage, when enabled, makes the relay prefer the usage + // reported by the upstream in streaming responses over the locally + // streamed token count. Streaming paths no longer buffer the full + // response text regardless of this flag; the local count is always + // available as a bounded-memory fallback. Defaults to false, so the + // locally streamed count is used unless the upstream usage is trusted. + TrustUpstreamUsage bool `json:"trust_upstream_usage,omitempty"` } type VertexKeyType string diff --git a/relay/channel/aws/relay-aws.go b/relay/channel/aws/relay-aws.go index 1f6ff7e69263..e4a05b2cea83 100644 --- a/relay/channel/aws/relay-aws.go +++ b/relay/channel/aws/relay-aws.go @@ -236,7 +236,6 @@ func awsHandler(c *gin.Context, info *relaycommon.RelayInfo, a *Adaptor) (*types ResponseId: helper.GetResponseID(c), Created: common.GetTimestamp(), Model: info.UpstreamModelName, - ResponseText: strings.Builder{}, Usage: &dto.Usage{}, } @@ -268,7 +267,6 @@ func awsStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, a *Adaptor) ( ResponseId: helper.GetResponseID(c), Created: common.GetTimestamp(), Model: info.UpstreamModelName, - ResponseText: strings.Builder{}, Usage: &dto.Usage{}, } diff --git a/relay/channel/claude/relay-claude.go b/relay/channel/claude/relay-claude.go index 18d7455e9f22..cab5034af864 100644 --- a/relay/channel/claude/relay-claude.go +++ b/relay/channel/claude/relay-claude.go @@ -583,12 +583,22 @@ func ResponseClaude2OpenAI(claudeResponse *dto.ClaudeResponse) *dto.OpenAITextRe } type ClaudeResponseInfo struct { - ResponseId string - Created int64 - Model string - ResponseText strings.Builder - Usage *dto.Usage - Done bool + ResponseId string + Created int64 + Model string + // usageAcc 流式累计 completion token（text + thinking 分离计数）， + // 替代原先用 strings.Builder 累积整段响应文本再估算的做法，避免大响应内存堆积。 + usageAcc *service.UsageAccumulator + Usage *dto.Usage + Done bool +} + +// ensureUsageAcc 懒初始化 usageAcc（首次累积时按 Model 创建）。 +func (cri *ClaudeResponseInfo) ensureUsageAcc() *service.UsageAccumulator { + if cri.usageAcc == nil { + cri.usageAcc = service.NewUsageAccumulator(cri.Model) + } + return cri.usageAcc } func cacheCreationTokensForOpenAIUsage(usage *dto.Usage) int { @@ -738,10 +748,10 @@ func FormatClaudeResponseInfo(claudeResponse *dto.ClaudeResponse, oaiResponse *d } else if claudeResponse.Type == "content_block_delta" { if claudeResponse.Delta != nil { if claudeResponse.Delta.Text != nil { - claudeInfo.ResponseText.WriteString(*claudeResponse.Delta.Text) + claudeInfo.ensureUsageAcc().Feed(*claudeResponse.Delta.Text) } if claudeResponse.Delta.Thinking != nil { - claudeInfo.ResponseText.WriteString(*claudeResponse.Delta.Thinking) + claudeInfo.ensureUsageAcc().FeedReasoning(*claudeResponse.Delta.Thinking) } } } else if claudeResponse.Type == "message_delta" { @@ -840,13 +850,13 @@ func HandleStreamFinalResponse(c *gin.Context, info *relaycommon.RelayInfo, clau common.SysLog("claude response usage is not complete, maybe upstream error") } // 只补缺失字段，不整份覆盖——保留 message_start 已拿到的 cache 字段 - fallback := service.ResponseText2Usage(c, claudeInfo.ResponseText.String(), info.UpstreamModelName, info.GetEstimatePromptTokens()) + fallbackCompletion := claudeInfo.ensureUsageAcc().LocalCompletionTokens() if claudeInfo.Usage.CompletionTokens == 0 || - (!claudeInfo.Done && fallback.CompletionTokens > claudeInfo.Usage.CompletionTokens) { - claudeInfo.Usage.CompletionTokens = fallback.CompletionTokens + (!claudeInfo.Done && fallbackCompletion > claudeInfo.Usage.CompletionTokens) { + claudeInfo.Usage.CompletionTokens = fallbackCompletion } if claudeInfo.Usage.PromptTokens == 0 { - claudeInfo.Usage.PromptTokens = fallback.PromptTokens + claudeInfo.Usage.PromptTokens = info.GetEstimatePromptTokens() } claudeInfo.Usage.TotalTokens = claudeInfo.Usage.PromptTokens + claudeInfo.Usage.CompletionTokens } @@ -874,7 +884,6 @@ func ClaudeStreamHandler(c *gin.Context, resp *http.Response, info *relaycommon. ResponseId: helper.GetResponseID(c), Created: common.GetTimestamp(), Model: info.UpstreamModelName, - ResponseText: strings.Builder{}, Usage: &dto.Usage{}, } var err *types.NewAPIError @@ -943,7 +952,6 @@ func ClaudeHandler(c *gin.Context, resp *http.Response, info *relaycommon.RelayI ResponseId: helper.GetResponseID(c), Created: common.GetTimestamp(), Model: info.UpstreamModelName, - ResponseText: strings.Builder{}, Usage: &dto.Usage{}, } responseBody, err := io.ReadAll(resp.Body) diff --git a/relay/channel/claude/relay_claude_test.go b/relay/channel/claude/relay_claude_test.go index 495c500bbd13..40bf24423132 100644 --- a/relay/channel/claude/relay_claude_test.go +++ b/relay/channel/claude/relay_claude_test.go @@ -2,7 +2,6 @@ package claude import ( "encoding/base64" - "strings" "testing" "github.com/QuantumNous/new-api/dto" @@ -161,8 +160,8 @@ func TestFormatClaudeResponseInfo_NilClaudeInfo(t *testing.T) { func TestFormatClaudeResponseInfo_ContentBlockDelta(t *testing.T) { text := "hello" claudeInfo := &ClaudeResponseInfo{ - Usage: &dto.Usage{}, - ResponseText: strings.Builder{}, + Model: "claude-3-5-sonnet", + Usage: &dto.Usage{}, } claudeResponse := &dto.ClaudeResponse{ Type: "content_block_delta", @@ -175,8 +174,9 @@ func TestFormatClaudeResponseInfo_ContentBlockDelta(t *testing.T) { if !ok { t.Fatal("expected true") } - if claudeInfo.ResponseText.String() != "hello" { - t.Errorf("ResponseText = %q, want %q", claudeInfo.ResponseText.String(), "hello") + // 文本通过流式累计器计数，应反映已喂入的 "hello" + if got := claudeInfo.ensureUsageAcc().LocalCompletionTokens(); got <= 0 { + t.Errorf("LocalCompletionTokens = %d, want > 0 for %q", got, text) } } diff --git a/relay/channel/claude/stream_handler_test.go b/relay/channel/claude/stream_handler_test.go new file mode 100644 index 000000000000..f993a67b60dd --- /dev/null +++ b/relay/channel/claude/stream_handler_test.go @@ -0,0 +1,147 @@ +package claude + +import ( + "io" + "net/http" + "net/http/httptest" + "os" + "strings" + "testing" + + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/dto" + relaycommon "github.com/QuantumNous/new-api/relay/common" + "github.com/QuantumNous/new-api/types" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +func TestMain(m *testing.M) { + // StreamScannerHandler uses time.NewTicker(StreamingTimeout); avoid zero-interval panic. + if constant.StreamingTimeout <= 0 { + constant.StreamingTimeout = 300 + } + os.Exit(m.Run()) +} + +func newStreamTestInfo(model string, trust bool) *relaycommon.RelayInfo { + info := &relaycommon.RelayInfo{ + RelayFormat: types.RelayFormatClaude, + ChannelMeta: &relaycommon.ChannelMeta{ + UpstreamModelName: model, + ChannelSetting: dto.ChannelSettings{TrustUpstreamUsage: trust}, + }, + } + info.SetEstimatePromptTokens(100) + return info +} + +func newSSEResp(sse string) *http.Response { + return &http.Response{ + Body: io.NopCloser(strings.NewReader(sse)), + StatusCode: http.StatusOK, + } +} + +func newStreamTestCtx() *gin.Context { + gin.SetMode(gin.TestMode) + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil) + return c +} + +// 上游提供 usage 且 trust=true：应直接用上游 usage（含 output_tokens）。 +func TestClaudeStreamHandler_TrustUpstreamUsage(t *testing.T) { + c := newStreamTestCtx() + info := newStreamTestInfo("claude-3-5-sonnet", true) + sse := `data: {"type":"message_start","message":{"id":"msg_1","model":"claude-3-5-sonnet","usage":{"input_tokens":100,"output_tokens":1}}} +data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hello world this is the answer"}} +data: {"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"input_tokens":100,"output_tokens":42}} +data: {"type":"message_stop"} +` + usage, apiErr := ClaudeStreamHandler(c, newSSEResp(sse), info) + require.Nil(t, apiErr) + require.NotNil(t, usage) + require.Equal(t, 100, usage.PromptTokens) + require.Equal(t, 42, usage.CompletionTokens, "trust=true 应采用上游 output_tokens=42") +} + +// 上游未给 output_tokens（异常/中断）：应回退到本地流式估算，且不为 0。 +func TestClaudeStreamHandler_LocalFallback(t *testing.T) { + c := newStreamTestCtx() + info := newStreamTestInfo("claude-3-5-sonnet", false) + // message_delta 不带 usage，message_stop 前断；本地需要根据文本估算 + sse := `data: {"type":"message_start","message":{"id":"msg_1","model":"claude-3-5-sonnet","usage":{"input_tokens":100,"output_tokens":0}}} +data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hello world this is a fairly long answer text"}} +` + usage, apiErr := ClaudeStreamHandler(c, newSSEResp(sse), info) + require.Nil(t, apiErr) + require.NotNil(t, usage) + require.Greater(t, usage.CompletionTokens, 0, "上游未给 usage 时本地估算应 > 0") +} + +// thinking + text 分离计数：thinking 也应计入 completion。 +func TestClaudeStreamHandler_ThinkingCounted(t *testing.T) { + c := newStreamTestCtx() + info := newStreamTestInfo("claude-3-5-sonnet", false) + textOnly := `data: {"type":"message_start","message":{"id":"m","model":"claude-3-5-sonnet","usage":{"input_tokens":10,"output_tokens":0}}} +data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"visible answer"}} +` + withThinking := `data: {"type":"message_start","message":{"id":"m","model":"claude-3-5-sonnet","usage":{"input_tokens":10,"output_tokens":0}}} +data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"visible answer"}} +data: {"type":"content_block_delta","index":0,"delta":{"type":"thinking_delta","thinking":"some internal reasoning content here that is fairly long"}} +` + u1, e1 := ClaudeStreamHandler(newStreamTestCtx(), newSSEResp(textOnly), info) + require.Nil(t, e1) + u2, e2 := ClaudeStreamHandler(c, newSSEResp(withThinking), newStreamTestInfo("claude-3-5-sonnet", false)) + require.Nil(t, e2) + require.Greater(t, u2.CompletionTokens, u1.CompletionTokens, "带 thinking 的 completion 应更大（thinking 被计入）") +} + +// cache 字段（read/creation）必须从 message_start 正确传递到最终 usage， +// 不被本次累积重构破坏。 +func TestClaudeStreamHandler_CacheTokensPreserved(t *testing.T) { + c := newStreamTestCtx() + info := newStreamTestInfo("claude-3-5-sonnet", true) + sse := `data: {"type":"message_start","message":{"id":"m","model":"claude-3-5-sonnet","usage":{"input_tokens":50,"output_tokens":1,"cache_read_input_tokens":4096,"cache_creation_input_tokens":256}}} +data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"answer text here"}} +data: {"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"input_tokens":50,"output_tokens":20}} +data: {"type":"message_stop"} +` + usage, apiErr := ClaudeStreamHandler(c, newSSEResp(sse), info) + require.Nil(t, apiErr) + require.NotNil(t, usage) + require.Equal(t, 4096, usage.PromptTokensDetails.CachedTokens, "cache_read 应保留") + require.Equal(t, 256, usage.PromptTokensDetails.CachedCreationTokens, "cache_creation 应保留") + require.Equal(t, 20, usage.CompletionTokens, "trust=true 用上游 output_tokens") +} + +// 使用从生产 sub2api 抓取的【真实】Claude SSE 响应（含 event: 行、ping、 +// cache_creation 嵌套结构、末尾空白），验证 handler 在真实上游格式下正确工作。 +// 这不是臆想的格式——是 2026-06 实际抓包内容（已脱敏 id）。 +func TestClaudeStreamHandler_RealUpstreamFormat(t *testing.T) { + c := newStreamTestCtx() + info := newStreamTestInfo("claude-haiku-4-5", true) + // 注意：真实流每个 data 行后有尾随空格、event: 行穿插、有 ping 事件。 + sse := "event: message_start\n" + + `data: {"type":"message_start","message":{"model":"claude-haiku-4-5","id":"msg_x","type":"message","role":"assistant","content":[],"usage":{"input_tokens":8,"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"cache_creation":{"ephemeral_5m_input_tokens":0,"ephemeral_1h_input_tokens":0},"output_tokens":1}} }` + "\n\n" + + "event: content_block_start\n" + + `data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""} }` + "\n\n" + + "event: ping\n" + + `data: {"type": "ping"}` + "\n\n" + + "event: content_block_delta\n" + + `data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hey"} }` + "\n\n" + + "event: content_block_delta\n" + + `data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"! How's it going?"} }` + "\n\n" + + "event: content_block_stop\n" + + `data: {"type":"content_block_stop","index":0 }` + "\n\n" + + "event: message_delta\n" + + `data: {"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"input_tokens":8,"output_tokens":11}}` + "\n\n" + + "event: message_stop\n" + + `data: {"type":"message_stop"}` + "\n\n" + usage, apiErr := ClaudeStreamHandler(c, newSSEResp(sse), info) + require.Nil(t, apiErr) + require.NotNil(t, usage) + require.Equal(t, 8, usage.PromptTokens, "真实 message_start input_tokens") + require.Equal(t, 11, usage.CompletionTokens, "真实 message_delta output_tokens (trust=true)") +} diff --git a/relay/channel/cloudflare/relay_cloudflare.go b/relay/channel/cloudflare/relay_cloudflare.go index 589ff1269689..2f2768a95118 100644 --- a/relay/channel/cloudflare/relay_cloudflare.go +++ b/relay/channel/cloudflare/relay_cloudflare.go @@ -35,7 +35,7 @@ func cfStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Res helper.SetEventStreamHeaders(c) id := helper.GetResponseID(c) - var responseText string + usageAcc := service.NewUsageAccumulator(info.UpstreamModelName) isFirst := true for scanner.Scan() { @@ -58,7 +58,7 @@ func cfStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Res } for _, choice := range response.Choices { choice.Delta.Role = "assistant" - responseText += choice.Delta.GetContentString() + usageAcc.Feed(choice.Delta.GetContentString()) } response.Id = id response.Model = info.UpstreamModelName @@ -75,7 +75,11 @@ func cfStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Res if err := scanner.Err(); err != nil { logger.LogError(c, "error_scanning_stream_response: "+err.Error()) } - usage := service.ResponseText2Usage(c, responseText, info.UpstreamModelName, info.GetEstimatePromptTokens()) + usage := &dto.Usage{ + PromptTokens: info.GetEstimatePromptTokens(), + CompletionTokens: usageAcc.LocalCompletionTokens(), + } + usage.TotalTokens = usage.PromptTokens + usage.CompletionTokens if info.ShouldIncludeUsage { response := helper.GenerateFinalUsageResponse(id, info.StartTime.Unix(), info.UpstreamModelName, *usage) err := helper.ObjectData(c, response) diff --git a/relay/channel/cloudflare/stream_handler_test.go b/relay/channel/cloudflare/stream_handler_test.go new file mode 100644 index 000000000000..a9ee5ee826dd --- /dev/null +++ b/relay/channel/cloudflare/stream_handler_test.go @@ -0,0 +1,62 @@ +package cloudflare + +import ( + "io" + "net/http" + "net/http/httptest" + "os" + "strings" + "testing" + "time" + + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/dto" + relaycommon "github.com/QuantumNous/new-api/relay/common" + "github.com/QuantumNous/new-api/types" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +func TestMain(m *testing.M) { + if constant.StreamingTimeout <= 0 { + constant.StreamingTimeout = 300 + } + os.Exit(m.Run()) +} + +func streamInfo(model string) *relaycommon.RelayInfo { + info := &relaycommon.RelayInfo{ + RelayFormat: types.RelayFormatOpenAI, + StartTime: time.Now(), + ChannelMeta: &relaycommon.ChannelMeta{ + UpstreamModelName: model, + ChannelSetting: dto.ChannelSettings{}, + }, + } + info.SetEstimatePromptTokens(100) + return info +} + +func sseResp(s string) *http.Response { + return &http.Response{Body: io.NopCloser(strings.NewReader(s)), StatusCode: http.StatusOK} +} + +func streamCtx() *gin.Context { + gin.SetMode(gin.TestMode) + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil) + return c +} + +// cloudflare 流式（无上游 usage）：本地估算 completion > 0。注意返回值顺序 (error, usage)。 +func TestCfStreamHandler_LocalEstimate(t *testing.T) { + sse := `data: {"id":"x","choices":[{"delta":{"role":"assistant","content":"Hello world this is"}}]} +data: {"id":"x","choices":[{"delta":{"content":" a cloudflare answer"}}]} +data: [DONE] +` + apiErr, usage := cfStreamHandler(streamCtx(), streamInfo("@cf/meta/llama"), sseResp(sse)) + require.Nil(t, apiErr) + require.NotNil(t, usage) + require.Greater(t, usage.CompletionTokens, 0) + require.Equal(t, 100, usage.PromptTokens) +} diff --git a/relay/channel/cohere/relay-cohere.go b/relay/channel/cohere/relay-cohere.go index 7b47789f869d..12da8f7ba57e 100644 --- a/relay/channel/cohere/relay-cohere.go +++ b/relay/channel/cohere/relay-cohere.go @@ -84,7 +84,7 @@ func cohereStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http responseId := helper.GetResponseID(c) createdTime := common.GetTimestamp() usage := &dto.Usage{} - responseText := "" + usageAcc := service.NewUsageAccumulator(info.UpstreamModelName) scanner := helper.NewStreamScanner(resp.Body) scanner.Split(func(data []byte, atEOF bool) (advance int, token []byte, err error) { if atEOF && len(data) == 0 { @@ -154,7 +154,7 @@ func cohereStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http Index: 0, }, } - responseText += cohereResp.Text + usageAcc.Feed(cohereResp.Text) } jsonStr, err := json.Marshal(openaiResp) if err != nil { @@ -169,8 +169,10 @@ func cohereStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http } }) if usage.PromptTokens == 0 { - usage = service.ResponseText2Usage(c, responseText, info.UpstreamModelName, info.GetEstimatePromptTokens()) + usage.PromptTokens = info.GetEstimatePromptTokens() + usage.CompletionTokens = usageAcc.LocalCompletionTokens() } + usage.TotalTokens = usage.PromptTokens + usage.CompletionTokens return usage, nil } diff --git a/relay/channel/cohere/stream_handler_test.go b/relay/channel/cohere/stream_handler_test.go new file mode 100644 index 000000000000..3cfcea1f2117 --- /dev/null +++ b/relay/channel/cohere/stream_handler_test.go @@ -0,0 +1,87 @@ +package cohere + +import ( + "io" + "net/http" + "net/http/httptest" + "os" + "strings" + "testing" + "time" + + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/dto" + relaycommon "github.com/QuantumNous/new-api/relay/common" + "github.com/QuantumNous/new-api/types" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +func TestMain(m *testing.M) { + if constant.StreamingTimeout <= 0 { + constant.StreamingTimeout = 300 + } + os.Exit(m.Run()) +} + +func streamInfo(model string) *relaycommon.RelayInfo { + info := &relaycommon.RelayInfo{ + RelayFormat: types.RelayFormatOpenAI, + StartTime: time.Now(), + ChannelMeta: &relaycommon.ChannelMeta{ + UpstreamModelName: model, + ChannelSetting: dto.ChannelSettings{}, + }, + } + info.SetEstimatePromptTokens(100) + return info +} + +func sseResp(s string) *http.Response { + return &http.Response{Body: io.NopCloser(strings.NewReader(s)), StatusCode: http.StatusOK} +} + +// closeNotifyRecorder 包装 httptest.ResponseRecorder 以满足 gin c.Stream 需要的 +// http.CloseNotifier 接口（ResponseRecorder 本身不实现）。 +type closeNotifyRecorder struct { + *httptest.ResponseRecorder + closed chan bool +} + +func newCloseNotifyRecorder() *closeNotifyRecorder { + return &closeNotifyRecorder{httptest.NewRecorder(), make(chan bool, 1)} +} + +func (c *closeNotifyRecorder) CloseNotify() <-chan bool { return c.closed } + +func streamCtx() *gin.Context { + gin.SetMode(gin.TestMode) + c, _ := gin.CreateTestContext(newCloseNotifyRecorder()) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat", nil) + return c +} + +// 上游 stream-end 帧带 billed_units（prompt/completion）但不含 total： +// 最终返回的 usage.TotalTokens 必须等于 prompt+completion，不能停在 0。 +// 复现并守护 #2：上游提供 usage 时跳过本地回退块，TotalTokens 之前会遗漏。 +func TestCohereStreamHandler_UpstreamUsageTotalTokens(t *testing.T) { + sse := `{"is_finished":false,"event_type":"text-generation","text":"Hello world answer"} +{"is_finished":true,"event_type":"stream-end","finish_reason":"COMPLETE","response":{"meta":{"billed_units":{"input_tokens":12,"output_tokens":5}}}}` + usage, apiErr := cohereStreamHandler(streamCtx(), streamInfo("command-r-plus"), sseResp(sse)) + require.Nil(t, apiErr) + require.NotNil(t, usage) + require.Equal(t, 12, usage.PromptTokens) + require.Equal(t, 5, usage.CompletionTokens) + require.Equal(t, 17, usage.TotalTokens, "上游有 usage 时 TotalTokens 必须 = prompt+completion") +} + +// cohere 流式（JSON-per-line），上游 finish 帧不带 usage prompt → 本地估算 > 0。 +func TestCohereStreamHandler_LocalEstimate(t *testing.T) { + sse := `{"is_finished":false,"event_type":"text-generation","text":"Hello world this is"} +{"is_finished":false,"event_type":"text-generation","text":" a generated answer"} +{"is_finished":true,"event_type":"stream-end","finish_reason":"COMPLETE"}` + usage, apiErr := cohereStreamHandler(streamCtx(), streamInfo("command-r-plus"), sseResp(sse)) + require.Nil(t, apiErr) + require.NotNil(t, usage) + require.Greater(t, usage.CompletionTokens, 0) +} diff --git a/relay/channel/coze/relay-coze.go b/relay/channel/coze/relay-coze.go index c2db5c609cde..923677ce7d6e 100644 --- a/relay/channel/coze/relay-coze.go +++ b/relay/channel/coze/relay-coze.go @@ -102,7 +102,7 @@ func cozeChatStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *ht scanner.Split(bufio.ScanLines) helper.SetEventStreamHeaders(c) id := helper.GetResponseID(c) - var responseText string + usageAcc := service.NewUsageAccumulator(info.UpstreamModelName) var currentEvent string var currentData string @@ -114,7 +114,7 @@ func cozeChatStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *ht if line == "" { if currentEvent != "" && currentData != "" { // handle last event - handleCozeEvent(c, currentEvent, currentData, &responseText, usage, id, info) + handleCozeEvent(c, currentEvent, currentData, usageAcc, usage, id, info) currentEvent = "" currentData = "" } @@ -134,7 +134,7 @@ func cozeChatStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *ht // Last event if currentEvent != "" && currentData != "" { - handleCozeEvent(c, currentEvent, currentData, &responseText, usage, id, info) + handleCozeEvent(c, currentEvent, currentData, usageAcc, usage, id, info) } if err := scanner.Err(); err != nil { @@ -143,13 +143,15 @@ func cozeChatStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *ht helper.Done(c) if usage.TotalTokens == 0 { - usage = service.ResponseText2Usage(c, responseText, info.UpstreamModelName, c.GetInt("coze_input_count")) + usage.PromptTokens = c.GetInt("coze_input_count") + usage.CompletionTokens = usageAcc.LocalCompletionTokens() + usage.TotalTokens = usage.PromptTokens + usage.CompletionTokens } return usage, nil } -func handleCozeEvent(c *gin.Context, event string, data string, responseText *string, usage *dto.Usage, id string, info *relaycommon.RelayInfo) { +func handleCozeEvent(c *gin.Context, event string, data string, usageAcc *service.UsageAccumulator, usage *dto.Usage, id string, info *relaycommon.RelayInfo) { switch event { case "conversation.chat.completed": // 将 data 解析为 CozeChatResponseData @@ -184,7 +186,7 @@ func handleCozeEvent(c *gin.Context, event string, data string, responseText *st return } - *responseText += content + usageAcc.Feed(content) openaiResponse := dto.ChatCompletionsStreamResponse{ Id: id, diff --git a/relay/channel/coze/stream_handler_test.go b/relay/channel/coze/stream_handler_test.go new file mode 100644 index 000000000000..f7bc7a368838 --- /dev/null +++ b/relay/channel/coze/stream_handler_test.go @@ -0,0 +1,81 @@ +package coze + +import ( + "io" + "net/http" + "net/http/httptest" + "os" + "strings" + "testing" + "time" + + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/dto" + relaycommon "github.com/QuantumNous/new-api/relay/common" + "github.com/QuantumNous/new-api/types" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +func TestMain(m *testing.M) { + if constant.StreamingTimeout <= 0 { + constant.StreamingTimeout = 300 + } + os.Exit(m.Run()) +} + +func streamInfo(model string) *relaycommon.RelayInfo { + info := &relaycommon.RelayInfo{ + RelayFormat: types.RelayFormatOpenAI, + StartTime: time.Now(), + ChannelMeta: &relaycommon.ChannelMeta{ + UpstreamModelName: model, + ChannelSetting: dto.ChannelSettings{}, + }, + } + info.SetEstimatePromptTokens(100) + return info +} + +func sseResp(s string) *http.Response { + return &http.Response{Body: io.NopCloser(strings.NewReader(s)), StatusCode: http.StatusOK} +} + +type closeNotifyRecorder struct { + *httptest.ResponseRecorder + closed chan bool +} + +func (c *closeNotifyRecorder) CloseNotify() <-chan bool { return c.closed } + +func streamCtx() *gin.Context { + gin.SetMode(gin.TestMode) + rec := &closeNotifyRecorder{httptest.NewRecorder(), make(chan bool, 1)} + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/v3/chat", nil) + c.Set("coze_input_count", 10) + return c +} + +// coze 上游 completed 帧带 usage：采用上游。 +func TestCozeStreamHandler_UpstreamUsage(t *testing.T) { + sse := "event: conversation.message.delta\n" + + `data: {"role":"assistant","type":"text","content":"\"Hello world\""}` + "\n\n" + + "event: conversation.chat.completed\n" + + `data: {"id":"chat_x","usage":{"token_count":15,"output_count":5,"input_count":10}}` + "\n\n" + usage, apiErr := cozeChatStreamHandler(streamCtx(), streamInfo("coze-bot"), sseResp(sse)) + require.Nil(t, apiErr) + require.NotNil(t, usage) + require.Equal(t, 5, usage.CompletionTokens, "应采用上游 completed 的 output_count") +} + +// coze 无 completed usage：本地估算 > 0，prompt 来自 coze_input_count。 +func TestCozeStreamHandler_LocalFallback(t *testing.T) { + sse := "event: conversation.message.delta\n" + + `data: {"role":"assistant","type":"text","content":"\"Hello world this is a fairly long coze generated answer\""}` + "\n\n" + usage, apiErr := cozeChatStreamHandler(streamCtx(), streamInfo("coze-bot"), sseResp(sse)) + require.Nil(t, apiErr) + require.NotNil(t, usage) + require.Greater(t, usage.CompletionTokens, 0) + require.Equal(t, 10, usage.PromptTokens, "prompt 应来自 coze_input_count") +} diff --git a/relay/channel/dify/relay-dify.go b/relay/channel/dify/relay-dify.go index 263ad66837b6..32dbc749a035 100644 --- a/relay/channel/dify/relay-dify.go +++ b/relay/channel/dify/relay-dify.go @@ -224,7 +224,7 @@ func streamResponseDify2OpenAI(difyResponse DifyChunkChatCompletionResponse) *dt } func difyStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.NewAPIError) { - var responseText string + usageAcc := service.NewUsageAccumulator(info.UpstreamModelName) usage := &dto.Usage{} var nodeToken int helper.SetEventStreamHeaders(c) @@ -245,7 +245,7 @@ func difyStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.R } openaiResponse := *streamResponseDify2OpenAI(difyResponse) if len(openaiResponse.Choices) != 0 { - responseText += openaiResponse.Choices[0].Delta.GetContentString() + usageAcc.Feed(openaiResponse.Choices[0].Delta.GetContentString()) if openaiResponse.Choices[0].Delta.ReasoningContent != nil { nodeToken += 1 } @@ -257,9 +257,12 @@ func difyStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.R }) helper.Done(c) if usage.TotalTokens == 0 { - usage = service.ResponseText2Usage(c, responseText, info.UpstreamModelName, info.GetEstimatePromptTokens()) + usage.PromptTokens = info.GetEstimatePromptTokens() + usage.CompletionTokens = usageAcc.LocalCompletionTokens() + usage.TotalTokens = usage.PromptTokens + usage.CompletionTokens } usage.CompletionTokens += nodeToken + usage.TotalTokens = usage.PromptTokens + usage.CompletionTokens return usage, nil } diff --git a/relay/channel/dify/stream_handler_test.go b/relay/channel/dify/stream_handler_test.go new file mode 100644 index 000000000000..e57fd4487c92 --- /dev/null +++ b/relay/channel/dify/stream_handler_test.go @@ -0,0 +1,89 @@ +package dify + +import ( + "io" + "net/http" + "net/http/httptest" + "os" + "strings" + "testing" + + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/dto" + relaycommon "github.com/QuantumNous/new-api/relay/common" + "github.com/QuantumNous/new-api/types" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +func TestMain(m *testing.M) { + if constant.StreamingTimeout <= 0 { + constant.StreamingTimeout = 300 + } + os.Exit(m.Run()) +} + +func streamInfo(model string) *relaycommon.RelayInfo { + info := &relaycommon.RelayInfo{ + RelayFormat: types.RelayFormatOpenAI, + ChannelMeta: &relaycommon.ChannelMeta{ + UpstreamModelName: model, + ChannelSetting: dto.ChannelSettings{}, + }, + } + info.SetEstimatePromptTokens(100) + return info +} + +func sseResp(s string) *http.Response { + return &http.Response{Body: io.NopCloser(strings.NewReader(s)), StatusCode: http.StatusOK} +} + +func streamCtx() *gin.Context { + gin.SetMode(gin.TestMode) + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat-messages", nil) + return c +} + +// dify 上游 message_end 带 usage：采用上游。 +func TestDifyStreamHandler_UpstreamUsage(t *testing.T) { + sse := `data: {"event":"message","answer":"Hello world answer"} +data: {"event":"message","answer":" more"} +data: {"event":"message_end","metadata":{"usage":{"prompt_tokens":10,"completion_tokens":6,"total_tokens":16}}} +` + usage, apiErr := difyStreamHandler(streamCtx(), streamInfo("dify-app"), sseResp(sse)) + require.Nil(t, apiErr) + require.NotNil(t, usage) + require.Equal(t, 6, usage.CompletionTokens, "应采用上游 message_end 的 usage") +} + +// dify 无 message_end usage：本地估算 > 0。 +func TestDifyStreamHandler_LocalFallback(t *testing.T) { + sse := `data: {"event":"message","answer":"Hello world this is a fairly long dify generated answer text"} +` + usage, apiErr := difyStreamHandler(streamCtx(), streamInfo("dify-app"), sseResp(sse)) + require.Nil(t, apiErr) + require.NotNil(t, usage) + require.Greater(t, usage.CompletionTokens, 0) +} + +// 上游 message_end 带完整 usage（含 total），其后 node_finished 帧又给 nodeToken +// 补偿到 completion：最终 TotalTokens 必须把 nodeToken 一并计入，不能停在上游的旧 total。 +// 复现并守护 #3：nodeToken 之前在 TotalTokens 算定后才加，导致 total 少算。 +func TestDifyStreamHandler_NodeTokenIncludedInTotal(t *testing.T) { + prev := constant.DifyDebug + constant.DifyDebug = true // node_finished 仅在 debug 下产生 reasoning → nodeToken + defer func() { constant.DifyDebug = prev }() + + sse := `data: {"event":"message","answer":"Hello answer"} +data: {"event":"node_finished","data":{"node_type":"llm","status":"succeeded"}} +data: {"event":"message_end","metadata":{"usage":{"prompt_tokens":10,"completion_tokens":6,"total_tokens":16}}} +` + usage, apiErr := difyStreamHandler(streamCtx(), streamInfo("dify-app"), sseResp(sse)) + require.Nil(t, apiErr) + require.NotNil(t, usage) + require.Equal(t, 7, usage.CompletionTokens, "completion 应为上游 6 + nodeToken 1") + require.Equal(t, usage.PromptTokens+usage.CompletionTokens, usage.TotalTokens, + "TotalTokens 必须 = prompt + (completion+nodeToken)，不能停在上游旧 total") +} diff --git a/relay/channel/gemini/relay-gemini.go b/relay/channel/gemini/relay-gemini.go index e39826dd64e7..fc2c76587bc0 100644 --- a/relay/channel/gemini/relay-gemini.go +++ b/relay/channel/gemini/relay-gemini.go @@ -1344,7 +1344,7 @@ func handleFinalStream(c *gin.Context, info *relaycommon.RelayInfo, resp *dto.Ch func geminiStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response, callback func(data string, geminiResponse *dto.GeminiChatResponse) bool) (*dto.Usage, *types.NewAPIError) { var usage = &dto.Usage{} var imageCount int - responseText := strings.Builder{} + usageAcc := service.NewUsageAccumulator(info.UpstreamModelName) helper.StreamScannerHandler(c, resp, info, func(data string, sr *helper.StreamResult) { var geminiResponse dto.GeminiChatResponse @@ -1364,7 +1364,7 @@ func geminiStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http imageCount++ } if part.Text != "" { - responseText.WriteString(part.Text) + usageAcc.Feed(part.Text) } } } @@ -1388,7 +1388,9 @@ func geminiStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http if usage.CompletionTokens <= 0 { if info.ReceivedResponseCount > 0 { - usage = service.ResponseText2Usage(c, responseText.String(), info.UpstreamModelName, info.GetEstimatePromptTokens()) + usage.PromptTokens = info.GetEstimatePromptTokens() + usage.CompletionTokens = usageAcc.LocalCompletionTokens() + usage.TotalTokens = usage.PromptTokens + usage.CompletionTokens } else { usage = &dto.Usage{} } diff --git a/relay/channel/gemini/stream_handler_test.go b/relay/channel/gemini/stream_handler_test.go new file mode 100644 index 000000000000..4c211a8cf9b1 --- /dev/null +++ b/relay/channel/gemini/stream_handler_test.go @@ -0,0 +1,81 @@ +package gemini + +import ( + "io" + "net/http" + "net/http/httptest" + "os" + "strings" + "testing" + + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/dto" + relaycommon "github.com/QuantumNous/new-api/relay/common" + "github.com/QuantumNous/new-api/types" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +func TestMain(m *testing.M) { + if constant.StreamingTimeout <= 0 { + constant.StreamingTimeout = 300 + } + os.Exit(m.Run()) +} + +func streamInfo(model string) *relaycommon.RelayInfo { + info := &relaycommon.RelayInfo{ + RelayFormat: types.RelayFormatOpenAI, + ChannelMeta: &relaycommon.ChannelMeta{ + UpstreamModelName: model, + ChannelSetting: dto.ChannelSettings{}, + }, + } + info.SetEstimatePromptTokens(100) + return info +} + +func sseResp(s string) *http.Response { + return &http.Response{Body: io.NopCloser(strings.NewReader(s)), StatusCode: http.StatusOK} +} + +func streamCtx() *gin.Context { + gin.SetMode(gin.TestMode) + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + c.Request = httptest.NewRequest(http.MethodPost, "/v1beta/models/x:streamGenerateContent", nil) + return c +} + +// 上游 usageMetadata 带 candidatesTokenCount：采用上游。 +func TestGeminiStreamHandler_UpstreamUsage(t *testing.T) { + sse := `data: {"candidates":[{"content":{"parts":[{"text":"Hello world answer"}]}}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":9,"totalTokenCount":19}} +data: {"candidates":[{"content":{"parts":[{"text":" more text"}]}}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":12,"totalTokenCount":22}} +data: [DONE] +` + usage, apiErr := GeminiChatStreamHandler(streamCtx(), streamInfo("gemini-1.5-flash"), sseResp(sse)) + require.Nil(t, apiErr) + require.NotNil(t, usage) + require.Equal(t, 12, usage.CompletionTokens, "应采用上游最后的 candidatesTokenCount=12") +} + +// 上游无 usageMetadata：本地估算 > 0。 +func TestGeminiStreamHandler_LocalFallback(t *testing.T) { + sse := `data: {"candidates":[{"content":{"parts":[{"text":"Hello world this is a fairly long generated gemini answer text"}]}}]} +data: [DONE] +` + usage, apiErr := GeminiChatStreamHandler(streamCtx(), streamInfo("gemini-1.5-flash"), sseResp(sse)) + require.Nil(t, apiErr) + require.NotNil(t, usage) + require.Greater(t, usage.CompletionTokens, 0) +} + +// 图片输出：无上游 usage 时按 imageCount*1400 计 completion。 +func TestGeminiStreamHandler_ImageCount(t *testing.T) { + sse := `data: {"candidates":[{"content":{"parts":[{"inlineData":{"mimeType":"image/png","data":"iVBORw0KGgo="}}]}}]} +data: [DONE] +` + usage, apiErr := GeminiChatStreamHandler(streamCtx(), streamInfo("gemini-2-flash-image"), sseResp(sse)) + require.Nil(t, apiErr) + require.NotNil(t, usage) + require.Equal(t, 1400, usage.CompletionTokens, "1 张图应计 1400 token") +} diff --git a/relay/channel/openai/chat_via_responses.go b/relay/channel/openai/chat_via_responses.go index 2c0752275daa..0d0fb0918e55 100644 --- a/relay/channel/openai/chat_via_responses.go +++ b/relay/channel/openai/chat_via_responses.go @@ -103,8 +103,8 @@ func OaiResponsesToChatStreamHandler(c *gin.Context, info *relaycommon.RelayInfo var ( usage = &dto.Usage{} - outputText strings.Builder - usageText strings.Builder + outputLen int // 累计输出文本长度（替代 outputText.Len()，不缓冲全文） + usageAcc = service.NewUsageAccumulator(info.UpstreamModelName) sentStart bool sentStop bool sawToolCall bool @@ -208,7 +208,7 @@ func OaiResponsesToChatStreamHandler(c *gin.Context, info *relaycommon.RelayInfo return false } - usageText.WriteString(delta) + usageAcc.Feed(delta) chunk := &dto.ChatCompletionsStreamResponse{ Id: responseId, Object: "chat.completion.chunk", @@ -234,7 +234,7 @@ func OaiResponsesToChatStreamHandler(c *gin.Context, info *relaycommon.RelayInfo if callID == "" { return true } - if outputText.Len() > 0 { + if outputLen > 0 { // Prefer streaming assistant text over tool calls to match non-stream behavior. return true } @@ -286,12 +286,12 @@ func OaiResponsesToChatStreamHandler(c *gin.Context, info *relaycommon.RelayInfo } sawToolCall = true - // Include tool call data in the local builder for fallback token estimation. + // Include tool call data in the local accumulator for fallback token estimation. if tool.Function.Name != "" { - usageText.WriteString(tool.Function.Name) + usageAcc.Feed(tool.Function.Name) } if argsDelta != "" { - usageText.WriteString(argsDelta) + usageAcc.Feed(argsDelta) } return true } @@ -364,8 +364,8 @@ func OaiResponsesToChatStreamHandler(c *gin.Context, info *relaycommon.RelayInfo } if streamResp.Delta != "" { - outputText.WriteString(streamResp.Delta) - usageText.WriteString(streamResp.Delta) + outputLen += len(streamResp.Delta) + usageAcc.Feed(streamResp.Delta) delta := streamResp.Delta chunk := &dto.ChatCompletionsStreamResponse{ Id: responseId, @@ -484,7 +484,7 @@ func OaiResponsesToChatStreamHandler(c *gin.Context, info *relaycommon.RelayInfo info.ClaudeConvertInfo.Usage = usage } finishReason := "stop" - if sawToolCall && outputText.Len() == 0 { + if sawToolCall && outputLen == 0 { finishReason = "tool_calls" } stop := helper.GenerateStopResponse(responseId, createAt, model, finishReason) @@ -516,7 +516,9 @@ func OaiResponsesToChatStreamHandler(c *gin.Context, info *relaycommon.RelayInfo } if usage.TotalTokens == 0 { - usage = service.ResponseText2Usage(c, usageText.String(), info.UpstreamModelName, info.GetEstimatePromptTokens()) + usage.PromptTokens = info.GetEstimatePromptTokens() + usage.CompletionTokens = usageAcc.LocalCompletionTokens() + usage.TotalTokens = usage.PromptTokens + usage.CompletionTokens } if !sentStart { @@ -529,7 +531,7 @@ func OaiResponsesToChatStreamHandler(c *gin.Context, info *relaycommon.RelayInfo info.ClaudeConvertInfo.Usage = usage } finishReason := "stop" - if sawToolCall && outputText.Len() == 0 { + if sawToolCall && outputLen == 0 { finishReason = "tool_calls" } stop := helper.GenerateStopResponse(responseId, createAt, model, finishReason) diff --git a/relay/channel/openai/helper.go b/relay/channel/openai/helper.go index 1a01d06da6dc..a807dad299c3 100644 --- a/relay/channel/openai/helper.go +++ b/relay/channel/openai/helper.go @@ -1,7 +1,7 @@ package openai import ( - "strings" + "io" "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/dto" @@ -74,7 +74,7 @@ func handleGeminiFormat(c *gin.Context, data string, info *relaycommon.RelayInfo return nil } -func ProcessStreamResponse(streamResponse dto.ChatCompletionsStreamResponse, responseTextBuilder *strings.Builder, toolCount *int) error { +func ProcessStreamResponse(streamResponse dto.ChatCompletionsStreamResponse, responseTextBuilder io.StringWriter, toolCount *int) error { for _, choice := range streamResponse.Choices { responseTextBuilder.WriteString(choice.Delta.GetContentString()) responseTextBuilder.WriteString(choice.Delta.GetReasoningContent()) @@ -91,7 +91,7 @@ func ProcessStreamResponse(streamResponse dto.ChatCompletionsStreamResponse, res return nil } -func processTokenData(relayMode int, data string, responseTextBuilder *strings.Builder, toolCount *int) error { +func processTokenData(relayMode int, data string, responseTextBuilder io.StringWriter, toolCount *int) error { switch relayMode { case relayconstant.RelayModeChatCompletions: var streamResponse dto.ChatCompletionsStreamResponse @@ -109,7 +109,7 @@ func processTokenData(relayMode int, data string, responseTextBuilder *strings.B return nil } -func processCompletionsStreamResponse(streamResponse dto.CompletionsStreamResponse, responseTextBuilder *strings.Builder) { +func processCompletionsStreamResponse(streamResponse dto.CompletionsStreamResponse, responseTextBuilder io.StringWriter) { for _, choice := range streamResponse.Choices { responseTextBuilder.WriteString(choice.Text) } diff --git a/relay/channel/openai/relay-openai.go b/relay/channel/openai/relay-openai.go index de40fe7071fc..b5b9fbef7c31 100644 --- a/relay/channel/openai/relay-openai.go +++ b/relay/channel/openai/relay-openai.go @@ -113,7 +113,7 @@ func OaiStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Re var createAt int64 = 0 var systemFingerprint string var containStreamUsage bool - var responseTextBuilder strings.Builder + usageAcc := service.NewUsageAccumulator(info.UpstreamModelName) var toolCount int var usage = &dto.Usage{} var lastStreamData string @@ -136,7 +136,7 @@ func OaiStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Re } lastStreamData = data - if err := processTokenData(info.RelayMode, data, &responseTextBuilder, &toolCount); err != nil { + if err := processTokenData(info.RelayMode, data, usageAcc, &toolCount); err != nil { logger.LogError(c, "error processing stream token data: "+err.Error()) sr.Error(err) } @@ -175,8 +175,12 @@ func OaiStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Re } if !containStreamUsage { - usage = service.ResponseText2Usage(c, responseTextBuilder.String(), info.UpstreamModelName, info.GetEstimatePromptTokens()) - usage.CompletionTokens += toolCount * 7 + // 上游未提供 usage：用流式估算的本地 token 数（不缓冲全文）+ 工具调用补偿 + if usage.PromptTokens == 0 { + usage.PromptTokens = info.GetEstimatePromptTokens() + } + usage.CompletionTokens = usageAcc.LocalCompletionTokens() + toolCount*7 + usage.TotalTokens = usage.PromptTokens + usage.CompletionTokens } applyUsagePostProcessing(info, usage, common.StringToByteSlice(lastStreamData)) diff --git a/relay/channel/openai/relay_responses.go b/relay/channel/openai/relay_responses.go index 2665b8d027e9..839b8f1ed44f 100644 --- a/relay/channel/openai/relay_responses.go +++ b/relay/channel/openai/relay_responses.go @@ -4,7 +4,6 @@ import ( "fmt" "io" "net/http" - "strings" "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/dto" @@ -77,7 +76,19 @@ func OaiResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp defer service.CloseResponseBodyGracefully(resp) var usage = &dto.Usage{} - var responseTextBuilder strings.Builder + + // Always stream the completion text through a bounded-memory, exact token + // counter instead of buffering the full response in a strings.Builder. This + // keeps heap residency flat (~tens of bytes) for large-context streaming + // responses regardless of the trust setting. TrustUpstreamUsage only decides + // whether the upstream-reported usage takes precedence over the locally + // counted tokens. + trustUpstreamUsage := info != nil && info.ChannelSetting.TrustUpstreamUsage + modelName := "" + if info != nil { + modelName = info.UpstreamModelName + } + usageAcc := service.NewUsageAccumulator(modelName) helper.StreamScannerHandler(c, resp, info, func(data string, sr *helper.StreamResult) { @@ -113,8 +124,8 @@ func OaiResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp } } case "response.output_text.delta": - // 处理输出文本 - responseTextBuilder.WriteString(streamResponse.Delta) + // 处理输出文本：流式估算，不缓冲全文 + usageAcc.Feed(streamResponse.Delta) case dto.ResponsesOutputTypeItemDone: // 函数调用处理 if streamResponse.Item != nil { @@ -130,17 +141,12 @@ func OaiResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp } }) - if usage.CompletionTokens == 0 { - // 计算输出文本的 token 数量 - tempStr := responseTextBuilder.String() - if len(tempStr) > 0 { - // 非正常结束，使用输出文本的 token 数量 - completionTokens := service.CountTextToken(tempStr, info.UpstreamModelName) - usage.CompletionTokens = completionTokens - } - } + // Resolve completion tokens via the unified accumulator: + // - trust=true: prefer upstream usage, fall back to local estimate if absent + // - trust=false: use the local streamed estimate (bounded memory, no buffering) + usage.CompletionTokens = usageAcc.Resolve(usage.CompletionTokens, trustUpstreamUsage) - if usage.PromptTokens == 0 && usage.CompletionTokens != 0 { + if info != nil && usage.PromptTokens == 0 && usage.CompletionTokens != 0 { usage.PromptTokens = info.GetEstimatePromptTokens() } diff --git a/relay/channel/openai/stream_handler_test.go b/relay/channel/openai/stream_handler_test.go new file mode 100644 index 000000000000..ff04698fa484 --- /dev/null +++ b/relay/channel/openai/stream_handler_test.go @@ -0,0 +1,123 @@ +package openai + +import ( + "io" + "net/http" + "net/http/httptest" + "os" + "strings" + "testing" + + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/dto" + relaycommon "github.com/QuantumNous/new-api/relay/common" + relayconstant "github.com/QuantumNous/new-api/relay/constant" + "github.com/QuantumNous/new-api/types" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +func TestMain(m *testing.M) { + if constant.StreamingTimeout <= 0 { + constant.StreamingTimeout = 300 + } + os.Exit(m.Run()) +} + +func streamInfo(model string) *relaycommon.RelayInfo { + info := &relaycommon.RelayInfo{ + RelayFormat: types.RelayFormatOpenAI, + RelayMode: relayconstant.RelayModeChatCompletions, + ChannelMeta: &relaycommon.ChannelMeta{ + UpstreamModelName: model, + ChannelSetting: dto.ChannelSettings{}, + }, + } + info.SetEstimatePromptTokens(100) + return info +} + +func sseResp(s string) *http.Response { + return &http.Response{Body: io.NopCloser(strings.NewReader(s)), StatusCode: http.StatusOK} +} + +func streamCtx(path string) *gin.Context { + gin.SetMode(gin.TestMode) + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + c.Request = httptest.NewRequest(http.MethodPost, path, nil) + return c +} + +// chat/completions 上游带 stream usage：采用上游。 +func TestOaiStreamHandler_UpstreamUsage(t *testing.T) { + sse := `data: {"id":"c","choices":[{"delta":{"role":"assistant","content":"Hello world answer text"}}]} +data: {"id":"c","choices":[{"delta":{},"finish_reason":"stop"}],"usage":{"prompt_tokens":10,"completion_tokens":8,"total_tokens":18}} +data: [DONE] +` + usage, apiErr := OaiStreamHandler(streamCtx("/v1/chat/completions"), streamInfo("gpt-4o"), sseResp(sse)) + require.Nil(t, apiErr) + require.NotNil(t, usage) + require.Equal(t, 10, usage.PromptTokens) + require.Equal(t, 8, usage.CompletionTokens) +} + +// chat/completions 上游无 usage：本地估算 > 0。 +func TestOaiStreamHandler_LocalFallback(t *testing.T) { + sse := `data: {"id":"c","choices":[{"delta":{"role":"assistant","content":"Hello world this is a long generated answer text without usage"}}]} +data: {"id":"c","choices":[{"delta":{},"finish_reason":"stop"}]} +data: [DONE] +` + usage, apiErr := OaiStreamHandler(streamCtx("/v1/chat/completions"), streamInfo("gpt-4o"), sseResp(sse)) + require.Nil(t, apiErr) + require.NotNil(t, usage) + require.Greater(t, usage.CompletionTokens, 0) +} + +// chat_via_responses 上游 responses 流：解析并计 usage。 +func TestOaiResponsesToChatStreamHandler_Basic(t *testing.T) { + sse := `data: {"type":"response.created","response":{"id":"r","status":"in_progress"}} +data: {"type":"response.output_text.delta","delta":"Hello world this is the answer"} +data: {"type":"response.completed","response":{"id":"r","status":"completed","usage":{"input_tokens":10,"output_tokens":6,"total_tokens":16}}} +` + usage, apiErr := OaiResponsesToChatStreamHandler(streamCtx("/v1/chat/completions"), streamInfo("gpt-5.5"), sseResp(sse)) + require.Nil(t, apiErr) + require.NotNil(t, usage) + require.Greater(t, usage.CompletionTokens, 0) +} + +// tool call（无上游 usage）：completion = 本地文本估算 + toolCount*7。 +func TestOaiStreamHandler_ToolCountCompensation(t *testing.T) { + // 无文本内容、只有一个 tool call：completion 应至少包含 toolCount*7=7 + sse := `data: {"id":"c","choices":[{"delta":{"role":"assistant","tool_calls":[{"index":0,"id":"call_1","type":"function","function":{"name":"get_weather","arguments":"{}"}}]}}]} +data: {"id":"c","choices":[{"delta":{},"finish_reason":"tool_calls"}]} +data: [DONE] +` + usage, apiErr := OaiStreamHandler(streamCtx("/v1/chat/completions"), streamInfo("gpt-4o"), sseResp(sse)) + require.Nil(t, apiErr) + require.NotNil(t, usage) + require.GreaterOrEqual(t, usage.CompletionTokens, 7, "至少含 toolCount*7 补偿") +} + +// 多 choice + 空 delta：不 panic，正常累计。 +func TestOaiStreamHandler_MultiChoiceAndEmptyDelta(t *testing.T) { + sse := `data: {"id":"c","choices":[{"index":0,"delta":{"role":"assistant","content":"Hello"}},{"index":1,"delta":{"content":"World"}}]} +data: {"id":"c","choices":[{"index":0,"delta":{}}]} +data: {"id":"c","choices":[{"index":0,"delta":{},"finish_reason":"stop"}],"usage":{"prompt_tokens":5,"completion_tokens":3,"total_tokens":8}} +data: [DONE] +` + usage, apiErr := OaiStreamHandler(streamCtx("/v1/chat/completions"), streamInfo("gpt-4o"), sseResp(sse)) + require.Nil(t, apiErr) + require.NotNil(t, usage) + require.Equal(t, 3, usage.CompletionTokens, "上游 usage 优先") +} + +// 流中途断（无结束帧、无 usage）：不 panic，本地估算已收到的文本。 +func TestOaiStreamHandler_TruncatedStream(t *testing.T) { + sse := `data: {"id":"c","choices":[{"delta":{"role":"assistant","content":"partial answer before"}}]} +data: {"id":"c","choices":[{"delta":{"content":" the stream was cut"}}]} +` + usage, apiErr := OaiStreamHandler(streamCtx("/v1/chat/completions"), streamInfo("gpt-4o"), sseResp(sse)) + require.Nil(t, apiErr) + require.NotNil(t, usage) + require.Greater(t, usage.CompletionTokens, 0, "截断流应本地估算已收文本") +} diff --git a/relay/channel/tencent/relay-tencent.go b/relay/channel/tencent/relay-tencent.go index 4cda7541a6f4..c64daddd5cb0 100644 --- a/relay/channel/tencent/relay-tencent.go +++ b/relay/channel/tencent/relay-tencent.go @@ -91,7 +91,7 @@ func streamResponseTencent2OpenAI(TencentResponse *TencentChatResponse) *dto.Cha } func tencentStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.NewAPIError) { - var responseText string + usageAcc := service.NewUsageAccumulator(info.UpstreamModelName) scanner := helper.NewStreamScanner(resp.Body) scanner.Split(bufio.ScanLines) @@ -113,7 +113,7 @@ func tencentStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *htt response := streamResponseTencent2OpenAI(&tencentResponse) if len(response.Choices) != 0 { - responseText += response.Choices[0].Delta.GetContentString() + usageAcc.Feed(response.Choices[0].Delta.GetContentString()) } err = helper.ObjectData(c, response) @@ -130,7 +130,12 @@ func tencentStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *htt service.CloseResponseBodyGracefully(resp) - return service.ResponseText2Usage(c, responseText, info.UpstreamModelName, info.GetEstimatePromptTokens()), nil + usage := &dto.Usage{ + PromptTokens: info.GetEstimatePromptTokens(), + CompletionTokens: usageAcc.LocalCompletionTokens(), + } + usage.TotalTokens = usage.PromptTokens + usage.CompletionTokens + return usage, nil } func tencentHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.NewAPIError) { diff --git a/relay/channel/tencent/stream_handler_test.go b/relay/channel/tencent/stream_handler_test.go new file mode 100644 index 000000000000..9d64e57054fc --- /dev/null +++ b/relay/channel/tencent/stream_handler_test.go @@ -0,0 +1,62 @@ +package tencent + +import ( + "io" + "net/http" + "net/http/httptest" + "os" + "strings" + "testing" + "time" + + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/dto" + relaycommon "github.com/QuantumNous/new-api/relay/common" + "github.com/QuantumNous/new-api/types" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +func TestMain(m *testing.M) { + if constant.StreamingTimeout <= 0 { + constant.StreamingTimeout = 300 + } + os.Exit(m.Run()) +} + +func streamInfo(model string) *relaycommon.RelayInfo { + info := &relaycommon.RelayInfo{ + RelayFormat: types.RelayFormatOpenAI, + StartTime: time.Now(), + ChannelMeta: &relaycommon.ChannelMeta{ + UpstreamModelName: model, + ChannelSetting: dto.ChannelSettings{}, + }, + } + info.SetEstimatePromptTokens(100) + return info +} + +func sseResp(s string) *http.Response { + return &http.Response{Body: io.NopCloser(strings.NewReader(s)), StatusCode: http.StatusOK} +} + +func streamCtx() *gin.Context { + gin.SetMode(gin.TestMode) + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil) + return c +} + +// tencent 流式：累积文本，本地估算 completion > 0。 +func TestTencentStreamHandler_LocalEstimate(t *testing.T) { + sse := `data: {"Choices":[{"Delta":{"Role":"assistant","Content":"Hello world this is"},"FinishReason":""}],"Id":"x"} +data: {"Choices":[{"Delta":{"Content":" a tencent answer"},"FinishReason":""}],"Id":"x"} +data: {"Choices":[{"Delta":{"Content":""},"FinishReason":"stop"}],"Id":"x"} +` + usage, apiErr := tencentStreamHandler(streamCtx(), streamInfo("hunyuan-standard"), sseResp(sse)) + require.Nil(t, apiErr) + require.NotNil(t, usage) + require.Greater(t, usage.CompletionTokens, 0) + require.Equal(t, 100, usage.PromptTokens) +} diff --git a/relay/channel/xai/stream_handler_test.go b/relay/channel/xai/stream_handler_test.go new file mode 100644 index 000000000000..6b397c669c8b --- /dev/null +++ b/relay/channel/xai/stream_handler_test.go @@ -0,0 +1,71 @@ +package xai + +import ( + "io" + "net/http" + "net/http/httptest" + "os" + "strings" + "testing" + + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/dto" + relaycommon "github.com/QuantumNous/new-api/relay/common" + "github.com/QuantumNous/new-api/types" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +func TestMain(m *testing.M) { + if constant.StreamingTimeout <= 0 { + constant.StreamingTimeout = 300 + } + os.Exit(m.Run()) +} + +func streamInfo(model string) *relaycommon.RelayInfo { + info := &relaycommon.RelayInfo{ + RelayFormat: types.RelayFormatOpenAI, + ChannelMeta: &relaycommon.ChannelMeta{ + UpstreamModelName: model, + ChannelSetting: dto.ChannelSettings{}, + }, + } + info.SetEstimatePromptTokens(100) + return info +} + +func sseResp(s string) *http.Response { + return &http.Response{Body: io.NopCloser(strings.NewReader(s)), StatusCode: http.StatusOK} +} + +func streamCtx() *gin.Context { + gin.SetMode(gin.TestMode) + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil) + return c +} + +// 上游带 usage：直接采用。 +func TestXAIStreamHandler_UpstreamUsage(t *testing.T) { + sse := `data: {"id":"x","choices":[{"delta":{"role":"assistant","content":"Hello world answer"}}]} +data: {"id":"x","choices":[{"delta":{"content":" more"}}],"usage":{"prompt_tokens":10,"completion_tokens":7,"total_tokens":17}} +data: [DONE] +` + usage, apiErr := xAIStreamHandler(streamCtx(), streamInfo("grok-2"), sseResp(sse)) + require.Nil(t, apiErr) + require.NotNil(t, usage) + require.Equal(t, 10, usage.PromptTokens) + require.Equal(t, 7, usage.CompletionTokens) +} + +// 上游无 usage：本地估算 > 0。 +func TestXAIStreamHandler_LocalFallback(t *testing.T) { + sse := `data: {"id":"x","choices":[{"delta":{"role":"assistant","content":"Hello world this is a fairly long generated answer text"}}]} +data: [DONE] +` + usage, apiErr := xAIStreamHandler(streamCtx(), streamInfo("grok-2"), sseResp(sse)) + require.Nil(t, apiErr) + require.NotNil(t, usage) + require.Greater(t, usage.CompletionTokens, 0) +} diff --git a/relay/channel/xai/text.go b/relay/channel/xai/text.go index f9a8ee2e6f96..94aaaec36106 100644 --- a/relay/channel/xai/text.go +++ b/relay/channel/xai/text.go @@ -3,7 +3,6 @@ package xai import ( "io" "net/http" - "strings" "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/dto" @@ -37,7 +36,7 @@ func streamResponseXAI2OpenAI(xAIResp *dto.ChatCompletionsStreamResponse, usage func xAIStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.NewAPIError) { usage := &dto.Usage{} - var responseTextBuilder strings.Builder + usageAcc := service.NewUsageAccumulator(info.UpstreamModelName) var toolCount int var containStreamUsage bool @@ -60,7 +59,7 @@ func xAIStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Re } openaiResponse := streamResponseXAI2OpenAI(xAIResp, usage) - _ = openai.ProcessStreamResponse(*openaiResponse, &responseTextBuilder, &toolCount) + _ = openai.ProcessStreamResponse(*openaiResponse, usageAcc, &toolCount) if err := helper.ObjectData(c, openaiResponse); err != nil { common.SysLog(err.Error()) sr.Error(err) @@ -68,8 +67,9 @@ func xAIStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Re }) if !containStreamUsage { - usage = service.ResponseText2Usage(c, responseTextBuilder.String(), info.UpstreamModelName, info.GetEstimatePromptTokens()) - usage.CompletionTokens += toolCount * 7 + usage.PromptTokens = info.GetEstimatePromptTokens() + usage.CompletionTokens = usageAcc.LocalCompletionTokens() + toolCount*7 + usage.TotalTokens = usage.PromptTokens + usage.CompletionTokens } helper.Done(c) diff --git a/service/estimator_reference_test.go b/service/estimator_reference_test.go new file mode 100644 index 000000000000..7dd775403e4c --- /dev/null +++ b/service/estimator_reference_test.go @@ -0,0 +1,169 @@ +package service + +import ( + "math" + "math/rand" + "strings" + "testing" + "unicode" +) + +// referenceEstimateToken 是【重构前】EstimateToken 的逐字符循环实现的精确副本 +// （取自 merge 提交 4f099c54 的 service/token_estimator.go）。 +// 它作为独立的参考实现，与重构后基于 streamingEstimator 的 EstimateToken 对拍， +// 用来证明：(1) 重构没有改变计费口径；(2) 测试不是自我循环验证 +// （若把 streamingEstimator.feed 改坏，本对拍会失败，因为参考实现不走 feed）。 +func referenceEstimateToken(provider Provider, text string) int { + m := getMultipliers(provider) + var count float64 + + type WordType int + const ( + None WordType = iota + Latin + Number + ) + currentWordType := None + + for _, r := range text { + if unicode.IsSpace(r) { + currentWordType = None + if r == '\n' || r == '\t' { + count += m.Newline + } else { + count += m.Space + } + continue + } + if isCJK(r) { + currentWordType = None + count += m.CJK + continue + } + if isEmoji(r) { + currentWordType = None + count += m.Emoji + continue + } + if isLatinOrNumber(r) { + isNum := unicode.IsNumber(r) + newType := Latin + if isNum { + newType = Number + } + if currentWordType == None || currentWordType != newType { + if newType == Number { + count += m.Number + } else { + count += m.Word + } + currentWordType = newType + } + continue + } + currentWordType = None + if isMathSymbol(r) { + count += m.MathSymbol + } else if r == '@' { + count += m.AtSign + } else if isURLDelim(r) { + count += m.URLDelim + } else { + count += m.Symbol + } + } + return int(math.Ceil(count)) + m.BasePad +} + +// 重构后的 EstimateToken 必须与独立的参考实现逐位相同（证明计费口径不变）。 +func TestEstimateToken_MatchesReferenceImpl(t *testing.T) { + corpus := []string{ + "hello", "hello world", "abc 123", "中文测试混合 English", + "a", " ", "Hello, World!\nNew line.\ttab", + "emoji 😀🎉 symbols ©® math ∑∫ url https://x.com/a?b=c user@host", + "VeryLongWordNoSpace", "123 456 789", "Mixed123Letters456", + strings.Repeat("word ", 100), strings.Repeat("中", 50), "", + } + for _, p := range []Provider{OpenAI, Gemini, Claude} { + for _, text := range corpus { + want := referenceEstimateToken(p, text) + got := EstimateToken(p, text) + if got != want { + t.Errorf("provider=%s text=%q: refactored=%d reference=%d", p, text, got, want) + } + } + } +} + +// 随机大样本对拍：确保重构在任意输入上都等于参考实现。 +func TestEstimateToken_FuzzMatchesReference(t *testing.T) { + alphabet := []rune(" \n\tabcXYZ012中文😀{}()@∑/.,!?") + r := rand.New(rand.NewSource(2026)) + for _, p := range []Provider{OpenAI, Gemini, Claude} { + for trial := 0; trial < 1000; trial++ { + n := r.Intn(300) + var b strings.Builder + for i := 0; i < n; i++ { + b.WriteRune(alphabet[r.Intn(len(alphabet))]) + } + text := b.String() + want := referenceEstimateToken(p, text) + got := EstimateToken(p, text) + if got != want { + t.Fatalf("provider=%s text=%q: refactored=%d reference=%d", p, text, got, want) + } + } + } +} + +// 尾部被切断的不完整 UTF-8 字节序列：一次性 EstimateToken 与流式 feed 都必须 +// 与参考实现（for range，对每个残留字节 emit utf8.RuneError）逐位相同。 +// 这覆盖一个真实场景：上游分块流式输出把一个多字节 rune 切在 chunk 末尾， +// 且流在补齐前就结束（截断/超时）——此时残留字节必须按 RuneError 计入， +// 否则流式计数会比旧实现少算，破坏计费口径一致性。 +func TestEstimateToken_TruncatedUTF8MatchesReference(t *testing.T) { + // 构造若干以不完整 UTF-8 结尾的输入（“中”=e4 b8 ad，“😀”=f0 9f 98 80）。 + truncated := []string{ + "ab" + string([]byte{0xe4, 0xb8}), // 2 个残留字节 + "hello " + string([]byte{0xe4}), // 1 个残留字节 + "x" + string([]byte{0xf0, 0x9f, 0x98}), // 3 个残留字节（emoji 切 3/4） + string([]byte{0xe4, 0xb8}), // 全是残留字节 + "中文" + string([]byte{0xf0, 0x9f}), // 完整 CJK + 2 残留字节 + } + for _, p := range []Provider{OpenAI, Gemini, Claude} { + for _, text := range truncated { + want := referenceEstimateToken(p, text) + // 一次性 + if got := EstimateToken(p, text); got != want { + t.Errorf("[oneshot] provider=%s bytes=% x: got=%d want=%d", p, []byte(text), got, want) + } + // 流式：把残留尾字节单独作为最后一个 chunk 喂入，再 result + e := newStreamingEstimator(p) + e.feed(text) + if got := e.result(); got != want { + t.Errorf("[stream] provider=%s bytes=% x: got=%d want=%d", p, []byte(text), got, want) + } + } + } +} + +// 绝对锚点：人工核算的已知值（不依赖任何被测实现，防止参考实现也被一起改坏）。 +// OpenAI: Word=1.02 Space=0.42 Number=1.55 Newline=0.5；BasePad=0；结果=ceil(sum)。 +func TestEstimateToken_HardcodedAnchors(t *testing.T) { + cases := []struct { + provider Provider + text string + want int + }{ + {OpenAI, "hello", 2}, // 1 word: ceil(1.02)=2 + {OpenAI, "hello world", 3}, // 1.02+0.42+1.02=2.46 -> 3 + {OpenAI, "", 0}, // 空 + {OpenAI, "a b", 3}, // 1.02+0.42+1.02=2.46 -> 3 + } + for _, c := range cases { + got := EstimateToken(c.provider, c.text) + if got != c.want { + t.Errorf("provider=%s text=%q: got=%d want=%d (hardcoded anchor)", c.provider, c.text, got, c.want) + } + } +} diff --git a/service/stream_token_counter.go b/service/stream_token_counter.go new file mode 100644 index 000000000000..a4eb7e0fd27f --- /dev/null +++ b/service/stream_token_counter.go @@ -0,0 +1,109 @@ +package service + +import "strings" + +// providerForModel maps a model name to the estimator Provider, matching the +// exact dispatch logic of EstimateTokenByModel so that streaming estimation +// produces identical results to the legacy ResponseText2Usage path. +func providerForModel(model string) Provider { + m := strings.ToLower(model) + switch { + case strings.Contains(m, "gemini"): + return Gemini + case strings.Contains(m, "claude"): + return Claude + default: + return OpenAI + } +} + +// UsageAccumulator replaces the legacy "accumulate the whole response into a +// strings.Builder, then EstimateToken(builder.String())" pattern used across +// every streaming relay handler. +// +// Instead of buffering the full response text (which makes heap residency grow +// with response size — hundreds of MB for large-context streams, never released +// under default Go GC), it feeds each delta through a streaming estimator that +// keeps only O(1) state. The token result is bit-for-bit identical to running +// EstimateTokenByModel over the concatenated text, because the streaming +// estimator shares the exact same per-rune state machine as EstimateToken. +// +// Billing semantics are controlled by TrustUpstreamUsage at the call site via +// Resolve(): +// - trust=true: prefer the upstream-reported completion tokens; fall back to +// the local streamed estimate only when the upstream omits usage. +// - trust=false: use the local streamed estimate (the legacy behavior, but +// without buffering the full text). +// +// For models/channels that emit separate reasoning ("thinking") content (e.g. +// Claude extended thinking), feed it via FeedReasoning so it is counted with a +// dedicated estimator and summed into the completion tokens — matching the +// legacy behavior of writing both text and thinking into the same builder, but +// without holding either in memory. +type UsageAccumulator struct { + text *streamingEstimator + reasoning *streamingEstimator // lazily created on first FeedReasoning + provider Provider +} + +// NewUsageAccumulator builds an accumulator for the given model. +func NewUsageAccumulator(model string) *UsageAccumulator { + p := providerForModel(model) + return &UsageAccumulator{ + text: newStreamingEstimator(p), + provider: p, + } +} + +// Feed accumulates a chunk of output text. The chunk may be any byte stream, +// including a multibyte rune split across two calls. +func (a *UsageAccumulator) Feed(delta string) { + if delta == "" { + return + } + a.text.feed(delta) +} + +// FeedReasoning accumulates a chunk of reasoning/thinking text, counted with a +// separate estimator so reasoning and visible text are tallied independently +// and then summed. +func (a *UsageAccumulator) FeedReasoning(delta string) { + if delta == "" { + return + } + if a.reasoning == nil { + a.reasoning = newStreamingEstimator(a.provider) + } + a.reasoning.feed(delta) +} + +// WriteString implements io.StringWriter so a *UsageAccumulator can be passed +// where the legacy code expected a *strings.Builder (e.g. openai +// ProcessStreamResponse / processTokenData). It accumulates into the text +// estimator and always reports the full input as consumed. +func (a *UsageAccumulator) WriteString(s string) (int, error) { + a.Feed(s) + return len(s), nil +} + +// LocalCompletionTokens returns the locally estimated completion tokens +// (text + reasoning), equal to EstimateTokenByModel over the concatenated +// streamed text. +func (a *UsageAccumulator) LocalCompletionTokens() int { + n := a.text.result() + if a.reasoning != nil { + n += a.reasoning.result() + } + return n +} + +// Resolve returns the final completion token count given the upstream-reported +// value and the channel's trust setting. +// - trustUpstream && upstreamCompletion > 0 -> use the upstream value +// - otherwise -> use the local streamed estimate +func (a *UsageAccumulator) Resolve(upstreamCompletion int, trustUpstream bool) int { + if trustUpstream && upstreamCompletion > 0 { + return upstreamCompletion + } + return a.LocalCompletionTokens() +} diff --git a/service/stream_token_counter_test.go b/service/stream_token_counter_test.go new file mode 100644 index 000000000000..d1a8ba945f29 --- /dev/null +++ b/service/stream_token_counter_test.go @@ -0,0 +1,192 @@ +package service + +import ( + "math/rand" + "strings" + "testing" +) + +func randomChunks(text string, r *rand.Rand) []string { + runes := []rune(text) + var chunks []string + i := 0 + for i < len(runes) { + n := r.Intn(8) + 1 + if i+n > len(runes) { + n = len(runes) - i + } + chunks = append(chunks, string(runes[i:i+n])) + i += n + } + return chunks +} + +// streamingEstimator must be bit-for-bit identical to one-shot EstimateToken +// across ordinary chunk splits. +func TestStreamingEstimatorMatchesOneShot(t *testing.T) { + corpus := []string{ + "Hello, world! This is a test.", + "function calculate(a, b) { return a + b; }", + "中文混合 English text 123456789 测试 tokenization", + " leading and multiple spaces ", + "newlines\nand\ttabs\r\nmixed", + "emoji 😀🎉 and symbols ©®™ €£¥", + "VeryLongWordWithoutAnySpaces", + "a b c d e f g h i j k l m n o p", + "https://example.com/path?query=value&foo=bar#frag", + "```go\nfunc main() {\n\tfmt.Println(\"hi\")\n}\n```", + strings.Repeat("repeat ", 200), + strings.Repeat("无空格连续中文", 100), + "Mixed123Numbers456And789Letters", + "\n\n\n\n\n", " ", "", "single", "a", " ", + "user@example.com sends ∑∫∂√ math", + } + providers := []Provider{OpenAI, Gemini, Claude} + r := rand.New(rand.NewSource(1)) + for _, p := range providers { + for _, text := range corpus { + whole := EstimateToken(p, text) + for trial := 0; trial < 30; trial++ { + e := newStreamingEstimator(p) + for _, ch := range randomChunks(text, r) { + e.feed(ch) + } + if got := e.result(); got != whole { + t.Errorf("provider=%s text=%q stream=%d whole=%d", p, text, got, whole) + } + } + } + } +} + +func TestStreamingEstimatorMatchesOneShotWhenMultibyteRunesAreSplit(t *testing.T) { + text := "prefix 中文 😀 suffix" + splits := [][]int{ + {1}, + {7, 8, 9}, // Split the first CJK rune byte-by-byte. + {7, 10, 11, 12}, // Complete one CJK rune, then split the next. + {14, 15, 16, 17}, // Split the emoji byte-by-byte. + {7, 8, 10, 14, 16, len([]byte(text))}, + } + for _, p := range []Provider{OpenAI, Gemini, Claude} { + want := EstimateToken(p, text) + for _, split := range splits { + e := newStreamingEstimator(p) + start := 0 + data := []byte(text) + for _, end := range split { + e.feed(string(data[start:end])) + start = end + } + if start < len(data) { + e.feed(string(data[start:])) + } + if got := e.result(); got != want { + t.Errorf("provider=%s split=%v: stream=%d whole=%d", p, split, got, want) + } + } + } +} + +// Random fuzz including bytes that force multibyte rune splits. +func TestStreamingEstimatorFuzz(t *testing.T) { + alphabet := []rune(" \n\tabcXYZ012中文测试😀{}()<>/@∑") + r := rand.New(rand.NewSource(42)) + for _, p := range []Provider{OpenAI, Gemini, Claude} { + for trial := 0; trial < 500; trial++ { + n := r.Intn(1500) + var b strings.Builder + for i := 0; i < n; i++ { + b.WriteRune(alphabet[r.Intn(len(alphabet))]) + } + text := b.String() + whole := EstimateToken(p, text) + // split by raw BYTES (not runes) to force multibyte boundary splits + e := newStreamingEstimator(p) + data := []byte(text) + i := 0 + for i < len(data) { + step := r.Intn(5) + 1 + if i+step > len(data) { + step = len(data) - i + } + e.feed(string(data[i : i+step])) + i += step + } + if got := e.result(); got != whole { + t.Errorf("provider=%s byte-split mismatch stream=%d whole=%d text=%q", p, got, whole, text[:min(80, len(text))]) + } + } + } +} + +// UsageAccumulator local count must equal legacy EstimateTokenByModel (gold standard). +func TestUsageAccumulatorGoldStandard(t *testing.T) { + cases := []struct{ model, text string }{ + {"gpt-5.5", strings.Repeat("This is generated output. ", 500)}, + {"claude-opus-4", strings.Repeat("Claude response text 中文 ", 300)}, + {"gemini-3-pro", strings.Repeat("Gemini output 12345 ", 400)}, + {"gpt-4o", "short"}, + } + r := rand.New(rand.NewSource(7)) + for _, c := range cases { + legacy := EstimateTokenByModel(c.model, c.text) + acc := NewUsageAccumulator(c.model) + for _, ch := range randomChunks(c.text, r) { + acc.Feed(ch) + } + if got := acc.LocalCompletionTokens(); got != legacy { + t.Errorf("model=%s acc=%d legacy=%d", c.model, got, legacy) + } + } +} + +// FeedReasoning: text + thinking counted separately then summed. +func TestUsageAccumulatorReasoning(t *testing.T) { + model := "claude-opus-4" + text := strings.Repeat("visible answer text ", 100) + thinking := strings.Repeat("internal reasoning step ", 200) + // legacy claude path concatenated both into one builder + legacyConcat := EstimateTokenByModel(model, text+thinking) + legacySeparate := EstimateTokenByModel(model, text) + EstimateTokenByModel(model, thinking) + + acc := NewUsageAccumulator(model) + acc.Feed(text) + acc.FeedReasoning(thinking) + got := acc.LocalCompletionTokens() + + // We chose separate counting (more accurate); assert it equals separate sum. + if got != legacySeparate { + t.Errorf("reasoning separate: acc=%d expected=%d", got, legacySeparate) + } + t.Logf("separate=%d concat=%d (diff is expected, separate is by design)", got, legacyConcat) +} + +// Resolve semantics: trust on/off. +func TestUsageAccumulatorResolve(t *testing.T) { + acc := NewUsageAccumulator("gpt-4o") + acc.Feed(strings.Repeat("word ", 50)) + local := acc.LocalCompletionTokens() + if local <= 0 { + t.Fatal("local should be > 0") + } + // trust=true, upstream provided -> use upstream + if got := acc.Resolve(999, true); got != 999 { + t.Errorf("trust+upstream: got %d want 999", got) + } + // trust=true, upstream missing -> use local + if got := acc.Resolve(0, true); got != local { + t.Errorf("trust+no-upstream: got %d want %d", got, local) + } + // trust=false -> always local (ignore upstream) + if got := acc.Resolve(999, false); got != local { + t.Errorf("no-trust: got %d want %d", got, local) + } +} + +func min(a, b int) int { + if a < b { + return a + } + return b +} diff --git a/service/token_estimator.go b/service/token_estimator.go index 9e27269ce3d8..453643993e74 100644 --- a/service/token_estimator.go +++ b/service/token_estimator.go @@ -5,6 +5,7 @@ import ( "strings" "sync" "unicode" + "unicode/utf8" ) // Provider 定义模型厂商大类 @@ -66,85 +67,131 @@ func getMultipliers(p Provider) multipliers { } // EstimateToken 计算 Token 数量 -func EstimateToken(provider Provider, text string) int { - m := getMultipliers(provider) - var count float64 - - // 状态机变量 - type WordType int - const ( - None WordType = iota - Latin - Number - ) - currentWordType := None - - for _, r := range text { - // 1. 处理空格和换行符 - if unicode.IsSpace(r) { - currentWordType = None - // 换行符和制表符使用Newline权重 - if r == '\n' || r == '\t' { - count += m.Newline - } else { - // 普通空格使用Space权重 - count += m.Space - } - continue - } +// wordType 是估算状态机里"当前是否处于一个连续单词/数字中"的状态。 +type wordType int - // 2. 处理 CJK (中日韩) - 按字符计费 - if isCJK(r) { - currentWordType = None - count += m.CJK - continue - } +const ( + wordNone wordType = iota + wordLatin + wordNumber +) - // 3. 处理Emoji - 使用专门的Emoji权重 - if isEmoji(r) { - currentWordType = None - count += m.Emoji - continue - } +// streamingEstimator 是 EstimateToken 的流式版本：逐 rune 喂入，维护与 +// EstimateToken 完全相同的状态机（count + currentWordType）。对任意切分方式， +// 分块喂入的结果与一次性整体估算【逐位相同】（BasePad 仅在 result 时加一次， +// 当前三个 provider 的 BasePad 均为 0）。内存为 O(1)（只有 count/状态/最多 +// 数字节的不完整 UTF-8 rune 缓冲），用于替代"用 strings.Builder 累积整个 +// 响应文本再 EstimateToken"的旧模式。 +type streamingEstimator struct { + m multipliers + count float64 + currentWordType wordType + pending []byte // 缓冲跨 chunk 边界被切断的不完整 UTF-8 rune +} - // 4. 处理拉丁字母/数字 (英文单词) - if isLatinOrNumber(r) { - isNum := unicode.IsNumber(r) - newType := Latin - if isNum { - newType = Number - } +func newStreamingEstimator(provider Provider) *streamingEstimator { + return &streamingEstimator{m: getMultipliers(provider)} +} - // 如果之前不在单词中，或者类型发生变化（字母<->数字），则视为新token - // 注意：对于OpenAI，通常"version 3.5"会切分，"abc123xyz"有时也会切分 - // 这里简单起见，字母和数字切换时增加权重 - if currentWordType == None || currentWordType != newType { - if newType == Number { - count += m.Number - } else { - count += m.Word - } - currentWordType = newType - } - // 单词中间的字符不额外计费 - continue +// feed 喂入一段文本（可以是任意 chunk，包括把一个多字节 rune 切成两半）。 +func (e *streamingEstimator) feed(s string) { + if s == "" { + return + } + var b []byte + if len(e.pending) > 0 { + b = append(e.pending, s...) + e.pending = nil + } else { + b = []byte(s) + } + for i := 0; i < len(b); { + if !utf8.FullRune(b[i:]) { + // 不完整的 UTF-8 rune 被切在 chunk 末尾，缓冲等待下次 + e.pending = append(e.pending[:0], b[i:]...) + return } + r, size := utf8.DecodeRune(b[i:]) + e.estimateRune(r) + i += size + } +} - // 5. 处理标点符号/特殊字符 - 按类型使用不同权重 - currentWordType = None - if isMathSymbol(r) { - count += m.MathSymbol - } else if r == '@' { - count += m.AtSign - } else if isURLDelim(r) { - count += m.URLDelim +// estimateRune 是从 EstimateToken 抽出的【单 rune】计费逻辑，二者共用，保证一致。 +func (e *streamingEstimator) estimateRune(r rune) { + m := e.m + // 1. 空格/换行 + if unicode.IsSpace(r) { + e.currentWordType = wordNone + if r == '\n' || r == '\t' { + e.count += m.Newline } else { - count += m.Symbol + e.count += m.Space } + return + } + // 2. CJK + if isCJK(r) { + e.currentWordType = wordNone + e.count += m.CJK + return } + // 3. Emoji + if isEmoji(r) { + e.currentWordType = wordNone + e.count += m.Emoji + return + } + // 4. 拉丁字母/数字（连续单词） + if isLatinOrNumber(r) { + newType := wordLatin + if unicode.IsNumber(r) { + newType = wordNumber + } + if e.currentWordType == wordNone || e.currentWordType != newType { + if newType == wordNumber { + e.count += m.Number + } else { + e.count += m.Word + } + e.currentWordType = newType + } + return + } + // 5. 标点/特殊字符 + e.currentWordType = wordNone + if isMathSymbol(r) { + e.count += m.MathSymbol + } else if r == '@' { + e.count += m.AtSign + } else if isURLDelim(r) { + e.count += m.URLDelim + } else { + e.count += m.Symbol + } +} - // 向上取整并加上基础 padding - return int(math.Ceil(count)) + m.BasePad +// result 返回当前累计的 token 估算（向上取整 + BasePad）。 +// 若仍有缓冲的不完整 UTF-8 尾字节（流在补齐前结束），按与一次性 +// EstimateToken（for range over string）一致的语义处理：每个残留字节 +// 解码为 utf8.RuneError，逐字节计入，保证流式与一次性结果逐位相同。 +func (e *streamingEstimator) result() int { + if len(e.pending) > 0 { + for range e.pending { + e.estimateRune(utf8.RuneError) + } + e.pending = nil + } + return int(math.Ceil(e.count)) + e.m.BasePad +} + +func EstimateToken(provider Provider, text string) int { + if text == "" { + return 0 + } + e := newStreamingEstimator(provider) + e.feed(text) + return e.result() } // 辅助：判断是否为 CJK 字符 diff --git a/web/default/src/features/channels/components/drawers/channel-mutate-drawer.tsx b/web/default/src/features/channels/components/drawers/channel-mutate-drawer.tsx index 87dcc84afea6..18f0bd36516b 100644 --- a/web/default/src/features/channels/components/drawers/channel-mutate-drawer.tsx +++ b/web/default/src/features/channels/components/drawers/channel-mutate-drawer.tsx @@ -3192,6 +3192,29 @@ export function ChannelMutateDrawer({ )} /> + + ( + +
+ {t('Trust Upstream Usage')} + + {t( + 'Use the usage reported by the upstream for billing instead of buffering the full response to recount tokens. Enable for channels whose upstream returns accurate usage (e.g. official OpenAI Responses) to keep memory flat on large streaming responses.' + )} + +
+ + + + + )} + />