Skip to content

feat(core): Migrate chatCompressionService to model configs. #12863

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
joshualitt merged 1 commit into from
Nov 24, 2025
Merged

feat(core): Migrate chatCompressionService to model configs. #12863

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
20 changes: 20 additions & 0 deletions docs/get-started/configuration.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -361,10 +361,10 @@
}
}
},
"gemini-3-pro-preview": {

Check warning on line 364 in docs/get-started/configuration.md

View workflow job for this annotation

GitHub Actions / Lint 

Found sensitive keyword "gemini-3". Please make sure this change is appropriate to submit.
"extends": "chat-base-3",
"modelConfig": {
"model": "gemini-3-pro-preview"

Check warning on line 367 in docs/get-started/configuration.md

View workflow job for this annotation

GitHub Actions / Lint 

Found sensitive keyword "gemini-3". Please make sure this change is appropriate to submit.
}
},
"gemini-2.5-pro": {
Expand Down Expand Up @@ -490,6 +490,26 @@
"next-speaker-checker": {
"extends": "gemini-2.5-flash-base",
"modelConfig": {}
},
"chat-compression-3-pro": {
"modelConfig": {
"model": "gemini-3-pro-preview"

Check warning on line 496 in docs/get-started/configuration.md

View workflow job for this annotation

GitHub Actions / Lint 

Found sensitive keyword "gemini-3". Please make sure this change is appropriate to submit.
}
},
"chat-compression-2.5-pro": {
"modelConfig": {
"model": "gemini-2.5-pro"
}
},
"chat-compression-2.5-flash": {
"modelConfig": {
"model": "gemini-2.5-flash"
}
},
"chat-compression-2.5-flash-lite": {
"modelConfig": {
"model": "gemini-2.5-flash-lite"
}
}
}
```
Expand Down
20 changes: 20 additions & 0 deletions packages/core/src/config/defaultModelConfigs.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -59,10 +59,10 @@
// ensure these model configs can be used interactively.
// TODO(joshualitt): Introduce internal base configs for the various models,
// note: we will have to think carefully about names.
'gemini-3-pro-preview': {

Check warning on line 62 in packages/core/src/config/defaultModelConfigs.ts

View workflow job for this annotation

GitHub Actions / Lint 

Found sensitive keyword "gemini-3". Please make sure this change is appropriate to submit.
extends: 'chat-base-3',
modelConfig: {
model: 'gemini-3-pro-preview',

Check warning on line 65 in packages/core/src/config/defaultModelConfigs.ts

View workflow job for this annotation

GitHub Actions / Lint 

Found sensitive keyword "gemini-3". Please make sure this change is appropriate to submit.
},
},
'gemini-2.5-pro': {
Expand Down Expand Up @@ -183,5 +183,25 @@
extends: 'gemini-2.5-flash-base',
modelConfig: {},
},
'chat-compression-3-pro': {
modelConfig: {
model: 'gemini-3-pro-preview',

Check warning on line 188 in packages/core/src/config/defaultModelConfigs.ts

View workflow job for this annotation

GitHub Actions / Lint 

Found sensitive keyword "gemini-3". Please make sure this change is appropriate to submit.
},
},
'chat-compression-2.5-pro': {
modelConfig: {
model: 'gemini-2.5-pro',
},
},
'chat-compression-2.5-flash': {
modelConfig: {
model: 'gemini-2.5-flash',
},
},
'chat-compression-2.5-flash-lite': {
modelConfig: {
model: 'gemini-2.5-flash-lite',
},
},
},
};
75 changes: 42 additions & 33 deletions packages/core/src/services/chatCompressionService.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,14 +8,15 @@
import {
ChatCompressionService,
findCompressSplitPoint,
modelStringToModelConfigAlias,
} from './chatCompressionService.js';
import type { Content, GenerateContentResponse } from '@google/genai';
import { CompressionStatus } from '../core/turn.js';
import { tokenLimit } from '../core/tokenLimits.js';
import type { GeminiChat } from '../core/geminiChat.js';
import type { Config } from '../config/config.js';
import { getInitialChatHistory } from '../utils/environmentContext.js';
import type { ContentGenerator } from '../core/contentGenerator.js';
import { DEFAULT_GEMINI_MODEL } from '../config/models.js';

vi.mock('../core/tokenLimits.js');
vi.mock('../telemetry/loggers.js');
Expand Down Expand Up @@ -101,11 +102,34 @@
});
});

describe('modelStringToModelConfigAlias', () => {
it('should return the default model for unexpected aliases', () => {
expect(modelStringToModelConfigAlias('gemini-flash-flash')).toBe(
DEFAULT_GEMINI_MODEL,
);
});

it('should handle valid names', () => {
expect(modelStringToModelConfigAlias('gemini-3-pro-preview')).toBe(

Check warning on line 113 in packages/core/src/services/chatCompressionService.test.ts

View workflow job for this annotation

GitHub Actions / Lint 

Found sensitive keyword "gemini-3". Please make sure this change is appropriate to submit.
'chat-compression-3-pro',
);
expect(modelStringToModelConfigAlias('gemini-2.5-pro')).toBe(
'chat-compression-2.5-pro',
);
expect(modelStringToModelConfigAlias('gemini-2.5-flash')).toBe(
'chat-compression-2.5-flash',
);
expect(modelStringToModelConfigAlias('gemini-2.5-flash-lite')).toBe(
'chat-compression-2.5-flash-lite',
);
});
});

describe('ChatCompressionService', () => {
let service: ChatCompressionService;
let mockChat: GeminiChat;
let mockConfig: Config;
const mockModel = 'gemini-pro';
const mockModel = 'gemini-2.5-pro';
const mockPromptId = 'test-prompt-id';

beforeEach(() => {
Expand All @@ -114,9 +138,22 @@
getHistory: vi.fn(),
getLastPromptTokenCount: vi.fn().mockReturnValue(500),
} as unknown as GeminiChat;

const mockGenerateContent = vi.fn().mockResolvedValue({
candidates: [
{
content: {
parts: [{ text: 'Summary' }],
},
},
],
} as unknown as GenerateContentResponse);

mockConfig = {
getCompressionThreshold: vi.fn(),
getContentGenerator: vi.fn(),
getBaseLlmClient: vi.fn().mockReturnValue({
generateContent: mockGenerateContent,
}),
isInteractive: vi.fn().mockReturnValue(false),
} as unknown as Config;

Expand Down Expand Up @@ -190,18 +227,6 @@
vi.mocked(mockChat.getHistory).mockReturnValue(history);
vi.mocked(mockChat.getLastPromptTokenCount).mockReturnValue(800);
vi.mocked(tokenLimit).mockReturnValue(1000);
const mockGenerateContent = vi.fn().mockResolvedValue({
candidates: [
{
content: {
parts: [{ text: 'Summary' }],
},
},
],
} as unknown as GenerateContentResponse);
vi.mocked(mockConfig.getContentGenerator).mockReturnValue({
generateContent: mockGenerateContent,
} as unknown as ContentGenerator);

const result = await service.compress(
mockChat,
Expand All @@ -215,7 +240,7 @@
expect(result.info.compressionStatus).toBe(CompressionStatus.COMPRESSED);
expect(result.newHistory).not.toBeNull();
expect(result.newHistory![0].parts![0].text).toBe('Summary');
expect(mockGenerateContent).toHaveBeenCalled();
expect(mockConfig.getBaseLlmClient().generateContent).toHaveBeenCalled();
});

it('should force compress even if under threshold', async () => {
Expand All @@ -229,19 +254,6 @@
vi.mocked(mockChat.getLastPromptTokenCount).mockReturnValue(100);
vi.mocked(tokenLimit).mockReturnValue(1000);

const mockGenerateContent = vi.fn().mockResolvedValue({
candidates: [
{
content: {
parts: [{ text: 'Summary' }],
},
},
],
} as unknown as GenerateContentResponse);
vi.mocked(mockConfig.getContentGenerator).mockReturnValue({
generateContent: mockGenerateContent,
} as unknown as ContentGenerator);

const result = await service.compress(
mockChat,
mockPromptId,
Expand All @@ -265,7 +277,7 @@
vi.mocked(tokenLimit).mockReturnValue(1000);

const longSummary = 'a'.repeat(1000); // Long summary to inflate token count
const mockGenerateContent = vi.fn().mockResolvedValue({
vi.mocked(mockConfig.getBaseLlmClient().generateContent).mockResolvedValue({
candidates: [
{
content: {
Expand All @@ -274,9 +286,6 @@
},
],
} as unknown as GenerateContentResponse);
vi.mocked(mockConfig.getContentGenerator).mockReturnValue({
generateContent: mockGenerateContent,
} as unknown as ContentGenerator);

const result = await service.compress(
mockChat,
Expand Down
55 changes: 37 additions & 18 deletions packages/core/src/services/chatCompressionService.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,12 @@ import { getResponseText } from '../utils/partUtils.js';
import { logChatCompression } from '../telemetry/loggers.js';
import { makeChatCompressionEvent } from '../telemetry/types.js';
import { getInitialChatHistory } from '../utils/environmentContext.js';
import {
DEFAULT_GEMINI_FLASH_LITE_MODEL,
DEFAULT_GEMINI_FLASH_MODEL,
DEFAULT_GEMINI_MODEL,
PREVIEW_GEMINI_MODEL,
} from '../config/models.js';

/**
* Default threshold for compression token count as a fraction of the model's
Expand Down Expand Up @@ -75,6 +81,21 @@ export function findCompressSplitPoint(
return lastSplitPoint;
}

export function modelStringToModelConfigAlias(model: string): string {
switch (model) {
case PREVIEW_GEMINI_MODEL:
return 'chat-compression-3-pro';
case DEFAULT_GEMINI_MODEL:
return 'chat-compression-2.5-pro';
case DEFAULT_GEMINI_FLASH_MODEL:
return 'chat-compression-2.5-flash';
case DEFAULT_GEMINI_FLASH_LITE_MODEL:
return 'chat-compression-2.5-flash-lite';
default:
return DEFAULT_GEMINI_MODEL;
}
}

export class ChatCompressionService {
async compress(
chat: GeminiChat,
Expand Down Expand Up @@ -139,26 +160,24 @@ export class ChatCompressionService {
};
}

const summaryResponse = await config.getContentGenerator().generateContent(
{
model,
contents: [
...historyToCompress,
{
role: 'user',
parts: [
{
text: 'First, reason in your scratchpad. Then, generate the <state_snapshot>.',
},
],
},
],
config: {
systemInstruction: { text: getCompressionPrompt() },
const summaryResponse = await config.getBaseLlmClient().generateContent({
modelConfigKey: { model: modelStringToModelConfigAlias(model) },
contents: [
...historyToCompress,
{
role: 'user',
parts: [
{
text: 'First, reason in your scratchpad. Then, generate the <state_snapshot>.',
},
],
},
},
],
systemInstruction: { text: getCompressionPrompt() },
promptId,
);
// TODO(joshualitt): wire up a sensible abort signal,
abortSignal: new AbortController().signal,
});
const summary = getResponseText(summaryResponse) ?? '';

const extraHistory: Content[] = [
Expand Down
16 changes: 16 additions & 0 deletions packages/core/src/services/test-data/resolved-aliases.golden.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,8 +37,8 @@
"topK": 64
}
},
"gemini-3-pro-preview": {

Check warning on line 40 in packages/core/src/services/test-data/resolved-aliases.golden.json

View workflow job for this annotation

GitHub Actions / Lint 

Found sensitive keyword "gemini-3". Please make sure this change is appropriate to submit.
"model": "gemini-3-pro-preview",

Check warning on line 41 in packages/core/src/services/test-data/resolved-aliases.golden.json

View workflow job for this annotation

GitHub Actions / Lint 

Found sensitive keyword "gemini-3". Please make sure this change is appropriate to submit.
"generateContentConfig": {
"temperature": 1,
"topP": 0.95,
Expand Down Expand Up @@ -198,5 +198,21 @@
"temperature": 0,
"topP": 1
}
},
"chat-compression-3-pro": {
"model": "gemini-3-pro-preview",

Check warning on line 203 in packages/core/src/services/test-data/resolved-aliases.golden.json

View workflow job for this annotation

GitHub Actions / Lint 

Found sensitive keyword "gemini-3". Please make sure this change is appropriate to submit.
"generateContentConfig": {}

@joshualitt joshualitt Nov 22, 2025

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

This is intentional, the existing requests do not pass any settings

},
"chat-compression-2.5-pro": {
"model": "gemini-2.5-pro",
"generateContentConfig": {}
},
"chat-compression-2.5-flash": {
"model": "gemini-2.5-flash",
"generateContentConfig": {}
},
"chat-compression-2.5-flash-lite": {
"model": "gemini-2.5-flash-lite",
"generateContentConfig": {}
}
}
Loading