Skip to content

Commit bd15549

Browse files
pyrytakalascidomino
authored andcommitted
Fix: Process all parts in response chunks when thought is first (#13539)
1 parent b1f7a7e commit bd15549

File tree

3 files changed

+113
-10
lines changed

3 files changed

+113
-10
lines changed

packages/core/src/core/turn.test.ts

Lines changed: 102 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -754,6 +754,108 @@ describe('Turn', () => {
754754

755755
expect(events).toEqual([expectedEvent]);
756756
});
757+
758+
it('should process all parts when thought is first part in chunk', async () => {
759+
const mockResponseStream = (async function* () {
760+
yield {
761+
type: StreamEventType.CHUNK,
762+
value: {
763+
candidates: [
764+
{
765+
content: {
766+
parts: [
767+
{ text: '**Planning** the solution', thought: 'planning' },
768+
{ text: 'I will help you with that.' },
769+
],
770+
},
771+
citationMetadata: {
772+
citations: [{ uri: 'https://example.com', title: 'Source' }],
773+
},
774+
finishReason: 'STOP',
775+
},
776+
],
777+
functionCalls: [
778+
{
779+
id: 'fc1',
780+
name: 'ReadFile',
781+
args: { path: 'file.txt' },
782+
},
783+
],
784+
responseId: 'trace-789',
785+
} as unknown as GenerateContentResponse,
786+
};
787+
})();
788+
mockSendMessageStream.mockResolvedValue(mockResponseStream);
789+
790+
const events = [];
791+
for await (const event of turn.run(
792+
{ model: 'gemini' },
793+
[{ text: 'Test mixed content' }],
794+
new AbortController().signal,
795+
)) {
796+
events.push(event);
797+
}
798+
799+
// Should yield:
800+
// 1. Thought event (from first part)
801+
// 2. Content event (from second part)
802+
// 3. ToolCallRequest event (from functionCalls)
803+
// 4. Citation event (from citationMetadata, emitted with finishReason)
804+
// 5. Finished event (from finishReason)
805+
806+
expect(events.length).toBe(5);
807+
808+
const thoughtEvent = events.find(
809+
(e) => e.type === GeminiEventType.Thought,
810+
);
811+
expect(thoughtEvent).toBeDefined();
812+
expect(thoughtEvent).toMatchObject({
813+
type: GeminiEventType.Thought,
814+
value: { subject: 'Planning', description: 'the solution' },
815+
traceId: 'trace-789',
816+
});
817+
818+
const contentEvent = events.find(
819+
(e) => e.type === GeminiEventType.Content,
820+
);
821+
expect(contentEvent).toBeDefined();
822+
expect(contentEvent).toMatchObject({
823+
type: GeminiEventType.Content,
824+
value: 'I will help you with that.',
825+
traceId: 'trace-789',
826+
});
827+
828+
const toolCallEvent = events.find(
829+
(e) => e.type === GeminiEventType.ToolCallRequest,
830+
);
831+
expect(toolCallEvent).toBeDefined();
832+
expect(toolCallEvent).toMatchObject({
833+
type: GeminiEventType.ToolCallRequest,
834+
value: expect.objectContaining({
835+
callId: 'fc1',
836+
name: 'ReadFile',
837+
args: { path: 'file.txt' },
838+
}),
839+
});
840+
841+
const citationEvent = events.find(
842+
(e) => e.type === GeminiEventType.Citation,
843+
);
844+
expect(citationEvent).toBeDefined();
845+
expect(citationEvent).toMatchObject({
846+
type: GeminiEventType.Citation,
847+
value: expect.stringContaining('https://example.com'),
848+
});
849+
850+
const finishedEvent = events.find(
851+
(e) => e.type === GeminiEventType.Finished,
852+
);
853+
expect(finishedEvent).toBeDefined();
854+
expect(finishedEvent).toMatchObject({
855+
type: GeminiEventType.Finished,
856+
value: { reason: 'STOP' },
857+
});
858+
});
757859
});
758860

759861
describe('getDebugResponses', () => {

packages/core/src/core/turn.ts

Lines changed: 10 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -290,15 +290,16 @@ export class Turn {
290290

291291
const traceId = resp.responseId;
292292

293-
const thoughtPart = resp.candidates?.[0]?.content?.parts?.[0];
294-
if (thoughtPart?.thought) {
295-
const thought = parseThought(thoughtPart.text ?? '');
296-
yield {
297-
type: GeminiEventType.Thought,
298-
value: thought,
299-
traceId,
300-
};
301-
continue;
293+
const parts = resp.candidates?.[0]?.content?.parts ?? [];
294+
for (const part of parts) {
295+
if (part.thought) {
296+
const thought = parseThought(part.text ?? '');
297+
yield {
298+
type: GeminiEventType.Thought,
299+
value: thought,
300+
traceId,
301+
};
302+
}
302303
}
303304

304305
const text = getResponseText(resp);

packages/core/src/utils/partUtils.ts

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -81,7 +81,7 @@ export function getResponseText(
8181
candidate.content.parts.length > 0
8282
) {
8383
return candidate.content.parts
84-
.filter((part) => part.text)
84+
.filter((part) => part.text && !part.thought)
8585
.map((part) => part.text)
8686
.join('');
8787
}

0 commit comments

Comments
 (0)