Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
66 changes: 66 additions & 0 deletions packages/core/src/core/geminiChat.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -864,6 +864,72 @@ describe('GeminiChat', () => {
expect(uiTelemetryService.setLastPromptTokenCount).not.toHaveBeenCalled();
});

it('should set temperature to 1 on retry', async () => {
// Use mockImplementationOnce to provide a fresh, promise-wrapped generator for each attempt.
vi.mocked(mockContentGenerator.generateContentStream)
.mockImplementationOnce(async () =>
// First call returns an invalid stream
(async function* () {
yield {
candidates: [{ content: { parts: [{ text: '' }] } }], // Invalid empty text part
} as unknown as GenerateContentResponse;
})(),
)
.mockImplementationOnce(async () =>
// Second call returns a valid stream
(async function* () {
yield {
candidates: [
{
content: { parts: [{ text: 'Successful response' }] },
finishReason: 'STOP',
},
],
} as unknown as GenerateContentResponse;
})(),
);

const stream = await chat.sendMessageStream(
'test-model',
{ message: 'test', config: { temperature: 0.5 } },
'prompt-id-retry-temperature',
);

for await (const _ of stream) {
// consume stream
}

expect(mockContentGenerator.generateContentStream).toHaveBeenCalledTimes(
2,
);

// First call should have original temperature
expect(
mockContentGenerator.generateContentStream,
).toHaveBeenNthCalledWith(
1,
expect.objectContaining({
config: expect.objectContaining({
temperature: 0.5,
}),
}),
'prompt-id-retry-temperature',
);

// Second call (retry) should have temperature 1
expect(
mockContentGenerator.generateContentStream,
).toHaveBeenNthCalledWith(
2,
expect.objectContaining({
config: expect.objectContaining({
temperature: 1,
}),
}),
'prompt-id-retry-temperature',
);
});

it('should fail after all retries on persistent invalid content and report metrics', async () => {
vi.mocked(mockContentGenerator.generateContentStream).mockImplementation(
async () =>
Expand Down
11 changes: 10 additions & 1 deletion packages/core/src/core/geminiChat.ts
Original file line number Diff line number Diff line change
Expand Up @@ -273,10 +273,19 @@ export class GeminiChat {
yield { type: StreamEventType.RETRY };
}

// If this is a retry, set temperature to 1 to encourage different output.
const currentParams = { ...params };
if (attempt > 0) {
currentParams.config = {
...currentParams.config,
temperature: 1,
};
Comment on lines +279 to +282

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

There's a potential issue here regarding the use of config versus generationConfig. The SendMessageParameters type is imported from @google/genai, which uses generationConfig for generation settings. This file, including this new change, consistently uses a config property instead. This discrepancy can lead to type-related bugs and confusion.

While this change is consistent with the existing code in this file, a broader refactoring to use generationConfig would be beneficial for long-term maintainability and correctness. A full fix would require changes outside of this diff (e.g., in makeApiCallAndProcessStream), so I am not providing a direct suggestion here.

}

const stream = await self.makeApiCallAndProcessStream(
model,
requestContents,
params,
currentParams,
prompt_id,
);

Expand Down