shumengya/sproutclaw
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge pull request #4247 from badlogic/separate-accumulators

Browse Source 
fix(ai): handle mixed chat completion deltas
This commit is contained in:
Mario Zechner
2026-05-07 11:37:28 +02:00
committed by GitHub
parent 6d4d2e9288 6b271842e2
commit fa35c5fabe
2 changed files with 367 additions and 121 deletions
Download Patch File Download Diff File Expand all files Collapse all files

256
packages/ai/src/providers/openai-completions.ts
View File

@@ -160,45 +160,104 @@ export const streamOpenAICompletions: StreamFunction<"openai-completions", OpenA
partialArgs?: string;
streamIndex?: number;
}
type StreamingBlock = TextContent | ThinkingContent | StreamingToolCallBlock;
type StreamingToolCallDelta = NonNullable<ChatCompletionChunk.Choice.Delta["tool_calls"]>[number];
let currentBlock: TextContent | ThinkingContent | StreamingToolCallBlock | null = null;
const blocks = output.content;
const getContentIndex = (block: typeof currentBlock) => (block ? blocks.indexOf(block) : -1);
const currentContentIndex = () => getContentIndex(currentBlock);
const finishCurrentBlock = (block?: typeof currentBlock) => {
if (block) {
const contentIndex = getContentIndex(block);
if (contentIndex === -1) {
return;
}
if (block.type === "text") {
stream.push({
type: "text_end",
contentIndex,
content: block.text,
partial: output,
});
} else if (block.type === "thinking") {
stream.push({
type: "thinking_end",
contentIndex,
content: block.thinking,
partial: output,
});
} else if (block.type === "toolCall") {
block.arguments = parseStreamingJson(block.partialArgs);
// Finalize in-place and strip the scratch buffers so replay only
// carries parsed arguments.
delete block.partialArgs;
delete block.streamIndex;
stream.push({
type: "toolcall_end",
contentIndex,
toolCall: block,
partial: output,
});
}
let textBlock: TextContent | null = null;
let thinkingBlock: ThinkingContent | null = null;
const toolCallBlocksByIndex = new Map<number, StreamingToolCallBlock>();
const toolCallBlocksById = new Map<string, StreamingToolCallBlock>();
const blocks = output.content as StreamingBlock[];
const getContentIndex = (block: StreamingBlock) => blocks.indexOf(block);
const finishBlock = (block: StreamingBlock) => {
const contentIndex = getContentIndex(block);
if (contentIndex === -1) {
return;
}
if (block.type === "text") {
stream.push({
type: "text_end",
contentIndex,
content: block.text,
partial: output,
});
} else if (block.type === "thinking") {
stream.push({
type: "thinking_end",
contentIndex,
content: block.thinking,
partial: output,
});
} else if (block.type === "toolCall") {
block.arguments = parseStreamingJson(block.partialArgs);
// Finalize in-place and strip the scratch buffers so replay only
// carries parsed arguments.
delete block.partialArgs;
delete block.streamIndex;
stream.push({
type: "toolcall_end",
contentIndex,
toolCall: block,
partial: output,
});
}
};
const ensureTextBlock = () => {
if (!textBlock) {
textBlock = { type: "text", text: "" };
blocks.push(textBlock);
stream.push({ type: "text_start", contentIndex: getContentIndex(textBlock), partial: output });
}
return textBlock;
};
const ensureThinkingBlock = (thinkingSignature: string) => {
if (!thinkingBlock) {
thinkingBlock = {
type: "thinking",
thinking: "",
thinkingSignature,
};
blocks.push(thinkingBlock);
stream.push({ type: "thinking_start", contentIndex: getContentIndex(thinkingBlock), partial: output });
}
return thinkingBlock;
};
const ensureToolCallBlock = (toolCall: StreamingToolCallDelta) => {
const streamIndex = typeof toolCall.index === "number" ? toolCall.index : undefined;
let block = streamIndex !== undefined ? toolCallBlocksByIndex.get(streamIndex) : undefined;
if (!block && toolCall.id) {
block = toolCallBlocksById.get(toolCall.id);
}
if (!block) {
block = {
type: "toolCall",
id: toolCall.id || "",
name: toolCall.function?.name || "",
arguments: {},
partialArgs: "",
streamIndex,
};
if (streamIndex !== undefined) {
toolCallBlocksByIndex.set(streamIndex, block);
}
if (toolCall.id) {
toolCallBlocksById.set(toolCall.id, block);
}
blocks.push(block);
stream.push({
type: "toolcall_start",
contentIndex: getContentIndex(block),
partial: output,
});
}
if (streamIndex !== undefined && block.streamIndex === undefined) {
block.streamIndex = streamIndex;
toolCallBlocksByIndex.set(streamIndex, block);
}
if (toolCall.id) {
toolCallBlocksById.set(toolCall.id, block);
}
return block;
};
for await (const chunk of openaiStream) {
@@ -237,22 +296,14 @@ export const streamOpenAICompletions: StreamFunction<"openai-completions", OpenA
choice.delta.content !== undefined &&
choice.delta.content.length > 0
) {
if (!currentBlock || currentBlock.type !== "text") {
finishCurrentBlock(currentBlock);
currentBlock = { type: "text", text: "" };
output.content.push(currentBlock);
stream.push({ type: "text_start", contentIndex: currentContentIndex(), partial: output });
}
if (currentBlock.type === "text") {
currentBlock.text += choice.delta.content;
stream.push({
type: "text_delta",
contentIndex: currentContentIndex(),
delta: choice.delta.content,
partial: output,
});
}
const block = ensureTextBlock();
block.text += choice.delta.content;
stream.push({
type: "text_delta",
contentIndex: getContentIndex(block),
delta: choice.delta.content,
partial: output,
});
}
// Some endpoints return reasoning in reasoning_content (llama.cpp),
@@ -260,38 +311,24 @@ export const streamOpenAICompletions: StreamFunction<"openai-completions", OpenA
// Use the first non-empty reasoning field to avoid duplication
// (e.g., chutes.ai returns both reasoning_content and reasoning with same content)
const reasoningFields = ["reasoning_content", "reasoning", "reasoning_text"];
const deltaFields = choice.delta as Record<string, unknown>;
let foundReasoningField: string | null = null;
for (const field of reasoningFields) {
if (
(choice.delta as any)[field] !== null &&
(choice.delta as any)[field] !== undefined &&
(choice.delta as any)[field].length > 0
) {
if (!foundReasoningField) {
foundReasoningField = field;
break;
}
const value = deltaFields[field];
if (typeof value === "string" && value.length > 0) {
foundReasoningField = field;
break;
}
}
if (foundReasoningField) {
if (!currentBlock || currentBlock.type !== "thinking") {
finishCurrentBlock(currentBlock);
currentBlock = {
type: "thinking",
thinking: "",
thinkingSignature: foundReasoningField,
};
output.content.push(currentBlock);
stream.push({ type: "thinking_start", contentIndex: currentContentIndex(), partial: output });
}
if (currentBlock.type === "thinking") {
const delta = (choice.delta as any)[foundReasoningField];
currentBlock.thinking += delta;
const delta = deltaFields[foundReasoningField];
if (typeof delta === "string" && delta.length > 0) {
const block = ensureThinkingBlock(foundReasoningField);
block.thinking += delta;
stream.push({
type: "thinking_delta",
contentIndex: currentContentIndex(),
contentIndex: getContentIndex(block),
delta,
partial: output,
});
@@ -300,52 +337,27 @@ export const streamOpenAICompletions: StreamFunction<"openai-completions", OpenA
if (choice?.delta?.tool_calls) {
for (const toolCall of choice.delta.tool_calls) {
const streamIndex = typeof toolCall.index === "number" ? toolCall.index : undefined;
const sameToolCall =
currentBlock?.type === "toolCall" &&
((streamIndex !== undefined && currentBlock.streamIndex === streamIndex) ||
(streamIndex === undefined && toolCall.id && currentBlock.id === toolCall.id));
if (!sameToolCall) {
finishCurrentBlock(currentBlock);
currentBlock = {
type: "toolCall",
id: toolCall.id || "",
name: toolCall.function?.name || "",
arguments: {},
partialArgs: "",
streamIndex,
};
output.content.push(currentBlock);
stream.push({
type: "toolcall_start",
contentIndex: getContentIndex(currentBlock),
partial: output,
});
const block = ensureToolCallBlock(toolCall);
if (!block.id && toolCall.id) {
block.id = toolCall.id;
toolCallBlocksById.set(toolCall.id, block);
}
if (!block.name && toolCall.function?.name) {
block.name = toolCall.function.name;
}
const currentToolCallBlock = currentBlock?.type === "toolCall" ? currentBlock : null;
if (currentToolCallBlock) {
if (!currentToolCallBlock.id && toolCall.id) currentToolCallBlock.id = toolCall.id;
if (!currentToolCallBlock.name && toolCall.function?.name) {
currentToolCallBlock.name = toolCall.function.name;
}
if (currentToolCallBlock.streamIndex === undefined && streamIndex !== undefined) {
currentToolCallBlock.streamIndex = streamIndex;
}
let delta = "";
if (toolCall.function?.arguments) {
delta = toolCall.function.arguments;
currentToolCallBlock.partialArgs += toolCall.function.arguments;
currentToolCallBlock.arguments = parseStreamingJson(currentToolCallBlock.partialArgs);
}
stream.push({
type: "toolcall_delta",
contentIndex: getContentIndex(currentToolCallBlock),
delta,
partial: output,
});
let delta = "";
if (toolCall.function?.arguments) {
delta = toolCall.function.arguments;
block.partialArgs = (block.partialArgs ?? "") + toolCall.function.arguments;
block.arguments = parseStreamingJson(block.partialArgs);
}
stream.push({
type: "toolcall_delta",
contentIndex: getContentIndex(block),
delta,
partial: output,
});
}
}
@@ -365,7 +377,9 @@ export const streamOpenAICompletions: StreamFunction<"openai-completions", OpenA
}
}
finishCurrentBlock(currentBlock);
for (const block of blocks) {
finishBlock(block);
}
if (options?.signal?.aborted) {
throw new Error("Request was aborted");
}

232
packages/ai/test/openai-completions-tool-choice.test.ts
View File

@@ -552,6 +552,238 @@ describe("openai-completions tool_choice", () => {
expect(toolCall).not.toHaveProperty("partialArgs");
});
it("accumulates mixed content, reasoning, and parallel tool call deltas independently", async () => {
mockState.chunks = [
{
id: "chatcmpl-mixed-deltas",
choices: [
{
delta: {
content: "answer 1",
reasoning_content: "think 1",
tool_calls: [
{
index: 0,
id: "tc_read_initial",
type: "function",
function: { name: "read", arguments: '{"path":"README' },
},
{
index: 1,
id: "tc_grep_initial",
type: "function",
function: { name: "grep", arguments: '{"pattern":"TODO' },
},
{
id: "tc_list_no_index",
type: "function",
function: { name: "list", arguments: '{"path":"packages' },
},
{
id: "tc_write_no_index",
type: "function",
function: { name: "write", arguments: '{"path":"out' },
},
],
},
finish_reason: null,
},
],
},
{
id: "chatcmpl-mixed-deltas",
choices: [
{
delta: {
content: " answer 2",
tool_calls: [
{
index: 1,
id: "tc_grep_changed",
type: "function",
function: { arguments: '","path":"src' },
},
{
id: "tc_write_no_index",
type: "function",
function: { arguments: '.txt","content":"ok"}' },
},
{
id: "tc_list_no_index",
type: "function",
function: { arguments: '/ai"}' },
},
],
},
finish_reason: null,
},
],
},
{
id: "chatcmpl-mixed-deltas",
choices: [
{
delta: {
content: "\n",
reasoning_content: " think 2",
tool_calls: [
{
index: 0,
id: "tc_read_changed",
type: "function",
function: { arguments: '.md"}' },
},
{
index: 1,
type: "function",
function: { arguments: '"}' },
},
],
},
finish_reason: "tool_calls",
},
],
usage: {
prompt_tokens: 10,
completion_tokens: 8,
prompt_tokens_details: { cached_tokens: 0 },
completion_tokens_details: { reasoning_tokens: 2 },
},
},
];
const { compat: _compat, ...baseModel } = getModel("openai", "gpt-4o-mini")!;
const model = { ...baseModel, api: "openai-completions" } as const;
const tools: Tool[] = [
{
name: "read",
description: "Read a file",
parameters: Type.Object({ path: Type.String() }),
},
{
name: "grep",
description: "Search a file",
parameters: Type.Object({ pattern: Type.String(), path: Type.String() }),
},
{
name: "list",
description: "List a directory",
parameters: Type.Object({ path: Type.String() }),
},
{
name: "write",
description: "Write a file",
parameters: Type.Object({ path: Type.String(), content: Type.String() }),
},
];
const s = streamSimple(
model,
{
messages: [
{
role: "user",
content: "Think, answer, and use tools.",
timestamp: Date.now(),
},
],
tools,
},
{ apiKey: "test" },
);
const eventTypes: string[] = [];
const toolEventsByContentIndex = new Map<number, string[]>();
for await (const event of s) {
eventTypes.push(event.type);
if (event.type === "toolcall_start" || event.type === "toolcall_delta" || event.type === "toolcall_end") {
const events = toolEventsByContentIndex.get(event.contentIndex) ?? [];
events.push(event.type);
toolEventsByContentIndex.set(event.contentIndex, events);
}
}
const response = await s.result();
expect(response.stopReason).toBe("toolUse");
expect(eventTypes.filter((type) => type === "text_start")).toHaveLength(1);
expect(eventTypes.filter((type) => type === "text_delta")).toHaveLength(3);
expect(eventTypes.filter((type) => type === "text_end")).toHaveLength(1);
expect(eventTypes.filter((type) => type === "thinking_start")).toHaveLength(1);
expect(eventTypes.filter((type) => type === "thinking_delta")).toHaveLength(2);
expect(eventTypes.filter((type) => type === "thinking_end")).toHaveLength(1);
expect(eventTypes.filter((type) => type === "toolcall_start")).toHaveLength(4);
expect(eventTypes.filter((type) => type === "toolcall_delta")).toHaveLength(9);
expect(eventTypes.filter((type) => type === "toolcall_end")).toHaveLength(4);
expect(toolEventsByContentIndex.get(2)).toEqual([
"toolcall_start",
"toolcall_delta",
"toolcall_delta",
"toolcall_end",
]);
expect(toolEventsByContentIndex.get(3)).toEqual([
"toolcall_start",
"toolcall_delta",
"toolcall_delta",
"toolcall_delta",
"toolcall_end",
]);
expect(toolEventsByContentIndex.get(4)).toEqual([
"toolcall_start",
"toolcall_delta",
"toolcall_delta",
"toolcall_end",
]);
expect(toolEventsByContentIndex.get(5)).toEqual([
"toolcall_start",
"toolcall_delta",
"toolcall_delta",
"toolcall_end",
]);
expect(response.content).toHaveLength(6);
expect(response.content[0]).toEqual({ type: "text", text: "answer 1 answer 2\n" });
expect(response.content[1]).toEqual({
type: "thinking",
thinking: "think 1 think 2",
thinkingSignature: "reasoning_content",
});
const readCall = response.content[2];
const grepCall = response.content[3];
const listCall = response.content[4];
const writeCall = response.content[5];
expect(readCall.type).toBe("toolCall");
expect(grepCall.type).toBe("toolCall");
expect(listCall.type).toBe("toolCall");
expect(writeCall.type).toBe("toolCall");
if (
readCall.type !== "toolCall" ||
grepCall.type !== "toolCall" ||
listCall.type !== "toolCall" ||
writeCall.type !== "toolCall"
) {
throw new Error("Expected toolCall content");
}
expect(readCall.id).toBe("tc_read_initial");
expect(readCall.name).toBe("read");
expect(readCall.arguments).toEqual({ path: "README.md" });
expect(readCall).not.toHaveProperty("streamIndex");
expect(readCall).not.toHaveProperty("partialArgs");
expect(grepCall.id).toBe("tc_grep_initial");
expect(grepCall.name).toBe("grep");
expect(grepCall.arguments).toEqual({ pattern: "TODO", path: "src" });
expect(grepCall).not.toHaveProperty("streamIndex");
expect(grepCall).not.toHaveProperty("partialArgs");
expect(listCall.id).toBe("tc_list_no_index");
expect(listCall.name).toBe("list");
expect(listCall.arguments).toEqual({ path: "packages/ai" });
expect(listCall).not.toHaveProperty("streamIndex");
expect(listCall).not.toHaveProperty("partialArgs");
expect(writeCall.id).toBe("tc_write_no_index");
expect(writeCall.name).toBe("write");
expect(writeCall.arguments).toEqual({ path: "out.txt", content: "ok" });
expect(writeCall).not.toHaveProperty("streamIndex");
expect(writeCall).not.toHaveProperty("partialArgs");
});
it("does not double-count reasoning tokens in completion usage", async () => {
mockState.chunks = [
{