diff options
| -rw-r--r-- | packages/core/src/core/geminiChat.test.ts | 50 | ||||
| -rw-r--r-- | packages/core/src/core/geminiChat.ts | 47 |
2 files changed, 67 insertions, 30 deletions
diff --git a/packages/core/src/core/geminiChat.test.ts b/packages/core/src/core/geminiChat.test.ts index 11e222c9..8d434dd4 100644 --- a/packages/core/src/core/geminiChat.test.ts +++ b/packages/core/src/core/geminiChat.test.ts @@ -69,10 +69,7 @@ describe('GeminiChat', () => { expect(history.length).toBe(2); expect(history[0]).toEqual(userInput); expect(history[1].role).toBe('model'); - expect(history[1].parts).toEqual([ - { text: 'Model part 1' }, - { text: 'Model part 2' }, - ]); + expect(history[1].parts).toEqual([{ text: 'Model part 1Model part 2' }]); }); it('should handle a mix of user and model roles in outputContents (though unusual)', () => { @@ -101,11 +98,7 @@ describe('GeminiChat', () => { chat.recordHistory(userInput, modelOutputParts); const history = chat.getHistory(); expect(history.length).toBe(2); - expect(history[1].parts).toEqual([ - { text: 'M1' }, - { text: 'M2' }, - { text: 'M3' }, - ]); + expect(history[1].parts).toEqual([{ text: 'M1M2M3' }]); }); it('should not consolidate if roles are different between model outputs', () => { @@ -177,8 +170,7 @@ describe('GeminiChat', () => { expect(finalHistory[2]).toEqual(secondUserInput); expect(finalHistory[3].role).toBe('model'); expect(finalHistory[3].parts).toEqual([ - { text: 'Second model response part 1' }, - { text: 'Second model response part 2' }, + { text: 'Second model response part 1Second model response part 2' }, ]); }); @@ -218,8 +210,7 @@ describe('GeminiChat', () => { expect(history[2]).toEqual(currentUserInput); expect(history[3].role).toBe('model'); expect(history[3].parts).toEqual([ - { text: 'Part A of new answer.' }, - { text: 'Part B of new answer.' }, + { text: 'Part A of new answer.Part B of new answer.' }, ]); }); @@ -235,6 +226,31 @@ describe('GeminiChat', () => { expect(history[1].parts).toEqual([]); }); + it('should handle aggregating modelOutput', () => { + const modelOutputUndefinedParts: Content[] = [ + { role: 'model', parts: [{ text: 'First model part' }] }, + { role: 'model', parts: [{ text: 'Second model part' }] }, + { role: 'model', parts: undefined as unknown as Part[] }, // Test undefined parts + { role: 'model', parts: [{ text: 'Third model part' }] }, + { role: 'model', parts: [] }, // Test empty parts array + ]; + // @ts-expect-error Accessing private method for testing purposes + chat.recordHistory(userInput, modelOutputUndefinedParts); + const history = chat.getHistory(); + expect(history.length).toBe(5); + expect(history[0]).toEqual(userInput); + expect(history[1].role).toBe('model'); + expect(history[1].parts).toEqual([ + { text: 'First model partSecond model part' }, + ]); + expect(history[2].role).toBe('model'); + expect(history[2].parts).toBeUndefined(); + expect(history[3].role).toBe('model'); + expect(history[3].parts).toEqual([{ text: 'Third model part' }]); + expect(history[4].role).toBe('model'); + expect(history[4].parts).toEqual([]); + }); + it('should handle modelOutput with parts being undefined or empty (if they pass initial every check)', () => { const modelOutputUndefinedParts: Content[] = [ { role: 'model', parts: [{ text: 'Text part' }] }, @@ -244,10 +260,14 @@ describe('GeminiChat', () => { // @ts-expect-error Accessing private method for testing purposes chat.recordHistory(userInput, modelOutputUndefinedParts); const history = chat.getHistory(); - expect(history.length).toBe(2); + expect(history.length).toBe(4); // userInput, model1 (text), model2 (undefined parts), model3 (empty parts) + expect(history[0]).toEqual(userInput); expect(history[1].role).toBe('model'); - // The consolidation logic should handle undefined/empty parts by spreading `|| []` expect(history[1].parts).toEqual([{ text: 'Text part' }]); + expect(history[2].role).toBe('model'); + expect(history[2].parts).toBeUndefined(); + expect(history[3].role).toBe('model'); + expect(history[3].parts).toEqual([]); }); it('should correctly handle automaticFunctionCallingHistory', () => { diff --git a/packages/core/src/core/geminiChat.ts b/packages/core/src/core/geminiChat.ts index b34b6f35..25ea6e52 100644 --- a/packages/core/src/core/geminiChat.ts +++ b/packages/core/src/core/geminiChat.ts @@ -15,6 +15,7 @@ import { SendMessageParameters, GoogleGenAI, createUserContent, + Part, } from '@google/genai'; import { retryWithBackoff } from '../utils/retry.js'; import { isFunctionResponse } from '../utils/messageInspectors.js'; @@ -343,13 +344,13 @@ export class GeminiChat { for (const content of outputContents) { const lastContent = consolidatedOutputContents[consolidatedOutputContents.length - 1]; - if ( - lastContent && - lastContent.role === 'model' && - content.role === 'model' && - lastContent.parts - ) { - lastContent.parts.push(...(content.parts || [])); + if (this.isTextContent(lastContent) && this.isTextContent(content)) { + // If both current and last are text, combine their text into the lastContent's first part + // and append any other parts from the current content. + lastContent.parts[0].text += content.parts[0].text || ''; + if (content.parts.length > 1) { + lastContent.parts.push(...content.parts.slice(1)); + } } else { consolidatedOutputContents.push(content); } @@ -357,24 +358,40 @@ export class GeminiChat { if (consolidatedOutputContents.length > 0) { const lastHistoryEntry = this.history[this.history.length - 1]; - // Only merge if AFC history was NOT just added, to prevent merging with last AFC model turn. const canMergeWithLastHistory = !automaticFunctionCallingHistory || automaticFunctionCallingHistory.length === 0; if ( canMergeWithLastHistory && - lastHistoryEntry && - lastHistoryEntry.role === 'model' && - lastHistoryEntry.parts && - consolidatedOutputContents[0].role === 'model' + this.isTextContent(lastHistoryEntry) && + this.isTextContent(consolidatedOutputContents[0]) ) { - lastHistoryEntry.parts.push( - ...(consolidatedOutputContents[0].parts || []), - ); + // If both current and last are text, combine their text into the lastHistoryEntry's first part + // and append any other parts from the current content. + lastHistoryEntry.parts[0].text += + consolidatedOutputContents[0].parts[0].text || ''; + if (consolidatedOutputContents[0].parts.length > 1) { + lastHistoryEntry.parts.push( + ...consolidatedOutputContents[0].parts.slice(1), + ); + } consolidatedOutputContents.shift(); // Remove the first element as it's merged } this.history.push(...consolidatedOutputContents); } } + + private isTextContent( + content: Content | undefined, + ): content is Content & { parts: [{ text: string }, ...Part[]] } { + return !!( + content && + content.role === 'model' && + content.parts && + content.parts.length > 0 && + typeof content.parts[0].text === 'string' && + content.parts[0].text !== '' + ); + } } |