mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-08-04 23:00:30 +08:00
fix(go-models): record Novita streaming usage without chatConfig (#17778)
Relate to #17284. ## Problem `novitaHandleStream` guarded usage recording with `if found && chatConfig != nil`. When a caller passes a nil `*ChatConfig` — common in the service layer (`model_chat.go`, `chat_pipeline.go`) — the streamed token usage is dropped entirely. The shared `HandleStreamingResponse` only uses `chatConfig` to expose `UsageResult` and records usage whenever the stream carries it. Novita's bespoke handler diverged from every other OpenAI-compatible streaming driver. ## Fix Record usage whenever the stream carries a usage event, mirroring `HandleStreamingResponse`. `applyStreamUsage` already handles a nil `chatConfig` internally (it only writes `chatConfig.UsageResult` when non-nil), so the extra guard was doing nothing but dropping usage. ## Test `TestNovitaStreamRecordsUsageWithoutChatConfig`: - nil `chatConfig` + usage event → stream completes without error (guard removed safely) - non-nil `chatConfig` + usage event → `UsageResult` populated with the streamed tokens Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -221,7 +221,7 @@ func novitaHandleStream(
|
||||
|
||||
done, err := ParseSSEStream[map[string]any](body, func(event map[string]any) error {
|
||||
tokenUsage, found := extractOpenAIStreamUsage(event)
|
||||
if found && chatConfig != nil {
|
||||
if found {
|
||||
applyStreamUsage(chatConfig, modelUsage, tokenUsage)
|
||||
}
|
||||
|
||||
|
||||
@@ -254,6 +254,64 @@ func TestNovitaStreamExtractsDeltaReasoningContent(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// TestNovitaStreamRecordsUsageWithoutChatConfig pins that the streaming
|
||||
// handler records a usage event even when chatConfig is nil. The old
|
||||
// `if found && chatConfig != nil` guard dropped the usage entirely for nil
|
||||
// chatConfig, diverging from the shared HandleStreamingResponse which
|
||||
// records usage whenever the stream carries it and only uses chatConfig to
|
||||
// expose UsageResult.
|
||||
func TestNovitaStreamRecordsUsageWithoutChatConfig(t *testing.T) {
|
||||
withSSRFBypass(t)
|
||||
ctx := t.Context()
|
||||
|
||||
sse := "" +
|
||||
`data: {"choices":[{"index":0,"delta":{"content":"hello"}}]}` + "\n" +
|
||||
`data: {"usage":{"prompt_tokens":3,"completion_tokens":5,"total_tokens":8}}` + "\n" +
|
||||
`data: {"choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}` + "\n" +
|
||||
`data: [DONE]` + "\n"
|
||||
|
||||
t.Run("nil chatConfig completes without dropping the usage event", func(t *testing.T) {
|
||||
srv := newNovitaSSEServer(t, "/openai/v1/chat/completions", sse)
|
||||
defer srv.Close()
|
||||
|
||||
apiKey := "test-key"
|
||||
err := newNovitaForTest(srv.URL).ChatStreamlyWithSender(
|
||||
ctx,
|
||||
"deepseek/deepseek-v3.1",
|
||||
[]Message{{Role: "user", Content: "x"}},
|
||||
&APIConfig{ApiKey: &apiKey}, nil, nil,
|
||||
func(c *string, r *string) error { return nil },
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatalf("stream: %v", err)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("non-nil chatConfig exposes the streamed usage", func(t *testing.T) {
|
||||
srv := newNovitaSSEServer(t, "/openai/v1/chat/completions", sse)
|
||||
defer srv.Close()
|
||||
|
||||
apiKey := "test-key"
|
||||
chatConfig := &ChatConfig{}
|
||||
err := newNovitaForTest(srv.URL).ChatStreamlyWithSender(
|
||||
ctx,
|
||||
"deepseek/deepseek-v3.1",
|
||||
[]Message{{Role: "user", Content: "x"}},
|
||||
&APIConfig{ApiKey: &apiKey}, chatConfig, nil,
|
||||
func(c *string, r *string) error { return nil },
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatalf("stream: %v", err)
|
||||
}
|
||||
if chatConfig.UsageResult == nil {
|
||||
t.Fatal("chatConfig.UsageResult is nil, want the streamed usage")
|
||||
}
|
||||
if chatConfig.UsageResult.PromptTokens != 3 || chatConfig.UsageResult.CompletionTokens != 5 || chatConfig.UsageResult.TotalTokens != 8 {
|
||||
t.Fatalf("UsageResult=%#v, want prompt=3 completion=5 total=8", chatConfig.UsageResult)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// TestNovitaChatPropagatesEnableThinking pins the maintainer's
|
||||
// requested behaviour: when ChatConfig.Thinking is set, the driver
|
||||
// MUST forward it as Novita's documented `enable_thinking` body field
|
||||
|
||||
Reference in New Issue
Block a user