fix: Refine model message consolidation for improved model interaction (#685)

This commit is contained in:
N. Taylor Mullen
2025-06-02 00:28:14 -07:00
committed by GitHub
parent 6d417186cb
commit 27ba28ef76
2 changed files with 67 additions and 30 deletions

View File

@@ -69,10 +69,7 @@ describe('GeminiChat', () => {
expect(history.length).toBe(2); expect(history.length).toBe(2);
expect(history[0]).toEqual(userInput); expect(history[0]).toEqual(userInput);
expect(history[1].role).toBe('model'); expect(history[1].role).toBe('model');
expect(history[1].parts).toEqual([ expect(history[1].parts).toEqual([{ text: 'Model part 1Model part 2' }]);
{ text: 'Model part 1' },
{ text: 'Model part 2' },
]);
}); });
it('should handle a mix of user and model roles in outputContents (though unusual)', () => { it('should handle a mix of user and model roles in outputContents (though unusual)', () => {
@@ -101,11 +98,7 @@ describe('GeminiChat', () => {
chat.recordHistory(userInput, modelOutputParts); chat.recordHistory(userInput, modelOutputParts);
const history = chat.getHistory(); const history = chat.getHistory();
expect(history.length).toBe(2); expect(history.length).toBe(2);
expect(history[1].parts).toEqual([ expect(history[1].parts).toEqual([{ text: 'M1M2M3' }]);
{ text: 'M1' },
{ text: 'M2' },
{ text: 'M3' },
]);
}); });
it('should not consolidate if roles are different between model outputs', () => { it('should not consolidate if roles are different between model outputs', () => {
@@ -177,8 +170,7 @@ describe('GeminiChat', () => {
expect(finalHistory[2]).toEqual(secondUserInput); expect(finalHistory[2]).toEqual(secondUserInput);
expect(finalHistory[3].role).toBe('model'); expect(finalHistory[3].role).toBe('model');
expect(finalHistory[3].parts).toEqual([ expect(finalHistory[3].parts).toEqual([
{ text: 'Second model response part 1' }, { text: 'Second model response part 1Second model response part 2' },
{ text: 'Second model response part 2' },
]); ]);
}); });
@@ -218,8 +210,7 @@ describe('GeminiChat', () => {
expect(history[2]).toEqual(currentUserInput); expect(history[2]).toEqual(currentUserInput);
expect(history[3].role).toBe('model'); expect(history[3].role).toBe('model');
expect(history[3].parts).toEqual([ expect(history[3].parts).toEqual([
{ text: 'Part A of new answer.' }, { text: 'Part A of new answer.Part B of new answer.' },
{ text: 'Part B of new answer.' },
]); ]);
}); });
@@ -235,6 +226,31 @@ describe('GeminiChat', () => {
expect(history[1].parts).toEqual([]); 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)', () => { it('should handle modelOutput with parts being undefined or empty (if they pass initial every check)', () => {
const modelOutputUndefinedParts: Content[] = [ const modelOutputUndefinedParts: Content[] = [
{ role: 'model', parts: [{ text: 'Text part' }] }, { role: 'model', parts: [{ text: 'Text part' }] },
@@ -244,10 +260,14 @@ describe('GeminiChat', () => {
// @ts-expect-error Accessing private method for testing purposes // @ts-expect-error Accessing private method for testing purposes
chat.recordHistory(userInput, modelOutputUndefinedParts); chat.recordHistory(userInput, modelOutputUndefinedParts);
const history = chat.getHistory(); 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'); 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[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', () => { it('should correctly handle automaticFunctionCallingHistory', () => {

View File

@@ -15,6 +15,7 @@ import {
SendMessageParameters, SendMessageParameters,
GoogleGenAI, GoogleGenAI,
createUserContent, createUserContent,
Part,
} from '@google/genai'; } from '@google/genai';
import { retryWithBackoff } from '../utils/retry.js'; import { retryWithBackoff } from '../utils/retry.js';
import { isFunctionResponse } from '../utils/messageInspectors.js'; import { isFunctionResponse } from '../utils/messageInspectors.js';
@@ -343,13 +344,13 @@ export class GeminiChat {
for (const content of outputContents) { for (const content of outputContents) {
const lastContent = const lastContent =
consolidatedOutputContents[consolidatedOutputContents.length - 1]; consolidatedOutputContents[consolidatedOutputContents.length - 1];
if ( if (this.isTextContent(lastContent) && this.isTextContent(content)) {
lastContent && // If both current and last are text, combine their text into the lastContent's first part
lastContent.role === 'model' && // and append any other parts from the current content.
content.role === 'model' && lastContent.parts[0].text += content.parts[0].text || '';
lastContent.parts if (content.parts.length > 1) {
) { lastContent.parts.push(...content.parts.slice(1));
lastContent.parts.push(...(content.parts || [])); }
} else { } else {
consolidatedOutputContents.push(content); consolidatedOutputContents.push(content);
} }
@@ -357,24 +358,40 @@ export class GeminiChat {
if (consolidatedOutputContents.length > 0) { if (consolidatedOutputContents.length > 0) {
const lastHistoryEntry = this.history[this.history.length - 1]; 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 = const canMergeWithLastHistory =
!automaticFunctionCallingHistory || !automaticFunctionCallingHistory ||
automaticFunctionCallingHistory.length === 0; automaticFunctionCallingHistory.length === 0;
if ( if (
canMergeWithLastHistory && canMergeWithLastHistory &&
lastHistoryEntry && this.isTextContent(lastHistoryEntry) &&
lastHistoryEntry.role === 'model' && this.isTextContent(consolidatedOutputContents[0])
lastHistoryEntry.parts &&
consolidatedOutputContents[0].role === 'model'
) { ) {
lastHistoryEntry.parts.push( // If both current and last are text, combine their text into the lastHistoryEntry's first part
...(consolidatedOutputContents[0].parts || []), // 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 consolidatedOutputContents.shift(); // Remove the first element as it's merged
} }
this.history.push(...consolidatedOutputContents); 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 !== ''
);
}
} }