Skip to content

Commit 8b8a19d

Browse files
committed
fix(streaming): deduplicate tool call emissions during streaming
The Go-side incremental JSON parser was emitting the same tool call on every streaming token because it lacked the len > lastEmittedCount guard that the XML parser had. On top of that, the post-streaming default: case re-emitted all tool calls from index 0, duplicating everything. This produced duplicate delta.tool_calls events causing clients to accumulate arguments as "{args}{args}" — invalid JSON. Fixes: - JSON incremental parser: add len(jsonResults) > lastEmittedCount guard and loop from lastEmittedCount (matching the XML parser pattern) - Post-streaming default: case: skip i < lastEmittedCount entries that were already emitted during streaming - JSON parser: use blocking channel send (matching XML parser behavior)
1 parent b0d9ce4 commit 8b8a19d

3 files changed

Lines changed: 223 additions & 44 deletions

File tree

core/http/endpoints/openai/chat.go

Lines changed: 48 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -265,55 +265,52 @@ func ChatEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, evaluator
265265
lastEmittedCount = len(partialResults)
266266
}
267267
} else {
268-
// Try JSON tool call parsing for streaming
269-
// Check if the result looks like JSON tool calls
268+
// Try JSON tool call parsing for streaming.
269+
// Only emit NEW tool calls (same guard as XML parser above).
270270
jsonResults, jsonErr := functions.ParseJSONIterative(cleanedResult, true)
271-
if jsonErr == nil && len(jsonResults) > 0 {
272-
// Check if these are tool calls (have "name" and optionally "arguments")
273-
for _, jsonObj := range jsonResults {
274-
if name, ok := jsonObj["name"].(string); ok && name != "" {
275-
// This looks like a tool call
276-
args := "{}"
277-
if argsVal, ok := jsonObj["arguments"]; ok {
278-
if argsStr, ok := argsVal.(string); ok {
279-
args = argsStr
280-
} else {
281-
argsBytes, _ := json.Marshal(argsVal)
282-
args = string(argsBytes)
283-
}
271+
if jsonErr == nil && len(jsonResults) > lastEmittedCount {
272+
for i := lastEmittedCount; i < len(jsonResults); i++ {
273+
jsonObj := jsonResults[i]
274+
name, ok := jsonObj["name"].(string)
275+
if !ok || name == "" {
276+
continue
277+
}
278+
args := "{}"
279+
if argsVal, ok := jsonObj["arguments"]; ok {
280+
if argsStr, ok := argsVal.(string); ok {
281+
args = argsStr
282+
} else {
283+
argsBytes, _ := json.Marshal(argsVal)
284+
args = string(argsBytes)
284285
}
285-
// Emit tool call
286-
initialMessage := schema.OpenAIResponse{
287-
ID: id,
288-
Created: created,
289-
Model: req.Model,
290-
Choices: []schema.Choice{{
291-
Delta: &schema.Message{
292-
Role: "assistant",
293-
ToolCalls: []schema.ToolCall{
294-
{
295-
Index: lastEmittedCount,
296-
ID: id,
297-
Type: "function",
298-
FunctionCall: schema.FunctionCall{
299-
Name: name,
300-
Arguments: args,
301-
},
286+
}
287+
initialMessage := schema.OpenAIResponse{
288+
ID: id,
289+
Created: created,
290+
Model: req.Model,
291+
Choices: []schema.Choice{{
292+
Delta: &schema.Message{
293+
Role: "assistant",
294+
ToolCalls: []schema.ToolCall{
295+
{
296+
Index: i,
297+
ID: id,
298+
Type: "function",
299+
FunctionCall: schema.FunctionCall{
300+
Name: name,
301+
Arguments: args,
302302
},
303303
},
304304
},
305-
Index: 0,
306-
FinishReason: nil,
307-
}},
308-
Object: "chat.completion.chunk",
309-
}
310-
select {
311-
case responses <- initialMessage:
312-
default:
313-
}
314-
lastEmittedCount++
305+
},
306+
Index: 0,
307+
FinishReason: nil,
308+
}},
309+
Object: "chat.completion.chunk",
315310
}
311+
responses <- initialMessage
316312
}
313+
lastEmittedCount = len(jsonResults)
317314
}
318315
}
319316
return true
@@ -426,10 +423,17 @@ func ChatEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, evaluator
426423
toolCallID = id
427424
}
428425

426+
if i < lastEmittedCount {
427+
// Already emitted during streaming by the incremental
428+
// JSON/XML parser — skip to avoid duplicate tool calls.
429+
continue
430+
}
431+
432+
// Tool call not yet emitted — send name + args (two chunks).
429433
initialMessage := schema.OpenAIResponse{
430434
ID: id,
431435
Created: created,
432-
Model: req.Model, // we have to return what the user sent here, due to OpenAI spec.
436+
Model: req.Model,
433437
Choices: []schema.Choice{{
434438
Delta: &schema.Message{
435439
Role: "assistant",
@@ -454,7 +458,7 @@ func ChatEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, evaluator
454458
responses <- schema.OpenAIResponse{
455459
ID: id,
456460
Created: created,
457-
Model: req.Model, // we have to return what the user sent here, due to OpenAI spec.
461+
Model: req.Model,
458462
Choices: []schema.Choice{{
459463
Delta: &schema.Message{
460464
Role: "assistant",

tests/e2e/mock-backend/main.go

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -95,6 +95,15 @@ func (m *MockBackend) Predict(ctx context.Context, in *pb.PredictOptions) (*pb.R
9595
}, nil
9696
}
9797

98+
// Simulate multiple tool calls in a single response (Go-side JSON parser path).
99+
if strings.Contains(in.Prompt, "MULTI_TOOL_CALL") {
100+
return &pb.Reply{
101+
Message: []byte(`{"name": "get_weather", "arguments": {"location": "Rome"}}
102+
{"name": "get_weather", "arguments": {"location": "Paris"}}`),
103+
Tokens: 30,
104+
PromptTokens: 10,
105+
}, nil
106+
}
98107
var response string
99108
toolName := mockToolNameFromRequest(in)
100109
if toolName != "" && !promptHasToolResults(in.Prompt) {
@@ -199,6 +208,38 @@ func (m *MockBackend) PredictStream(in *pb.PredictOptions, stream pb.Backend_Pre
199208
return nil
200209
}
201210

211+
// Simulate tool calls streamed as whole JSON objects (Go-side parser path).
212+
// Each object is sent as a complete chunk so the incremental parser can
213+
// detect tool calls mid-stream (unlike char-by-char which only parses after
214+
// streaming completes).
215+
if strings.Contains(in.Prompt, "MULTI_TOOL_CALL") {
216+
chunks := []string{
217+
`{"name": "get_weather", "arguments": {"location": "Rome"}}`,
218+
"\n",
219+
`{"name": "get_weather", "arguments": {"location": "Paris"}}`,
220+
}
221+
for i, chunk := range chunks {
222+
if err := stream.Send(&pb.Reply{
223+
Message: []byte(chunk),
224+
Tokens: int32(i + 1),
225+
}); err != nil {
226+
return err
227+
}
228+
}
229+
return nil
230+
}
231+
232+
// Simulate single tool call streamed as whole JSON (Go-side parser path).
233+
if strings.Contains(in.Prompt, "SINGLE_TOOL_CALL") {
234+
if err := stream.Send(&pb.Reply{
235+
Message: []byte(`{"name": "get_weather", "arguments": {"location": "San Francisco"}}`),
236+
Tokens: 1,
237+
}); err != nil {
238+
return err
239+
}
240+
return nil
241+
}
242+
202243
var toStream string
203244
toolName := mockToolNameFromRequest(in)
204245
if toolName != "" && !promptHasToolResults(in.Prompt) {

tests/e2e/mock_backend_test.go

Lines changed: 134 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -463,4 +463,138 @@ var _ = Describe("Mock Backend E2E Tests", Label("MockBackend"), func() {
463463
})
464464
})
465465
})
466+
467+
// Tests for duplicate tool call emissions during streaming.
468+
// The Go-side incremental JSON parser was emitting the same tool call on
469+
// every streaming token, and the post-streaming default: case re-emitted
470+
// all tool calls again, producing massive duplication.
471+
Describe("Streaming Tool Call Deduplication", Label("ToolDedup"), func() {
472+
// Helper: parse SSE lines and count tool call name/arguments chunks
473+
parseToolCallChunks := func(data []byte) (nameChunks int, argChunks int) {
474+
for _, line := range strings.Split(string(data), "\n") {
475+
line = strings.TrimSpace(line)
476+
if !strings.HasPrefix(line, "data: ") || line == "data: [DONE]" {
477+
continue
478+
}
479+
var chunk map[string]any
480+
if err := json.Unmarshal([]byte(strings.TrimPrefix(line, "data: ")), &chunk); err != nil {
481+
continue
482+
}
483+
choices, _ := chunk["choices"].([]any)
484+
if len(choices) == 0 {
485+
continue
486+
}
487+
delta, _ := choices[0].(map[string]any)["delta"].(map[string]any)
488+
if delta == nil {
489+
continue
490+
}
491+
toolCalls, _ := delta["tool_calls"].([]any)
492+
for _, tc := range toolCalls {
493+
tcMap, _ := tc.(map[string]any)
494+
fn, _ := tcMap["function"].(map[string]any)
495+
if fn == nil {
496+
continue
497+
}
498+
if name, _ := fn["name"].(string); name != "" {
499+
nameChunks++
500+
}
501+
if args, _ := fn["arguments"].(string); args != "" {
502+
argChunks++
503+
}
504+
}
505+
}
506+
return
507+
}
508+
509+
Context("Single tool call via Go-side JSON parser", func() {
510+
It("should emit exactly one tool call name without duplicates", func() {
511+
body := `{
512+
"model": "mock-model-autoparser",
513+
"messages": [{"role": "user", "content": "SINGLE_TOOL_CALL"}],
514+
"tools": [{"type": "function", "function": {"name": "get_weather", "description": "Get weather", "parameters": {"type": "object", "properties": {"location": {"type": "string"}}, "required": ["location"]}}}],
515+
"stream": true
516+
}`
517+
req, err := http.NewRequest("POST", apiURL+"/chat/completions", strings.NewReader(body))
518+
Expect(err).ToNot(HaveOccurred())
519+
req.Header.Set("Content-Type", "application/json")
520+
521+
httpClient := &http.Client{Timeout: 60 * time.Second}
522+
resp, err := httpClient.Do(req)
523+
Expect(err).ToNot(HaveOccurred())
524+
defer resp.Body.Close()
525+
Expect(resp.StatusCode).To(Equal(200))
526+
527+
data, err := io.ReadAll(resp.Body)
528+
Expect(err).ToNot(HaveOccurred())
529+
530+
nameChunks, argChunks := parseToolCallChunks(data)
531+
532+
Expect(nameChunks).To(Equal(1),
533+
"Expected exactly 1 tool call name chunk, got %d. Full SSE:\n%s",
534+
nameChunks, string(data))
535+
Expect(argChunks).To(BeNumerically(">=", 1),
536+
"Expected at least 1 arguments chunk. Full SSE:\n%s", string(data))
537+
})
538+
})
539+
540+
Context("ChatDelta tool calls (regression guard)", func() {
541+
It("should emit exactly one tool call name per tool", func() {
542+
body := `{
543+
"model": "mock-model-autoparser",
544+
"messages": [{"role": "user", "content": "AUTOPARSER_TOOL_CALL"}],
545+
"tools": [{"type": "function", "function": {"name": "search_collections", "description": "Search documents", "parameters": {"type": "object", "properties": {"query": {"type": "string"}}, "required": ["query"]}}}],
546+
"stream": true
547+
}`
548+
req, err := http.NewRequest("POST", apiURL+"/chat/completions", strings.NewReader(body))
549+
Expect(err).ToNot(HaveOccurred())
550+
req.Header.Set("Content-Type", "application/json")
551+
552+
httpClient := &http.Client{Timeout: 60 * time.Second}
553+
resp, err := httpClient.Do(req)
554+
Expect(err).ToNot(HaveOccurred())
555+
defer resp.Body.Close()
556+
Expect(resp.StatusCode).To(Equal(200))
557+
558+
data, err := io.ReadAll(resp.Body)
559+
Expect(err).ToNot(HaveOccurred())
560+
561+
nameChunks, _ := parseToolCallChunks(data)
562+
563+
Expect(nameChunks).To(Equal(1),
564+
"Expected exactly 1 tool call name chunk from ChatDeltas, got %d. Full SSE:\n%s",
565+
nameChunks, string(data))
566+
})
567+
})
568+
569+
Context("Multiple tool calls via Go-side JSON parser", func() {
570+
It("should emit exactly two tool call names without duplicates", func() {
571+
body := `{
572+
"model": "mock-model-autoparser",
573+
"messages": [{"role": "user", "content": "MULTI_TOOL_CALL"}],
574+
"tools": [{"type": "function", "function": {"name": "get_weather", "description": "Get weather", "parameters": {"type": "object", "properties": {"location": {"type": "string"}}, "required": ["location"]}}}],
575+
"stream": true
576+
}`
577+
req, err := http.NewRequest("POST", apiURL+"/chat/completions", strings.NewReader(body))
578+
Expect(err).ToNot(HaveOccurred())
579+
req.Header.Set("Content-Type", "application/json")
580+
581+
httpClient := &http.Client{Timeout: 60 * time.Second}
582+
resp, err := httpClient.Do(req)
583+
Expect(err).ToNot(HaveOccurred())
584+
defer resp.Body.Close()
585+
Expect(resp.StatusCode).To(Equal(200))
586+
587+
data, err := io.ReadAll(resp.Body)
588+
Expect(err).ToNot(HaveOccurred())
589+
590+
nameChunks, argChunks := parseToolCallChunks(data)
591+
592+
Expect(nameChunks).To(Equal(2),
593+
"Expected exactly 2 tool call name chunks (one per tool), got %d. Full SSE:\n%s",
594+
nameChunks, string(data))
595+
Expect(argChunks).To(BeNumerically(">=", 2),
596+
"Expected at least 2 arguments chunks. Full SSE:\n%s", string(data))
597+
})
598+
})
599+
})
466600
})

0 commit comments

Comments
 (0)