386 lines
16 KiB
TypeScript
386 lines
16 KiB
TypeScript
/**
|
|
* @license
|
|
* Copyright 2025 Google LLC
|
|
* SPDX-License-Identifier: Apache-2.0
|
|
*/
|
|
|
|
import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
|
|
import { Content, Models, GenerateContentConfig, Part } from '@google/genai';
|
|
import { GeminiChat } from './geminiChat.js';
|
|
|
|
// Mocks
|
|
const mockModelsModule = {
|
|
generateContent: vi.fn(),
|
|
generateContentStream: vi.fn(),
|
|
countTokens: vi.fn(),
|
|
embedContent: vi.fn(),
|
|
batchEmbedContents: vi.fn(),
|
|
} as unknown as Models;
|
|
|
|
describe('GeminiChat', () => {
|
|
let chat: GeminiChat;
|
|
const model = 'gemini-pro';
|
|
const config: GenerateContentConfig = {};
|
|
|
|
beforeEach(() => {
|
|
vi.clearAllMocks();
|
|
// Reset history for each test by creating a new instance
|
|
chat = new GeminiChat(mockModelsModule, model, config, []);
|
|
});
|
|
|
|
afterEach(() => {
|
|
vi.restoreAllMocks();
|
|
});
|
|
|
|
describe('recordHistory', () => {
|
|
const userInput: Content = {
|
|
role: 'user',
|
|
parts: [{ text: 'User input' }],
|
|
};
|
|
|
|
it('should add user input and a single model output to history', () => {
|
|
const modelOutput: Content[] = [
|
|
{ role: 'model', parts: [{ text: 'Model output' }] },
|
|
];
|
|
// @ts-expect-error Accessing private method for testing purposes
|
|
chat.recordHistory(userInput, modelOutput);
|
|
const history = chat.getHistory();
|
|
expect(history).toEqual([userInput, modelOutput[0]]);
|
|
});
|
|
|
|
it('should consolidate adjacent model outputs', () => {
|
|
const modelOutputParts: Content[] = [
|
|
{ role: 'model', parts: [{ text: 'Model part 1' }] },
|
|
{ role: 'model', parts: [{ text: 'Model part 2' }] },
|
|
];
|
|
// @ts-expect-error Accessing private method for testing purposes
|
|
chat.recordHistory(userInput, modelOutputParts);
|
|
const history = chat.getHistory();
|
|
expect(history.length).toBe(2);
|
|
expect(history[0]).toEqual(userInput);
|
|
expect(history[1].role).toBe('model');
|
|
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)', () => {
|
|
const mixedOutput: Content[] = [
|
|
{ role: 'model', parts: [{ text: 'Model 1' }] },
|
|
{ role: 'user', parts: [{ text: 'Unexpected User' }] }, // This should be pushed as is
|
|
{ role: 'model', parts: [{ text: 'Model 2' }] },
|
|
];
|
|
// @ts-expect-error Accessing private method for testing purposes
|
|
chat.recordHistory(userInput, mixedOutput);
|
|
const history = chat.getHistory();
|
|
expect(history.length).toBe(4); // user, model1, user_unexpected, model2
|
|
expect(history[0]).toEqual(userInput);
|
|
expect(history[1]).toEqual(mixedOutput[0]);
|
|
expect(history[2]).toEqual(mixedOutput[1]);
|
|
expect(history[3]).toEqual(mixedOutput[2]);
|
|
});
|
|
|
|
it('should consolidate multiple adjacent model outputs correctly', () => {
|
|
const modelOutputParts: Content[] = [
|
|
{ role: 'model', parts: [{ text: 'M1' }] },
|
|
{ role: 'model', parts: [{ text: 'M2' }] },
|
|
{ role: 'model', parts: [{ text: 'M3' }] },
|
|
];
|
|
// @ts-expect-error Accessing private method for testing purposes
|
|
chat.recordHistory(userInput, modelOutputParts);
|
|
const history = chat.getHistory();
|
|
expect(history.length).toBe(2);
|
|
expect(history[1].parts).toEqual([{ text: 'M1M2M3' }]);
|
|
});
|
|
|
|
it('should not consolidate if roles are different between model outputs', () => {
|
|
const modelOutputParts: Content[] = [
|
|
{ role: 'model', parts: [{ text: 'M1' }] },
|
|
{ role: 'user', parts: [{ text: 'Interjecting User' }] },
|
|
{ role: 'model', parts: [{ text: 'M2' }] },
|
|
];
|
|
// @ts-expect-error Accessing private method for testing purposes
|
|
chat.recordHistory(userInput, modelOutputParts);
|
|
const history = chat.getHistory();
|
|
expect(history.length).toBe(4); // user, M1, Interjecting User, M2
|
|
expect(history[1].parts).toEqual([{ text: 'M1' }]);
|
|
expect(history[3].parts).toEqual([{ text: 'M2' }]);
|
|
});
|
|
|
|
it('should merge with last history entry if it is also a model output', () => {
|
|
// @ts-expect-error Accessing private property for test setup
|
|
chat.history = [
|
|
userInput,
|
|
{ role: 'model', parts: [{ text: 'Initial Model Output' }] },
|
|
]; // Prime the history
|
|
|
|
const newModelOutput: Content[] = [
|
|
{ role: 'model', parts: [{ text: 'New Model Part 1' }] },
|
|
{ role: 'model', parts: [{ text: 'New Model Part 2' }] },
|
|
];
|
|
// @ts-expect-error Accessing private method for testing purposes
|
|
chat.recordHistory(userInput, newModelOutput); // userInput here is for the *next* turn, but history is already primed
|
|
|
|
// Reset and set up a more realistic scenario for merging with existing history
|
|
chat = new GeminiChat(mockModelsModule, model, config, []);
|
|
const firstUserInput: Content = {
|
|
role: 'user',
|
|
parts: [{ text: 'First user input' }],
|
|
};
|
|
const firstModelOutput: Content[] = [
|
|
{ role: 'model', parts: [{ text: 'First model response' }] },
|
|
];
|
|
// @ts-expect-error Accessing private method for testing purposes
|
|
chat.recordHistory(firstUserInput, firstModelOutput);
|
|
|
|
const secondUserInput: Content = {
|
|
role: 'user',
|
|
parts: [{ text: 'Second user input' }],
|
|
};
|
|
const secondModelOutput: Content[] = [
|
|
{ role: 'model', parts: [{ text: 'Second model response part 1' }] },
|
|
{ role: 'model', parts: [{ text: 'Second model response part 2' }] },
|
|
];
|
|
// @ts-expect-error Accessing private method for testing purposes
|
|
chat.recordHistory(secondUserInput, secondModelOutput);
|
|
|
|
const finalHistory = chat.getHistory();
|
|
expect(finalHistory.length).toBe(4); // user1, model1, user2, model2(consolidated)
|
|
expect(finalHistory[0]).toEqual(firstUserInput);
|
|
expect(finalHistory[1]).toEqual(firstModelOutput[0]);
|
|
expect(finalHistory[2]).toEqual(secondUserInput);
|
|
expect(finalHistory[3].role).toBe('model');
|
|
expect(finalHistory[3].parts).toEqual([
|
|
{ text: 'Second model response part 1Second model response part 2' },
|
|
]);
|
|
});
|
|
|
|
it('should correctly merge consolidated new output with existing model history', () => {
|
|
// Setup: history ends with a model turn
|
|
const initialUser: Content = {
|
|
role: 'user',
|
|
parts: [{ text: 'Initial user query' }],
|
|
};
|
|
const initialModel: Content = {
|
|
role: 'model',
|
|
parts: [{ text: 'Initial model answer.' }],
|
|
};
|
|
chat = new GeminiChat(mockModelsModule, model, config, [
|
|
initialUser,
|
|
initialModel,
|
|
]);
|
|
|
|
// New interaction
|
|
const currentUserInput: Content = {
|
|
role: 'user',
|
|
parts: [{ text: 'Follow-up question' }],
|
|
};
|
|
const newModelParts: Content[] = [
|
|
{ role: 'model', parts: [{ text: 'Part A of new answer.' }] },
|
|
{ role: 'model', parts: [{ text: 'Part B of new answer.' }] },
|
|
];
|
|
|
|
// @ts-expect-error Accessing private method for testing purposes
|
|
chat.recordHistory(currentUserInput, newModelParts);
|
|
const history = chat.getHistory();
|
|
|
|
// Expected: initialUser, initialModel, currentUserInput, consolidatedNewModelParts
|
|
expect(history.length).toBe(4);
|
|
expect(history[0]).toEqual(initialUser);
|
|
expect(history[1]).toEqual(initialModel);
|
|
expect(history[2]).toEqual(currentUserInput);
|
|
expect(history[3].role).toBe('model');
|
|
expect(history[3].parts).toEqual([
|
|
{ text: 'Part A of new answer.Part B of new answer.' },
|
|
]);
|
|
});
|
|
|
|
it('should handle empty modelOutput array', () => {
|
|
// @ts-expect-error Accessing private method for testing purposes
|
|
chat.recordHistory(userInput, []);
|
|
const history = chat.getHistory();
|
|
// If modelOutput is empty, it might push a default empty model part depending on isFunctionResponse
|
|
// Assuming isFunctionResponse(userInput) is false for this simple text input
|
|
expect(history.length).toBe(2);
|
|
expect(history[0]).toEqual(userInput);
|
|
expect(history[1].role).toBe('model');
|
|
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' }] },
|
|
{ role: 'model', parts: undefined as unknown as Part[] }, // Test undefined parts
|
|
{ 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(4); // userInput, model1 (text), model2 (undefined parts), model3 (empty parts)
|
|
expect(history[0]).toEqual(userInput);
|
|
expect(history[1].role).toBe('model');
|
|
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', () => {
|
|
const afcHistory: Content[] = [
|
|
{ role: 'user', parts: [{ text: 'AFC User' }] },
|
|
{ role: 'model', parts: [{ text: 'AFC Model' }] },
|
|
];
|
|
const modelOutput: Content[] = [
|
|
{ role: 'model', parts: [{ text: 'Regular Model Output' }] },
|
|
];
|
|
// @ts-expect-error Accessing private method for testing purposes
|
|
chat.recordHistory(userInput, modelOutput, afcHistory);
|
|
const history = chat.getHistory();
|
|
expect(history.length).toBe(3);
|
|
expect(history[0]).toEqual(afcHistory[0]);
|
|
expect(history[1]).toEqual(afcHistory[1]);
|
|
expect(history[2]).toEqual(modelOutput[0]);
|
|
});
|
|
|
|
it('should add userInput if AFC history is present but empty', () => {
|
|
const modelOutput: Content[] = [
|
|
{ role: 'model', parts: [{ text: 'Model Output' }] },
|
|
];
|
|
// @ts-expect-error Accessing private method for testing purposes
|
|
chat.recordHistory(userInput, modelOutput, []); // Empty AFC history
|
|
const history = chat.getHistory();
|
|
expect(history.length).toBe(2);
|
|
expect(history[0]).toEqual(userInput);
|
|
expect(history[1]).toEqual(modelOutput[0]);
|
|
});
|
|
|
|
it('should skip "thought" content from modelOutput', () => {
|
|
const modelOutputWithThought: Content[] = [
|
|
{ role: 'model', parts: [{ thought: true }, { text: 'Visible text' }] },
|
|
{ role: 'model', parts: [{ text: 'Another visible text' }] },
|
|
];
|
|
// @ts-expect-error Accessing private method for testing purposes
|
|
chat.recordHistory(userInput, modelOutputWithThought);
|
|
const history = chat.getHistory();
|
|
expect(history.length).toBe(2); // User input + consolidated model output
|
|
expect(history[0]).toEqual(userInput);
|
|
expect(history[1].role).toBe('model');
|
|
// The 'thought' part is skipped, 'Another visible text' becomes the first part.
|
|
expect(history[1].parts).toEqual([{ text: 'Another visible text' }]);
|
|
});
|
|
|
|
it('should skip "thought" content even if it is the only content', () => {
|
|
const modelOutputOnlyThought: Content[] = [
|
|
{ role: 'model', parts: [{ thought: true }] },
|
|
];
|
|
// @ts-expect-error Accessing private method for testing purposes
|
|
chat.recordHistory(userInput, modelOutputOnlyThought);
|
|
const history = chat.getHistory();
|
|
expect(history.length).toBe(1); // User input + default empty model part
|
|
expect(history[0]).toEqual(userInput);
|
|
});
|
|
|
|
it('should correctly consolidate text parts when a thought part is in between', () => {
|
|
const modelOutputMixed: Content[] = [
|
|
{ role: 'model', parts: [{ text: 'Part 1.' }] },
|
|
{
|
|
role: 'model',
|
|
parts: [{ thought: true }, { text: 'Should be skipped' }],
|
|
},
|
|
{ role: 'model', parts: [{ text: 'Part 2.' }] },
|
|
];
|
|
// @ts-expect-error Accessing private method for testing purposes
|
|
chat.recordHistory(userInput, modelOutputMixed);
|
|
const history = chat.getHistory();
|
|
expect(history.length).toBe(2);
|
|
expect(history[0]).toEqual(userInput);
|
|
expect(history[1].role).toBe('model');
|
|
expect(history[1].parts).toEqual([{ text: 'Part 1.Part 2.' }]);
|
|
});
|
|
|
|
it('should handle multiple thought parts correctly', () => {
|
|
const modelOutputMultipleThoughts: Content[] = [
|
|
{ role: 'model', parts: [{ thought: true }] },
|
|
{ role: 'model', parts: [{ text: 'Visible 1' }] },
|
|
{ role: 'model', parts: [{ thought: true }] },
|
|
{ role: 'model', parts: [{ text: 'Visible 2' }] },
|
|
];
|
|
// @ts-expect-error Accessing private method for testing purposes
|
|
chat.recordHistory(userInput, modelOutputMultipleThoughts);
|
|
const history = chat.getHistory();
|
|
expect(history.length).toBe(2);
|
|
expect(history[0]).toEqual(userInput);
|
|
expect(history[1].role).toBe('model');
|
|
expect(history[1].parts).toEqual([{ text: 'Visible 1Visible 2' }]);
|
|
});
|
|
|
|
it('should handle thought part at the end of outputContents', () => {
|
|
const modelOutputThoughtAtEnd: Content[] = [
|
|
{ role: 'model', parts: [{ text: 'Visible text' }] },
|
|
{ role: 'model', parts: [{ thought: true }] },
|
|
];
|
|
// @ts-expect-error Accessing private method for testing purposes
|
|
chat.recordHistory(userInput, modelOutputThoughtAtEnd);
|
|
const history = chat.getHistory();
|
|
expect(history.length).toBe(2);
|
|
expect(history[0]).toEqual(userInput);
|
|
expect(history[1].role).toBe('model');
|
|
expect(history[1].parts).toEqual([{ text: 'Visible text' }]);
|
|
});
|
|
});
|
|
|
|
describe('addHistory', () => {
|
|
it('should add a new content item to the history', () => {
|
|
const newContent: Content = {
|
|
role: 'user',
|
|
parts: [{ text: 'A new message' }],
|
|
};
|
|
chat.addHistory(newContent);
|
|
const history = chat.getHistory();
|
|
expect(history.length).toBe(1);
|
|
expect(history[0]).toEqual(newContent);
|
|
});
|
|
|
|
it('should add multiple items correctly', () => {
|
|
const content1: Content = {
|
|
role: 'user',
|
|
parts: [{ text: 'Message 1' }],
|
|
};
|
|
const content2: Content = {
|
|
role: 'model',
|
|
parts: [{ text: 'Message 2' }],
|
|
};
|
|
chat.addHistory(content1);
|
|
chat.addHistory(content2);
|
|
const history = chat.getHistory();
|
|
expect(history.length).toBe(2);
|
|
expect(history[0]).toEqual(content1);
|
|
expect(history[1]).toEqual(content2);
|
|
});
|
|
});
|
|
});
|