Skip to content

Commit 383d5a0

Browse files
committed
Use structured prompt for compression.
1 parent dbd6260 commit 383d5a0

File tree

3 files changed

+103
-55
lines changed

3 files changed

+103
-55
lines changed

packages/core/src/core/client.ts

Lines changed: 38 additions & 51 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@ import {
2222
ChatCompressionInfo,
2323
} from './turn.js';
2424
import { Config } from '../config/config.js';
25-
import { getCoreSystemPrompt } from './prompts.js';
25+
import { getCoreSystemPrompt, getCompressionPrompt } from './prompts.js';
2626
import { ReadManyFilesTool } from '../tools/read-many-files.js';
2727
import { getResponseText } from '../utils/generateContentResponseUtilities.js';
2828
import { checkNextSpeaker } from '../utils/nextSpeakerChecker.js';
@@ -171,7 +171,7 @@ export class GeminiClient {
171171
const toolRegistry = await this.config.getToolRegistry();
172172
const toolDeclarations = toolRegistry.getFunctionDeclarations();
173173
const tools: Tool[] = [{ functionDeclarations: toolDeclarations }];
174-
const initialHistory: Content[] = [
174+
const history: Content[] = [
175175
{
176176
role: 'user',
177177
parts: envParts,
@@ -180,8 +180,8 @@ export class GeminiClient {
180180
role: 'model',
181181
parts: [{ text: 'Got it. Thanks for the context!' }],
182182
},
183+
...(extraHistory ?? []),
183184
];
184-
const history = initialHistory.concat(extraHistory ?? []);
185185
try {
186186
const userMemory = this.config.getUserMemory();
187187
const systemInstruction = getCoreSystemPrompt(userMemory);
@@ -428,74 +428,61 @@ export class GeminiClient {
428428
async tryCompressChat(
429429
force: boolean = false,
430430
): Promise<ChatCompressionInfo | null> {
431-
const history = this.getChat().getHistory(true); // Get curated history
431+
const curatedHistory = this.getChat().getHistory(true);
432432

433433
// Regardless of `force`, don't do anything if the history is empty.
434-
if (history.length === 0) {
434+
if (curatedHistory.length === 0) {
435435
return null;
436436
}
437437

438-
const { totalTokens: originalTokenCount } =
438+
let { totalTokens: originalTokenCount } =
439439
await this.getContentGenerator().countTokens({
440440
model: this.model,
441-
contents: history,
441+
contents: curatedHistory,
442442
});
443+
if (originalTokenCount === undefined) {
444+
console.warn(`Could not determine token count for model ${this.model}.`);
445+
originalTokenCount = 0;
446+
}
443447

444-
// If not forced, check if we should compress based on context size.
445-
if (!force) {
446-
if (originalTokenCount === undefined) {
447-
// If token count is undefined, we can't determine if we need to compress.
448-
console.warn(
449-
`Could not determine token count for model ${this.model}. Skipping compression check.`,
450-
);
451-
return null;
452-
}
453-
const tokenCount = originalTokenCount; // Now guaranteed to be a number
454-
455-
const limit = tokenLimit(this.model);
456-
if (!limit) {
457-
// If no limit is defined for the model, we can't compress.
458-
console.warn(
459-
`No token limit defined for model ${this.model}. Skipping compression check.`,
460-
);
461-
return null;
462-
}
463-
464-
if (tokenCount < 0.95 * limit) {
465-
return null;
466-
}
448+
// Don't compress if not forced and we are under the limit.
449+
if (!force && originalTokenCount < 0.95 * tokenLimit(this.model)) {
450+
return null;
467451
}
468452

469-
const summarizationRequestMessage = {
470-
text: 'Summarize our conversation up to this point. The summary should be a concise yet comprehensive overview of all key topics, questions, answers, and important details discussed. This summary will replace the current chat history to conserve tokens, so it must capture everything essential to understand the context and continue our conversation effectively as if no information was lost.',
471-
};
472-
const response = await this.getChat().sendMessage({
473-
message: summarizationRequestMessage,
453+
const { text: summary } = await this.getChat().sendMessage({
454+
message: {
455+
text: 'First, reason in your scratchpad. Then, generate the <state_snapshot>.',
456+
},
457+
config: {
458+
systemInstruction: { text: getCompressionPrompt() },
459+
},
474460
});
475-
const newHistory = [
461+
this.chat = await this.startChat([
476462
{
477463
role: 'user',
478-
parts: [summarizationRequestMessage],
464+
parts: [{ text: summary }],
479465
},
480466
{
481467
role: 'model',
482-
parts: [{ text: response.text }],
468+
parts: [{ text: 'Got it. Thanks for the additional context!' }],
483469
},
484-
];
485-
this.chat = await this.startChat(newHistory);
486-
const newTokenCount = (
470+
]);
471+
472+
const { totalTokens: newTokenCount } =
487473
await this.getContentGenerator().countTokens({
488474
model: this.model,
489-
contents: newHistory,
490-
})
491-
).totalTokens;
492-
493-
return originalTokenCount && newTokenCount
494-
? {
495-
originalTokenCount,
496-
newTokenCount,
497-
}
498-
: null;
475+
contents: this.getChat().getHistory(),
476+
});
477+
if (newTokenCount === undefined) {
478+
console.warn('Could not determine compressed history token count.');
479+
return null;
480+
}
481+
482+
return {
483+
originalTokenCount,
484+
newTokenCount,
485+
};
499486
}
500487

501488
/**

packages/core/src/core/geminiChat.ts

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -72,10 +72,6 @@ function isValidContent(content: Content): boolean {
7272
* @throws Error if the history contains an invalid role.
7373
*/
7474
function validateHistory(history: Content[]) {
75-
// Empty history is valid.
76-
if (history.length === 0) {
77-
return;
78-
}
7975
for (const content of history) {
8076
if (content.role !== 'user' && content.role !== 'model') {
8177
throw new Error(`Role must be user or model, but got ${content.role}.`);

packages/core/src/core/prompts.ts

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -271,3 +271,68 @@ Your core function is efficient and safe assistance. Balance extreme conciseness
271271

272272
return `${basePrompt}${memorySuffix}`;
273273
}
274+
275+
/**
276+
* Provides the system prompt for the history compression process.
277+
* This prompt instructs the model to act as a specialized state manager,
278+
* think in a scratchpad, and produce a structured XML summary.
279+
*/
280+
export function getCompressionPrompt(): string {
281+
return `
282+
You are the component that summarizes internal chat history into a given structure.
283+
284+
When the conversation history grows too large, you will be invoked to distill the entire history into a concise, structured XML snapshot. This snapshot is CRITICAL, as it will become the agent's *only* memory of the past. The agent will resume its work based solely on this snapshot. All crucial details, plans, errors, and user directives MUST be preserved.
285+
286+
First, you will think through the entire history in a private <scratchpad>. Review the user's overall goal, the agent's actions, tool outputs, file modifications, and any unresolved questions. Identify every piece of information that is essential for future actions.
287+
288+
After your reasoning is complete, generate the final <compressed_chat_history> XML object. Be incredibly dense with information. Omit any irrelevant conversational filler.
289+
290+
The structure MUST be as follows:
291+
292+
<compressed_chat_history>
293+
<overall_goal>
294+
<!-- A single, concise sentence describing the user's high-level objective. -->
295+
<!-- Example: "Refactor the authentication service to use a new JWT library." -->
296+
</overall_goal>
297+
298+
<key_knowledge>
299+
<!-- Crucial facts, conventions, and constraints the agent must remember based on the conversation history and interaction with the user. Use bullet points. -->
300+
<!-- Example:
301+
- Build Command: \`npm run build\`
302+
- Testing: Tests are run with \`npm test\`. Test files must end in \`.test.ts\`.
303+
- API Endpoint: The primary API endpoint is \`https://api.example.com/v2\`.
304+
305+
-->
306+
</key_knowledge>
307+
308+
<file_system_state>
309+
<!-- List files that have been created, read, modified, or deleted. Note their status and critical learnings. -->
310+
<!-- Example:
311+
- CWD: \`/home/user/project/src\`
312+
- READ: \`package.json\` - Confirmed 'axios' is a dependency.
313+
- MODIFIED: \`services/auth.ts\` - Replaced 'jsonwebtoken' with 'jose'.
314+
- CREATED: \`tests/new-feature.test.ts\` - Initial test structure for the new feature.
315+
-->
316+
</file_system_state>
317+
318+
<recent_actions>
319+
<!-- A summary of the last few significant agent actions and their outcomes. Focus on facts. -->
320+
<!-- Example:
321+
- Ran \`grep 'old_function'\` which returned 3 results in 2 files.
322+
- Ran \`npm run test\`, which failed due to a snapshot mismatch in \`UserProfile.test.ts\`.
323+
- Ran \`ls -F static/\` and discovered image assets are stored as \`.webp\`.
324+
-->
325+
</recent_actions>
326+
327+
<current_plan>
328+
<!-- The agent's step-by-step plan. Mark completed steps. -->
329+
<!-- Example:
330+
1. [DONE] Identify all files using the deprecated 'UserAPI'.
331+
2. [IN PROGRESS] Refactor \`src/components/UserProfile.tsx\` to use the new 'ProfileAPI'.
332+
3. [TODO] Refactor the remaining files.
333+
4. [TODO] Update tests to reflect the API change.
334+
-->
335+
</current_plan>
336+
</compressed_chat_history>
337+
`.trim();
338+
}

0 commit comments

Comments
 (0)