Skip to content

Commit 8af963b

Browse files
localai-botmudler
localai-bot
and
mudler
authored
fix(streaming): comply with OpenAI usage / stream_options spec (#9815)
* fix(streaming): comply with OpenAI usage / stream_options spec (#8546) LocalAI emitted `"usage":{"prompt_tokens":0,...}` on every streamed chunk because `OpenAIResponse.Usage` was a value type without `omitempty`. The official OpenAI Node SDK and its consumers (continuedev/continue, Kilo Code, Roo Code, Zed, IntelliJ Continue) filter on a truthy `result.usage` to detect the trailing usage chunk; LocalAI's zero-but-non-null usage on every intermediate chunk made that filter swallow every content chunk and surface an empty chat response while the server log looked successful. Changes: - `core/schema/openai.go`: `Usage *OpenAIUsage \`json:"usage,omitempty"\`` so intermediate chunks no longer carry a `usage` key. Add `OpenAIRequest.StreamOptions` with `include_usage` to mirror OpenAI's request field. - `core/http/endpoints/openai/chat.go` and `completion.go`: keep using the `Usage` struct field as an in-process channel for the running cumulative, but strip it before JSON marshalling. When the request set `stream_options.include_usage: true`, emit a dedicated trailing chunk with `"choices": []` and the populated usage (matching the OpenAI spec and llama.cpp's server behavior). - `chat_emit.go`: new `streamUsageTrailerJSON` helper; drop the `usage` parameter from `buildNoActionFinalChunks` since chunks no longer carry usage. - Update `image.go`, `inpainting.go`, `edit.go` to wrap their Usage values with `&` for the new pointer field. - UI: send `stream_options:{include_usage:true}` from the React (`useChat.js`) and legacy (`static/chat.js`) chat clients so the token-count badge keeps populating now that the server is spec-compliant. Tests: - New `chat_stream_usage_test.go` pins the spec invariants: intermediate chunks have no `usage` key, the trailer JSON has `"choices":[]` and a populated `usage`, and `OpenAIRequest` parses `stream_options.include_usage`. - Update `chat_emit_test.go` to reflect that finals no longer embed usage. Verified against the live LocalAI instance: before the fix Continue's filter logic swallowed 16/16 token chunks; with the new shape it yields 4/5 and routes usage through the dedicated trailer chunk. Fixes #8546 Assisted-by: Claude:opus-4.7 [Claude Code] Signed-off-by: Ettore Di Giacinto <mudler@localai.io> * fix(streaming): silence errcheck on usage trailer Fprintf The new spec-compliant `stream_options.include_usage` trailer writes were flagged by errcheck since they're new code (golangci-lint runs new-from-merge-base on master); the surrounding `fmt.Fprintf` data: writes are grandfathered. Drop the return values explicitly to match the linter's contract without adding a nolint shim. Assisted-by: Claude:opus-4.7 [Claude Code] Signed-off-by: Ettore Di Giacinto <mudler@localai.io> --------- Signed-off-by: Ettore Di Giacinto <mudler@localai.io> Co-authored-by: Ettore Di Giacinto <mudler@localai.io>
1 parent 6e1dbae commit 8af963b

11 files changed

Lines changed: 342 additions & 63 deletions

File tree

‎core/http/endpoints/openai/chat.go‎

Lines changed: 39 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -131,13 +131,19 @@ func ChatEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, evaluator
131131
delta.Reasoning = &reasoningDelta
132132
}
133133

134+
// Usage rides as a struct field for the consumer to track the
135+
// running cumulative — it is stripped before JSON marshal so the
136+
// wire chunk stays spec-compliant (no `usage` on intermediate
137+
// chunks). The dedicated trailer chunk (when include_usage=true)
138+
// carries the final totals.
139+
usageForChunk := usage
134140
resp := schema.OpenAIResponse{
135141
ID: id,
136142
Created: created,
137143
Model: req.Model, // we have to return what the user sent here, due to OpenAI spec.
138144
Choices: []schema.Choice{{Delta: delta, Index: 0, FinishReason: nil}},
139145
Object: "chat.completion.chunk",
140-
Usage: usage,
146+
Usage: &usageForChunk,
141147
}
142148

143149
responses <- resp
@@ -164,7 +170,7 @@ func ChatEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, evaluator
164170
hasChatDeltaToolCalls := false
165171
hasChatDeltaContent := false
166172

167-
_, tokenUsage, chatDeltas, err := ComputeChoices(req, prompt, config, cl, startupOptions, loader, func(s string, c *[]schema.Choice) {}, func(s string, usage backend.TokenUsage) bool {
173+
_, _, chatDeltas, err := ComputeChoices(req, prompt, config, cl, startupOptions, loader, func(s string, c *[]schema.Choice) {}, func(s string, usage backend.TokenUsage) bool {
168174
result += s
169175

170176
// Track whether ChatDeltas from the C++ autoparser contain
@@ -387,16 +393,11 @@ func ChatEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, evaluator
387393

388394
switch {
389395
case noActionToRun:
390-
usage := schema.OpenAIUsage{
391-
PromptTokens: tokenUsage.Prompt,
392-
CompletionTokens: tokenUsage.Completion,
393-
TotalTokens: tokenUsage.Prompt + tokenUsage.Completion,
394-
}
395-
if extraUsage {
396-
usage.TimingTokenGeneration = tokenUsage.TimingTokenGeneration
397-
usage.TimingPromptProcessing = tokenUsage.TimingPromptProcessing
398-
}
399-
396+
// Token-cumulative usage is communicated to the streaming
397+
// consumer via the per-token callback's chunk struct (stripped
398+
// before wire marshal). The final usage trailer — when the
399+
// caller opted in with stream_options.include_usage — is built
400+
// by the outer streaming loop, not here.
400401
var result string
401402
if !sentInitialRole {
402403
var hqErr error
@@ -409,7 +410,7 @@ func ChatEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, evaluator
409410
for _, chunk := range buildNoActionFinalChunks(
410411
id, req.Model, created,
411412
sentInitialRole, sentReasoning,
412-
result, reasoning, usage,
413+
result, reasoning,
413414
) {
414415
responses <- chunk
415416
}
@@ -724,7 +725,13 @@ func ChatEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, evaluator
724725
xlog.Debug("No choices in the response, skipping")
725726
continue
726727
}
727-
usage = &ev.Usage // Copy a pointer to the latest usage chunk so that the stop message can reference it
728+
// Capture the running cumulative usage from this chunk
729+
// (when present) so the include_usage trailer can carry
730+
// the final totals. Usage is stripped before marshal
731+
// below so the wire chunk stays spec-compliant.
732+
if ev.Usage != nil {
733+
usage = ev.Usage
734+
}
728735
if len(ev.Choices[0].Delta.ToolCalls) > 0 {
729736
toolsCalled = true
730737
// Collect and merge tool call deltas for MCP execution
@@ -740,6 +747,11 @@ func ChatEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, evaluator
740747
collectedContent += *sp
741748
}
742749
}
750+
// OpenAI streaming spec: intermediate chunks must NOT
751+
// carry a `usage` field. Strip the tracking copy
752+
// before marshalling — usage is delivered via the
753+
// dedicated trailer chunk when include_usage=true.
754+
ev.Usage = nil
743755
respData, err := json.Marshal(ev)
744756
if err != nil {
745757
xlog.Debug("Failed to marshal response", "error", err)
@@ -888,6 +900,9 @@ func ChatEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, evaluator
888900
finishReason = FinishReasonFunctionCall
889901
}
890902

903+
// Final delta chunk: empty delta with finish_reason set. Per
904+
// OpenAI streaming spec this chunk does NOT carry usage —
905+
// the optional trailer (below) does, gated on include_usage.
891906
resp := &schema.OpenAIResponse{
892907
ID: id,
893908
Created: created,
@@ -899,11 +914,18 @@ func ChatEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, evaluator
899914
Delta: &schema.Message{},
900915
}},
901916
Object: "chat.completion.chunk",
902-
Usage: *usage,
903917
}
904918
respData, _ := json.Marshal(resp)
905-
906919
fmt.Fprintf(c.Response().Writer, "data: %s\n\n", respData)
920+
921+
// Trailing usage chunk per OpenAI spec: emit only when the
922+
// caller opted in via stream_options.include_usage. Shape:
923+
// {"choices":[],"usage":{...},"object":"chat.completion.chunk",...}
924+
if input.StreamOptions != nil && input.StreamOptions.IncludeUsage && usage != nil {
925+
trailer := streamUsageTrailerJSON(id, input.Model, created, *usage)
926+
_, _ = fmt.Fprintf(c.Response().Writer, "data: %s\n\n", trailer)
927+
}
928+
907929
fmt.Fprintf(c.Response().Writer, "data: [DONE]\n\n")
908930
c.Response().Flush()
909931
xlog.Debug("Stream ended")
@@ -1263,7 +1285,7 @@ func ChatEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, evaluator
12631285
Model: input.Model, // we have to return what the user sent here, due to OpenAI spec.
12641286
Choices: result,
12651287
Object: "chat.completion",
1266-
Usage: usage,
1288+
Usage: &usage,
12671289
}
12681290
respData, _ := json.Marshal(resp)
12691291
xlog.Debug("Response", "response", string(respData))

‎core/http/endpoints/openai/chat_emit.go‎

Lines changed: 50 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,45 @@
11
package openai
22

33
import (
4+
"encoding/json"
45
"fmt"
56

67
"github.com/mudler/LocalAI/core/schema"
78
"github.com/mudler/LocalAI/pkg/functions"
89
)
910

11+
// streamUsageTrailerJSON returns the bytes of the OpenAI-spec trailing usage
12+
// chunk emitted in streaming completions when the request opts in via
13+
// `stream_options.include_usage: true`. The shape is:
14+
//
15+
// {"id":"...","object":"chat.completion.chunk","created":N,
16+
// "model":"...","choices":[],"usage":{...}}
17+
//
18+
// `choices` is intentionally an empty array (not absent, not null) — that is
19+
// what the OpenAI spec mandates, and what consumers like the official OpenAI
20+
// SDK and Continue's openai-adapter look for to recognise this as the usage
21+
// chunk rather than a content chunk. schema.OpenAIResponse has `omitempty`
22+
// on Choices, so we cannot reuse it for the trailer.
23+
func streamUsageTrailerJSON(id, model string, created int, usage schema.OpenAIUsage) []byte {
24+
trailer := struct {
25+
ID string `json:"id"`
26+
Created int `json:"created"`
27+
Model string `json:"model"`
28+
Object string `json:"object"`
29+
Choices []schema.Choice `json:"choices"`
30+
Usage schema.OpenAIUsage `json:"usage"`
31+
}{
32+
ID: id,
33+
Created: created,
34+
Model: model,
35+
Object: "chat.completion.chunk",
36+
Choices: []schema.Choice{},
37+
Usage: usage,
38+
}
39+
b, _ := json.Marshal(trailer)
40+
return b
41+
}
42+
1043
// hasRealCall reports whether functionResults contains at least one
1144
// entry whose Name is something other than the noAction sentinel.
1245
// Used by processTools to decide between the "answer the question"
@@ -25,44 +58,48 @@ func hasRealCall(functionResults []functions.FuncCallResults, noAction string) b
2558
// pseudo-function or emitted no tool calls at all).
2659
//
2760
// When content was already streamed (contentAlreadyStreamed=true) the
28-
// helper emits a single trailing usage chunk, optionally carrying
29-
// reasoning that was produced but not streamed incrementally. When
30-
// content was not streamed it emits a role chunk followed by a
31-
// content+reasoning+usage chunk — the "send everything at once" fallback.
61+
// helper emits a trailing reasoning chunk if any non-streamed reasoning
62+
// remains, else nothing. When content was not streamed it emits a role
63+
// chunk followed by a content (+reasoning) chunk — the "send everything
64+
// at once" fallback.
3265
//
3366
// Reasoning re-emission is guarded by reasoningAlreadyStreamed, not by
3467
// probing the extractor's Go-side state: the C++ autoparser delivers
3568
// reasoning through ProcessChatDeltaReasoning which populates a
3669
// separate accumulator that extractor.Reasoning() does not expose.
3770
// Without this guard the callback would stream reasoning incrementally
3871
// and the final chunk would duplicate it.
72+
//
73+
// The returned chunks intentionally do NOT carry a `usage` field. The
74+
// usage trailer is emitted separately by the streaming handler when
75+
// `stream_options.include_usage` is true, per OpenAI spec.
3976
func buildNoActionFinalChunks(
4077
id, model string,
4178
created int,
4279
contentAlreadyStreamed bool,
4380
reasoningAlreadyStreamed bool,
4481
content string,
4582
reasoning string,
46-
usage schema.OpenAIUsage,
4783
) []schema.OpenAIResponse {
4884
var out []schema.OpenAIResponse
4985

5086
if contentAlreadyStreamed {
51-
delta := &schema.Message{}
52-
if reasoning != "" && !reasoningAlreadyStreamed {
53-
r := reasoning
54-
delta.Reasoning = &r
87+
if reasoning == "" || reasoningAlreadyStreamed {
88+
return nil
5589
}
90+
r := reasoning
5691
out = append(out, schema.OpenAIResponse{
5792
ID: id, Created: created, Model: model,
58-
Choices: []schema.Choice{{Delta: delta, Index: 0}},
59-
Object: "chat.completion.chunk",
60-
Usage: usage,
93+
Choices: []schema.Choice{{
94+
Delta: &schema.Message{Reasoning: &r},
95+
Index: 0,
96+
}},
97+
Object: "chat.completion.chunk",
6198
})
6299
return out
63100
}
64101

65-
// Content was not streamed — send role, then content (+reasoning) + usage.
102+
// Content was not streamed — send role, then content (+reasoning).
66103
out = append(out, schema.OpenAIResponse{
67104
ID: id, Created: created, Model: model,
68105
Choices: []schema.Choice{{
@@ -82,7 +119,6 @@ func buildNoActionFinalChunks(
82119
ID: id, Created: created, Model: model,
83120
Choices: []schema.Choice{{Delta: delta, Index: 0}},
84121
Object: "chat.completion.chunk",
85-
Usage: usage,
86122
})
87123
return out
88124
}

‎core/http/endpoints/openai/chat_emit_test.go‎

Lines changed: 22 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -609,54 +609,52 @@ var _ = Describe("buildNoActionFinalChunks", func() {
609609
testModel = "test-model"
610610
testCreated = 1700000000
611611
)
612-
usage := schema.OpenAIUsage{PromptTokens: 5, CompletionTokens: 7, TotalTokens: 12}
613612

614-
Describe("Content streamed — trailing usage chunk", func() {
615-
It("emits just one chunk with usage, no content, no reasoning when reasoning was streamed", func() {
613+
Describe("Content streamed — trailing reasoning only", func() {
614+
It("emits nothing when content and reasoning were already streamed", func() {
615+
// Before the streaming-usage-spec fix this branch emitted a
616+
// content-less chunk solely to carry `usage`. Per the OpenAI
617+
// spec usage no longer rides on delta chunks; the dedicated
618+
// trailer (when include_usage=true) carries it instead — so
619+
// with nothing to deliver the helper returns no chunks.
616620
chunks := buildNoActionFinalChunks(
617621
testID, testModel, testCreated,
618622
true, true,
619-
"", "already-streamed-reasoning", usage,
623+
"", "already-streamed-reasoning",
620624
)
621-
622-
Expect(chunks).To(HaveLen(1))
623-
Expect(chunks[0].Usage.TotalTokens).To(Equal(12))
624-
Expect(contentOf(chunks[0])).To(BeEmpty())
625-
Expect(reasoningOf(chunks[0])).To(BeEmpty(),
626-
"reasoning must not be re-emitted once it was streamed via the callback")
625+
Expect(chunks).To(BeEmpty())
627626
})
628627

629628
It("emits a trailing reasoning delivery when reasoning came only at end", func() {
630629
chunks := buildNoActionFinalChunks(
631630
testID, testModel, testCreated,
632631
true, false,
633-
"", "autoparser final reasoning", usage,
632+
"", "autoparser final reasoning",
634633
)
635634

636635
Expect(chunks).To(HaveLen(1))
637636
Expect(reasoningOf(chunks[0])).To(Equal("autoparser final reasoning"))
638637
Expect(contentOf(chunks[0])).To(BeEmpty())
639-
Expect(chunks[0].Usage.TotalTokens).To(Equal(12))
638+
Expect(chunks[0].Usage).To(BeNil(),
639+
"intermediate chunks must not carry usage per OpenAI spec")
640640
})
641641

642-
It("omits reasoning when it's empty regardless of streamed flag", func() {
642+
It("returns no chunks when reasoning is empty and content was streamed", func() {
643643
chunks := buildNoActionFinalChunks(
644644
testID, testModel, testCreated,
645645
true, false,
646-
"", "", usage,
646+
"", "",
647647
)
648-
649-
Expect(chunks).To(HaveLen(1))
650-
Expect(reasoningOf(chunks[0])).To(BeEmpty())
648+
Expect(chunks).To(BeEmpty())
651649
})
652650
})
653651

654-
Describe("Content not streamed — role, then content+usage", func() {
652+
Describe("Content not streamed — role, then content", func() {
655653
It("emits role chunk then content chunk without reasoning when reasoning was streamed", func() {
656654
chunks := buildNoActionFinalChunks(
657655
testID, testModel, testCreated,
658656
false, true,
659-
"the answer", "already-streamed-reasoning", usage,
657+
"the answer", "already-streamed-reasoning",
660658
)
661659

662660
Expect(chunks).To(HaveLen(2))
@@ -666,29 +664,29 @@ var _ = Describe("buildNoActionFinalChunks", func() {
666664
Expect(contentOf(chunks[1])).To(Equal("the answer"))
667665
Expect(reasoningOf(chunks[1])).To(BeEmpty(),
668666
"reasoning must not be re-emitted if it was streamed earlier")
669-
Expect(chunks[1].Usage.TotalTokens).To(Equal(12))
667+
Expect(chunks[1].Usage).To(BeNil())
670668
})
671669

672670
It("emits role, then content+reasoning when reasoning was not streamed", func() {
673671
chunks := buildNoActionFinalChunks(
674672
testID, testModel, testCreated,
675673
false, false,
676-
"the answer", "autoparser final reasoning", usage,
674+
"the answer", "autoparser final reasoning",
677675
)
678676

679677
Expect(chunks).To(HaveLen(2))
680678
Expect(chunks[0].Choices[0].Delta.Role).To(Equal("assistant"))
681679

682680
Expect(contentOf(chunks[1])).To(Equal("the answer"))
683681
Expect(reasoningOf(chunks[1])).To(Equal("autoparser final reasoning"))
684-
Expect(chunks[1].Usage.TotalTokens).To(Equal(12))
682+
Expect(chunks[1].Usage).To(BeNil())
685683
})
686684

687685
It("still emits content even when reasoning is empty", func() {
688686
chunks := buildNoActionFinalChunks(
689687
testID, testModel, testCreated,
690688
false, false,
691-
"just an answer", "", usage,
689+
"just an answer", "",
692690
)
693691

694692
Expect(chunks).To(HaveLen(2))
@@ -702,7 +700,7 @@ var _ = Describe("buildNoActionFinalChunks", func() {
702700
chunks := buildNoActionFinalChunks(
703701
testID, testModel, testCreated,
704702
false, false,
705-
"hi", "reasoning", usage,
703+
"hi", "reasoning",
706704
)
707705
for i, ch := range chunks {
708706
Expect(ch.ID).To(Equal(testID), "chunk[%d] ID", i)

0 commit comments

Comments
 (0)