Skip to content

Commit 716ddd6

Browse files
authored
feat(autoparser): prefer chat deltas from backends when emitted (#9224)
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
1 parent 223deb9 commit 716ddd6

5 files changed

Lines changed: 295 additions & 4 deletions

File tree

core/backend/llm.go

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,27 @@ type TokenUsage struct {
3636
Completion int
3737
TimingPromptProcessing float64
3838
TimingTokenGeneration float64
39+
ChatDeltas []*proto.ChatDelta // per-chunk deltas from C++ autoparser (only set during streaming)
40+
}
41+
42+
// HasChatDeltaContent returns true if any chat delta carries content or reasoning text.
43+
// Used to decide whether to prefer C++ autoparser deltas over Go-side tag extraction.
44+
func (t TokenUsage) HasChatDeltaContent() bool {
45+
for _, d := range t.ChatDeltas {
46+
if d.Content != "" || d.ReasoningContent != "" {
47+
return true
48+
}
49+
}
50+
return false
51+
}
52+
53+
// ChatDeltaReasoningAndContent extracts accumulated reasoning and content from chat deltas.
54+
func (t TokenUsage) ChatDeltaReasoningAndContent() (reasoning, content string) {
55+
for _, d := range t.ChatDeltas {
56+
content += d.Content
57+
reasoning += d.ReasoningContent
58+
}
59+
return reasoning, content
3960
}
4061

4162
// ModelInferenceFunc is a test-friendly indirection to call model inference logic.
@@ -171,6 +192,9 @@ func ModelInference(ctx context.Context, s string, messages schema.Messages, ima
171192
allChatDeltas = append(allChatDeltas, reply.ChatDeltas...)
172193
}
173194

195+
// Attach per-chunk chat deltas to tokenUsage so the callback can use them
196+
tokenUsage.ChatDeltas = reply.ChatDeltas
197+
174198
// Parse logprobs from reply if present (collect from last chunk that has them)
175199
if len(reply.Logprobs) > 0 {
176200
var parsedLogprobs schema.Logprobs
@@ -200,6 +224,9 @@ func ModelInference(ctx context.Context, s string, messages schema.Messages, ima
200224
if len(msg) == 0 {
201225
tokenCallback("", tokenUsage)
202226
}
227+
228+
// Clear per-chunk deltas so they don't leak to the next chunk
229+
tokenUsage.ChatDeltas = nil
203230
})
204231
if len(allChatDeltas) > 0 {
205232
xlog.Debug("[ChatDeltas] streaming completed, accumulated deltas from C++ autoparser", "total_deltas", len(allChatDeltas))

core/backend/llm_test.go

Lines changed: 109 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ import (
44
. "github.com/mudler/LocalAI/core/backend"
55
"github.com/mudler/LocalAI/core/config"
66
"github.com/mudler/LocalAI/core/schema"
7+
pb "github.com/mudler/LocalAI/pkg/grpc/proto"
78

89
. "github.com/onsi/ginkgo/v2"
910
. "github.com/onsi/gomega"
@@ -107,3 +108,111 @@ var _ = Describe("LLM tests", func() {
107108
})
108109
})
109110
})
111+
112+
var _ = Describe("TokenUsage ChatDelta helpers", func() {
113+
Describe("HasChatDeltaContent", func() {
114+
It("should return false when ChatDeltas is nil", func() {
115+
usage := TokenUsage{}
116+
Expect(usage.HasChatDeltaContent()).To(BeFalse())
117+
})
118+
119+
It("should return false when ChatDeltas is empty", func() {
120+
usage := TokenUsage{ChatDeltas: []*pb.ChatDelta{}}
121+
Expect(usage.HasChatDeltaContent()).To(BeFalse())
122+
})
123+
124+
It("should return false when all deltas have empty content and reasoning", func() {
125+
usage := TokenUsage{
126+
ChatDeltas: []*pb.ChatDelta{
127+
{Content: "", ReasoningContent: ""},
128+
{Content: ""},
129+
},
130+
}
131+
Expect(usage.HasChatDeltaContent()).To(BeFalse())
132+
})
133+
134+
It("should return true when a delta has content", func() {
135+
usage := TokenUsage{
136+
ChatDeltas: []*pb.ChatDelta{
137+
{Content: "hello"},
138+
},
139+
}
140+
Expect(usage.HasChatDeltaContent()).To(BeTrue())
141+
})
142+
143+
It("should return true when a delta has reasoning content", func() {
144+
usage := TokenUsage{
145+
ChatDeltas: []*pb.ChatDelta{
146+
{ReasoningContent: "thinking..."},
147+
},
148+
}
149+
Expect(usage.HasChatDeltaContent()).To(BeTrue())
150+
})
151+
152+
It("should return true when a delta has both content and reasoning", func() {
153+
usage := TokenUsage{
154+
ChatDeltas: []*pb.ChatDelta{
155+
{Content: "hello", ReasoningContent: "thinking..."},
156+
},
157+
}
158+
Expect(usage.HasChatDeltaContent()).To(BeTrue())
159+
})
160+
})
161+
162+
Describe("ChatDeltaReasoningAndContent", func() {
163+
It("should return empty strings when ChatDeltas is nil", func() {
164+
usage := TokenUsage{}
165+
reasoning, content := usage.ChatDeltaReasoningAndContent()
166+
Expect(reasoning).To(BeEmpty())
167+
Expect(content).To(BeEmpty())
168+
})
169+
170+
It("should concatenate content from multiple deltas", func() {
171+
usage := TokenUsage{
172+
ChatDeltas: []*pb.ChatDelta{
173+
{Content: "Hello"},
174+
{Content: " world"},
175+
},
176+
}
177+
reasoning, content := usage.ChatDeltaReasoningAndContent()
178+
Expect(content).To(Equal("Hello world"))
179+
Expect(reasoning).To(BeEmpty())
180+
})
181+
182+
It("should concatenate reasoning from multiple deltas", func() {
183+
usage := TokenUsage{
184+
ChatDeltas: []*pb.ChatDelta{
185+
{ReasoningContent: "step 1"},
186+
{ReasoningContent: " step 2"},
187+
},
188+
}
189+
reasoning, content := usage.ChatDeltaReasoningAndContent()
190+
Expect(reasoning).To(Equal("step 1 step 2"))
191+
Expect(content).To(BeEmpty())
192+
})
193+
194+
It("should separate reasoning and content from mixed deltas", func() {
195+
usage := TokenUsage{
196+
ChatDeltas: []*pb.ChatDelta{
197+
{ReasoningContent: "thinking"},
198+
{Content: "answer"},
199+
},
200+
}
201+
reasoning, content := usage.ChatDeltaReasoningAndContent()
202+
Expect(reasoning).To(Equal("thinking"))
203+
Expect(content).To(Equal("answer"))
204+
})
205+
206+
It("should handle deltas with both fields set", func() {
207+
usage := TokenUsage{
208+
ChatDeltas: []*pb.ChatDelta{
209+
{Content: "a", ReasoningContent: "r1"},
210+
{Content: "b", ReasoningContent: "r2"},
211+
},
212+
}
213+
reasoning, content := usage.ChatDeltaReasoningAndContent()
214+
Expect(reasoning).To(Equal("r1r2"))
215+
Expect(content).To(Equal("ab"))
216+
})
217+
})
218+
})

core/http/endpoints/openai/chat.go

Lines changed: 23 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -81,7 +81,17 @@ func ChatEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, evaluator
8181
extractor := reason.NewReasoningExtractor(thinkingStartToken, config.ReasoningConfig)
8282

8383
_, _, _, err := ComputeChoices(req, s, config, cl, startupOptions, loader, func(s string, c *[]schema.Choice) {}, func(s string, tokenUsage backend.TokenUsage) bool {
84-
reasoningDelta, contentDelta := extractor.ProcessToken(s)
84+
var reasoningDelta, contentDelta string
85+
86+
// Prefer pre-parsed chat deltas from C++ autoparser when available
87+
if tokenUsage.HasChatDeltaContent() {
88+
reasoningDelta, contentDelta = tokenUsage.ChatDeltaReasoningAndContent()
89+
// Keep extractor state consistent for fallback
90+
extractor.ProcessToken(s)
91+
} else {
92+
// Fallback: Go-side extraction from raw text
93+
reasoningDelta, contentDelta = extractor.ProcessToken(s)
94+
}
8595

8696
usage := schema.OpenAIUsage{
8797
PromptTokens: tokenUsage.Prompt,
@@ -133,7 +143,18 @@ func ChatEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, evaluator
133143

134144
_, tokenUsage, chatDeltas, err := ComputeChoices(req, prompt, config, cl, startupOptions, loader, func(s string, c *[]schema.Choice) {}, func(s string, usage backend.TokenUsage) bool {
135145
result += s
136-
reasoningDelta, contentDelta := extractor.ProcessToken(s)
146+
147+
var reasoningDelta, contentDelta string
148+
149+
// Prefer pre-parsed chat deltas from C++ autoparser when available
150+
if usage.HasChatDeltaContent() {
151+
reasoningDelta, contentDelta = usage.ChatDeltaReasoningAndContent()
152+
// Keep extractor state consistent for fallback
153+
extractor.ProcessToken(s)
154+
} else {
155+
// Fallback: Go-side extraction from raw text
156+
reasoningDelta, contentDelta = extractor.ProcessToken(s)
157+
}
137158

138159
// Emit reasoning deltas in their own SSE chunks before any tool-call chunks
139160
// (OpenAI spec: reasoning and tool_calls never share a delta)

core/http/endpoints/openai/inference_test.go

Lines changed: 119 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -398,5 +398,124 @@ var _ = Describe("ComputeChoices", func() {
398398
Expect(choices).To(HaveLen(1))
399399
Expect(streamedTokens).To(Equal([]string{"Hello", " world"}))
400400
})
401+
402+
It("should pass chat deltas through TokenUsage during streaming", func() {
403+
var receivedDeltas [][]*pb.ChatDelta
404+
backend.ModelInferenceFunc = func(
405+
ctx context.Context, s string, messages schema.Messages,
406+
images, videos, audios []string,
407+
loader *model.ModelLoader, c *config.ModelConfig, cl *config.ModelConfigLoader,
408+
o *config.ApplicationConfig,
409+
tokenCallback func(string, backend.TokenUsage) bool,
410+
tools, toolChoice string,
411+
logprobs, topLogprobs *int,
412+
logitBias map[string]float64,
413+
metadata map[string]string,
414+
) (func() (backend.LLMResponse, error), error) {
415+
predFunc := func() (backend.LLMResponse, error) {
416+
if tokenCallback != nil {
417+
// Simulate C++ autoparser sending reasoning in chat deltas
418+
tokenCallback("<|channel>thought\nthinking\n<channel|>", backend.TokenUsage{
419+
Prompt: 5,
420+
ChatDeltas: []*pb.ChatDelta{
421+
{ReasoningContent: "thinking"},
422+
},
423+
})
424+
tokenCallback("Hello!", backend.TokenUsage{
425+
Prompt: 5, Completion: 3,
426+
ChatDeltas: []*pb.ChatDelta{
427+
{Content: "Hello!"},
428+
},
429+
})
430+
}
431+
return backend.LLMResponse{
432+
Response: "<|channel>thought\nthinking\n<channel|>Hello!",
433+
Usage: backend.TokenUsage{Prompt: 5, Completion: 3},
434+
ChatDeltas: []*pb.ChatDelta{
435+
{ReasoningContent: "thinking"},
436+
{Content: "Hello!"},
437+
},
438+
}, nil
439+
}
440+
return predFunc, nil
441+
}
442+
443+
choices, _, deltas, err := ComputeChoices(
444+
makeReq(), "test", cfg, nil, appCfg, nil,
445+
func(s string, c *[]schema.Choice) {
446+
*c = append(*c, schema.Choice{Text: s})
447+
},
448+
func(s string, usage backend.TokenUsage) bool {
449+
// Capture chat deltas received per-chunk
450+
if len(usage.ChatDeltas) > 0 {
451+
receivedDeltas = append(receivedDeltas, usage.ChatDeltas)
452+
}
453+
return true
454+
},
455+
)
456+
Expect(err).ToNot(HaveOccurred())
457+
Expect(choices).To(HaveLen(1))
458+
459+
// Verify per-chunk deltas were received during streaming
460+
Expect(receivedDeltas).To(HaveLen(2))
461+
Expect(receivedDeltas[0][0].ReasoningContent).To(Equal("thinking"))
462+
Expect(receivedDeltas[1][0].Content).To(Equal("Hello!"))
463+
464+
// Verify final accumulated deltas are also returned
465+
Expect(deltas).To(HaveLen(2))
466+
Expect(deltas[0].ReasoningContent).To(Equal("thinking"))
467+
Expect(deltas[1].Content).To(Equal("Hello!"))
468+
})
469+
470+
It("should prefer chat deltas over raw text when HasChatDeltaContent is true", func() {
471+
// Verify that the callback can distinguish between
472+
// chunks with and without chat deltas
473+
var withDeltas, withoutDeltas int
474+
backend.ModelInferenceFunc = func(
475+
ctx context.Context, s string, messages schema.Messages,
476+
images, videos, audios []string,
477+
loader *model.ModelLoader, c *config.ModelConfig, cl *config.ModelConfigLoader,
478+
o *config.ApplicationConfig,
479+
tokenCallback func(string, backend.TokenUsage) bool,
480+
tools, toolChoice string,
481+
logprobs, topLogprobs *int,
482+
logitBias map[string]float64,
483+
metadata map[string]string,
484+
) (func() (backend.LLMResponse, error), error) {
485+
predFunc := func() (backend.LLMResponse, error) {
486+
if tokenCallback != nil {
487+
// Chunk with chat deltas (C++ autoparser active)
488+
tokenCallback("raw-text", backend.TokenUsage{
489+
ChatDeltas: []*pb.ChatDelta{{Content: "parsed-content"}},
490+
})
491+
// Chunk without chat deltas (fallback)
492+
tokenCallback("fallback-text", backend.TokenUsage{})
493+
}
494+
return backend.LLMResponse{Response: "raw-textfallback-text"}, nil
495+
}
496+
return predFunc, nil
497+
}
498+
499+
_, _, _, err := ComputeChoices(
500+
makeReq(), "test", cfg, nil, appCfg, nil,
501+
func(s string, c *[]schema.Choice) {
502+
*c = append(*c, schema.Choice{Text: s})
503+
},
504+
func(s string, usage backend.TokenUsage) bool {
505+
if usage.HasChatDeltaContent() {
506+
withDeltas++
507+
r, c := usage.ChatDeltaReasoningAndContent()
508+
Expect(c).To(Equal("parsed-content"))
509+
Expect(r).To(BeEmpty())
510+
} else {
511+
withoutDeltas++
512+
}
513+
return true
514+
},
515+
)
516+
Expect(err).ToNot(HaveOccurred())
517+
Expect(withDeltas).To(Equal(1))
518+
Expect(withoutDeltas).To(Equal(1))
519+
})
401520
})
402521
})

core/http/endpoints/openresponses/responses.go

Lines changed: 17 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1819,7 +1819,14 @@ func handleOpenResponsesStream(c echo.Context, responseID string, createdAt int6
18191819

18201820
// If no tool calls detected yet, handle reasoning and text
18211821
if !inToolCallMode {
1822-
reasoningDelta, contentDelta := extractor.ProcessToken(token)
1822+
var reasoningDelta, contentDelta string
1823+
// Prefer pre-parsed chat deltas from C++ autoparser when available
1824+
if tokenUsage.HasChatDeltaContent() {
1825+
reasoningDelta, contentDelta = tokenUsage.ChatDeltaReasoningAndContent()
1826+
extractor.ProcessToken(token) // keep state consistent
1827+
} else {
1828+
reasoningDelta, contentDelta = extractor.ProcessToken(token)
1829+
}
18231830

18241831
// Handle reasoning item
18251832
if extractor.Reasoning() != "" {
@@ -2338,7 +2345,15 @@ func handleOpenResponsesStream(c echo.Context, responseID string, createdAt int6
23382345
// Stream text deltas with reasoning extraction
23392346
tokenCallback := func(token string, tokenUsage backend.TokenUsage) bool {
23402347
accumulatedText += token
2341-
reasoningDelta, contentDelta := extractor.ProcessToken(token)
2348+
2349+
var reasoningDelta, contentDelta string
2350+
// Prefer pre-parsed chat deltas from C++ autoparser when available
2351+
if tokenUsage.HasChatDeltaContent() {
2352+
reasoningDelta, contentDelta = tokenUsage.ChatDeltaReasoningAndContent()
2353+
extractor.ProcessToken(token) // keep state consistent
2354+
} else {
2355+
reasoningDelta, contentDelta = extractor.ProcessToken(token)
2356+
}
23422357

23432358
// Handle reasoning item
23442359
if extractor.Reasoning() != "" {

0 commit comments

Comments
 (0)