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
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 67 additions and 30 deletions

View File

@ -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', () => {

View File

@ -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 !== ''
);
}
}