Skip to content

Commit 5529931

Browse files
AlemTuzlakjpr5
AlemTuzlak
authored and
jpr5
committed
feat: preserve both content and toolCalls in stream collapse functions
1 parent 07b1a12 commit 5529931

3 files changed

Lines changed: 146 additions & 24 deletions

File tree

‎src/__tests__/content-with-toolcalls.test.ts‎

Lines changed: 122 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -425,3 +425,125 @@ describe("Gemini — content + toolCalls", () => {
425425
expect(body.candidates[0].finishReason).toBe("FUNCTION_CALL");
426426
});
427427
});
428+
429+
import {
430+
collapseOpenAISSE,
431+
collapseAnthropicSSE,
432+
collapseGeminiSSE,
433+
collapseOllamaNDJSON,
434+
} from "../stream-collapse.js";
435+
436+
describe("stream-collapse — content + toolCalls coexistence", () => {
437+
it("OpenAI: preserves both content and toolCalls", () => {
438+
const body = [
439+
`data: ${JSON.stringify({ id: "c1", choices: [{ delta: { role: "assistant" } }] })}`,
440+
"",
441+
`data: ${JSON.stringify({ id: "c1", choices: [{ delta: { content: "Hello" } }] })}`,
442+
"",
443+
`data: ${JSON.stringify({
444+
id: "c1",
445+
choices: [
446+
{
447+
delta: {
448+
tool_calls: [
449+
{
450+
index: 0,
451+
id: "call_abc",
452+
type: "function",
453+
function: { name: "get_weather", arguments: '{"city":"NYC"}' },
454+
},
455+
],
456+
},
457+
},
458+
],
459+
})}`,
460+
"",
461+
"data: [DONE]",
462+
"",
463+
].join("\n");
464+
465+
const result = collapseOpenAISSE(body);
466+
expect(result.content).toBe("Hello");
467+
expect(result.toolCalls).toHaveLength(1);
468+
expect(result.toolCalls![0].name).toBe("get_weather");
469+
});
470+
471+
it("Anthropic: preserves both content and toolCalls", () => {
472+
const body = [
473+
`event: message_start\ndata: ${JSON.stringify({ type: "message_start", message: {} })}`,
474+
"",
475+
`event: content_block_start\ndata: ${JSON.stringify({ type: "content_block_start", index: 0, content_block: { type: "text", text: "" } })}`,
476+
"",
477+
`event: content_block_delta\ndata: ${JSON.stringify({ type: "content_block_delta", index: 0, delta: { type: "text_delta", text: "Hello" } })}`,
478+
"",
479+
`event: content_block_stop\ndata: ${JSON.stringify({ type: "content_block_stop", index: 0 })}`,
480+
"",
481+
`event: content_block_start\ndata: ${JSON.stringify({ type: "content_block_start", index: 1, content_block: { type: "tool_use", id: "toolu_abc", name: "get_weather", input: {} } })}`,
482+
"",
483+
`event: content_block_delta\ndata: ${JSON.stringify({ type: "content_block_delta", index: 1, delta: { type: "input_json_delta", partial_json: '{"city":"NYC"}' } })}`,
484+
"",
485+
`event: content_block_stop\ndata: ${JSON.stringify({ type: "content_block_stop", index: 1 })}`,
486+
"",
487+
`event: message_delta\ndata: ${JSON.stringify({ type: "message_delta", delta: { stop_reason: "tool_use" } })}`,
488+
"",
489+
`event: message_stop\ndata: ${JSON.stringify({ type: "message_stop" })}`,
490+
"",
491+
].join("\n");
492+
493+
const result = collapseAnthropicSSE(body);
494+
expect(result.content).toBe("Hello");
495+
expect(result.toolCalls).toHaveLength(1);
496+
expect(result.toolCalls![0].name).toBe("get_weather");
497+
});
498+
499+
it("Gemini: preserves both content and toolCalls", () => {
500+
const body = [
501+
`data: ${JSON.stringify({
502+
candidates: [
503+
{ content: { role: "model", parts: [{ text: "Hello" }] }, index: 0 },
504+
],
505+
})}`,
506+
"",
507+
`data: ${JSON.stringify({
508+
candidates: [
509+
{
510+
content: {
511+
role: "model",
512+
parts: [
513+
{ functionCall: { name: "get_weather", args: { city: "NYC" } } },
514+
],
515+
},
516+
finishReason: "FUNCTION_CALL",
517+
index: 0,
518+
},
519+
],
520+
})}`,
521+
"",
522+
].join("\n");
523+
524+
const result = collapseGeminiSSE(body);
525+
expect(result.content).toBe("Hello");
526+
expect(result.toolCalls).toHaveLength(1);
527+
expect(result.toolCalls![0].name).toBe("get_weather");
528+
});
529+
530+
it("Ollama: preserves both content and toolCalls", () => {
531+
const body = [
532+
JSON.stringify({
533+
model: "llama3",
534+
message: {
535+
role: "assistant",
536+
content: "Hello",
537+
tool_calls: [{ function: { name: "get_weather", arguments: { city: "NYC" } } }],
538+
},
539+
done: false,
540+
}),
541+
JSON.stringify({ model: "llama3", message: { role: "assistant", content: "" }, done: true }),
542+
].join("\n");
543+
544+
const result = collapseOllamaNDJSON(body);
545+
expect(result.content).toBe("Hello");
546+
expect(result.toolCalls).toHaveLength(1);
547+
expect(result.toolCalls![0].name).toBe("get_weather");
548+
});
549+
});

‎src/__tests__/stream-collapse.test.ts‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1568,7 +1568,7 @@ describe("collapseOllamaNDJSON with tool_calls", () => {
15681568
expect(result.content).toBeUndefined();
15691569
});
15701570

1571-
it("returns toolCalls (not content) when both tool_calls and text are present", () => {
1571+
it("preserves both content and toolCalls when both tool_calls and text are present", () => {
15721572
const body = [
15731573
JSON.stringify({
15741574
model: "llama3",
@@ -1594,11 +1594,11 @@ describe("collapseOllamaNDJSON with tool_calls", () => {
15941594
].join("\n");
15951595

15961596
const result = collapseOllamaNDJSON(body);
1597-
// When toolCalls are present, they take priority over content
1597+
// When toolCalls are present alongside content, both are preserved
15981598
expect(result.toolCalls).toBeDefined();
15991599
expect(result.toolCalls).toHaveLength(1);
16001600
expect(result.toolCalls![0].name).toBe("get_weather");
1601-
expect(result.content).toBeUndefined();
1601+
expect(result.content).toBe("Let me check the weather.");
16021602
});
16031603

16041604
it("extracts multiple tool_calls across chunks", () => {

‎src/stream-collapse.ts‎

Lines changed: 21 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -140,6 +140,7 @@ export function collapseOpenAISSE(body: string): CollapseResult {
140140
if (toolCallMap.size > 0) {
141141
const sorted = Array.from(toolCallMap.entries()).sort(([a], [b]) => a - b);
142142
return {
143+
...(content ? { content } : {}),
143144
toolCalls: sorted.map(([, tc]) => ({
144145
name: tc.name,
145146
arguments: tc.arguments,
@@ -229,6 +230,7 @@ export function collapseAnthropicSSE(body: string): CollapseResult {
229230
if (toolCallMap.size > 0) {
230231
const sorted = Array.from(toolCallMap.entries()).sort(([a], [b]) => a - b);
231232
return {
233+
...(content ? { content } : {}),
232234
toolCalls: sorted.map(([, tc]) => ({
233235
name: tc.name,
234236
arguments: tc.arguments,
@@ -260,6 +262,7 @@ export function collapseGeminiSSE(body: string): CollapseResult {
260262
const lines = body.split("\n\n").filter((l) => l.trim().length > 0);
261263
let content = "";
262264
let droppedChunks = 0;
265+
const toolCalls: ToolCall[] = [];
263266

264267
for (const line of lines) {
265268
const dataLine = line.split("\n").find((l) => l.startsWith("data:"));
@@ -284,32 +287,25 @@ export function collapseGeminiSSE(body: string): CollapseResult {
284287
const parts = candidateContent.parts as Array<Record<string, unknown>> | undefined;
285288
if (!parts || parts.length === 0) continue;
286289

287-
// Handle functionCall parts
288-
const fnCallParts = parts.filter((p) => p.functionCall);
289-
if (fnCallParts.length > 0) {
290-
const toolCallMap = new Map<number, { name: string; arguments: string }>();
291-
for (let i = 0; i < fnCallParts.length; i++) {
292-
const fc = fnCallParts[i].functionCall as Record<string, unknown>;
293-
toolCallMap.set(i, {
290+
for (const part of parts) {
291+
if (part.functionCall) {
292+
const fc = part.functionCall as Record<string, unknown>;
293+
toolCalls.push({
294294
name: String(fc.name ?? ""),
295295
arguments: typeof fc.args === "string" ? (fc.args as string) : JSON.stringify(fc.args),
296296
});
297-
}
298-
if (toolCallMap.size > 0) {
299-
const sorted = Array.from(toolCallMap.entries()).sort(([a], [b]) => a - b);
300-
return {
301-
toolCalls: sorted.map(([, tc]) => ({
302-
name: tc.name,
303-
arguments: tc.arguments,
304-
})),
305-
...(droppedChunks > 0 ? { droppedChunks } : {}),
306-
};
297+
} else if (typeof part.text === "string") {
298+
content += part.text;
307299
}
308300
}
301+
}
309302

310-
if (typeof parts[0].text === "string") {
311-
content += parts[0].text;
312-
}
303+
if (toolCalls.length > 0) {
304+
return {
305+
...(content ? { content } : {}),
306+
toolCalls,
307+
...(droppedChunks > 0 ? { droppedChunks } : {}),
308+
};
313309
}
314310

315311
return { content, ...(droppedChunks > 0 ? { droppedChunks } : {}) };
@@ -372,7 +368,11 @@ export function collapseOllamaNDJSON(body: string): CollapseResult {
372368
}
373369

374370
if (toolCalls.length > 0) {
375-
return { toolCalls, ...(droppedChunks > 0 ? { droppedChunks } : {}) };
371+
return {
372+
...(content ? { content } : {}),
373+
toolCalls,
374+
...(droppedChunks > 0 ? { droppedChunks } : {}),
375+
};
376376
}
377377

378378
return { content, ...(droppedChunks > 0 ? { droppedChunks } : {}) };

0 commit comments

Comments
 (0)