Skip to content

Commit b29c02d

Browse files
authored
Fix not using flash for next speaker check (google-gemini#4016)
1 parent 44ef040 commit b29c02d

File tree

2 files changed

+21
-0
lines changed

2 files changed

+21
-0
lines changed

packages/core/src/utils/nextSpeakerChecker.test.ts

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66

77
import { describe, it, expect, vi, beforeEach, Mock, afterEach } from 'vitest';
88
import { Content, GoogleGenAI, Models } from '@google/genai';
9+
import { DEFAULT_GEMINI_FLASH_MODEL } from '../config/models.js';
910
import { GeminiClient } from '../core/client.js';
1011
import { Config } from '../config/config.js';
1112
import { checkNextSpeaker, NextSpeakerResponse } from './nextSpeakerChecker.js';
@@ -231,4 +232,22 @@ describe('checkNextSpeaker', () => {
231232
);
232233
expect(result).toBeNull();
233234
});
235+
236+
it('should call generateJson with DEFAULT_GEMINI_FLASH_MODEL', async () => {
237+
(chatInstance.getHistory as Mock).mockReturnValue([
238+
{ role: 'model', parts: [{ text: 'Some model output.' }] },
239+
] as Content[]);
240+
const mockApiResponse: NextSpeakerResponse = {
241+
reasoning: 'Model made a statement, awaiting user input.',
242+
next_speaker: 'user',
243+
};
244+
(mockGeminiClient.generateJson as Mock).mockResolvedValue(mockApiResponse);
245+
246+
await checkNextSpeaker(chatInstance, mockGeminiClient, abortSignal);
247+
248+
expect(mockGeminiClient.generateJson).toHaveBeenCalled();
249+
const generateJsonCall = (mockGeminiClient.generateJson as Mock).mock
250+
.calls[0];
251+
expect(generateJsonCall[3]).toBe(DEFAULT_GEMINI_FLASH_MODEL);
252+
});
234253
});

packages/core/src/utils/nextSpeakerChecker.ts

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
*/
66

77
import { Content, SchemaUnion, Type } from '@google/genai';
8+
import { DEFAULT_GEMINI_FLASH_MODEL } from '../config/models.js';
89
import { GeminiClient } from '../core/client.js';
910
import { GeminiChat } from '../core/geminiChat.js';
1011
import { isFunctionResponse } from './messageInspectors.js';
@@ -131,6 +132,7 @@ export async function checkNextSpeaker(
131132
contents,
132133
RESPONSE_SCHEMA,
133134
abortSignal,
135+
DEFAULT_GEMINI_FLASH_MODEL,
134136
)) as unknown as NextSpeakerResponse;
135137

136138
if (

0 commit comments

Comments
 (0)