MediaMetz/gemini-cli
1
0
Fork 0
mirror of https://github.com/google-gemini/gemini-cli.git synced 2026-07-12 19:10:55 -07:00
Code Issues Packages Projects Releases Wiki Activity

token counting service

Browse Source
This commit is contained in:
Your Name
2026-04-06 19:48:44 +00:00
parent d7433ddd03
commit c1b06fec0d
9 changed files with 99 additions and 97 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/core/src/context/historyObserver.ts
+1 -1
View File
@@ -36,7 +36,7 @@ export class HistoryObserver {
this.unsubscribeHistory = this.chatHistory.subscribe((_event: HistoryEvent) => {
// Rebuild the pristine IR graph from the full source history on every change.
const pristineEpisodes = IrMapper.toIr(this.chatHistory.get());
const pristineEpisodes = IrMapper.toIr(this.chatHistory.get(), this.sidecar.tokenCalculator);
this.tracer.logEvent('HistoryObserver', 'Rebuilt pristine graph from chat history update', { episodeCount: pristineEpisodes.length });
this.onIrRebuilt(pristineEpisodes);
packages/core/src/context/ir/mapper.ts
+4 -7
View File
@@ -6,20 +6,17 @@
import type { Content } from '@google/genai';
import type { Episode } from './types.js';
import { toIr, setMapperConfig } from './toIr.js';
import { toIr } from './toIr.js';
import { fromIr } from './fromIr.js';
import type { ContextTokenCalculator } from '../utils/contextTokenCalculator.js';
export class IrMapper {
static setConfig(cfg: { charsPerToken?: number }) {
setMapperConfig(cfg);
}
/**
* Translates a flat Gemini Content[] array into our rich Episodic Intermediate Representation.
* Groups adjacent function calls and responses into unified ToolExecution nodes.
*/
static toIr(history: readonly Content[]): Episode[] {
return toIr(history);
static toIr(history: readonly Content[], tokenCalculator: ContextTokenCalculator): Episode[] {
return toIr(history, tokenCalculator);
}
/**
packages/core/src/context/ir/toIr.ts
+20 -23
View File
@@ -16,7 +16,7 @@ import type {
UserPrompt,
SystemEvent,
} from './types.js';
import { estimateContextTokenCountSync } from '../utils/contextTokenCalculator.js';
import type { ContextTokenCalculator } from '../utils/contextTokenCalculator.js';
// WeakMap to provide stable, deterministic identity across parses for the exact same Content/Part references
const nodeIdentityMap = new WeakMap<object, string>();
@@ -30,26 +30,20 @@ export function getStableId(obj: object): string {
return id;
}
export let charsPerTokenConfig: { charsPerToken?: number } | undefined;
export function setMapperConfig(cfg: { charsPerToken?: number }) {
charsPerTokenConfig = cfg;
}
export function createMetadata(parts: Part[]): IrMetadata {
const tokens = estimateContextTokenCountSync(parts, 0, charsPerTokenConfig);
return {
originalTokens: tokens,
currentTokens: tokens,
transformations: [],
};
}
export function toIr(history: readonly Content[]): Episode[] {
export function toIr(history: readonly Content[], tokenCalculator: ContextTokenCalculator): Episode[] {
const episodes: Episode[] = [];
let currentEpisode: Partial<Episode> | null = null;
const pendingCallParts: Map<string, Part> = new Map();
const createMetadata = (parts: Part[]): IrMetadata => {
const tokens = tokenCalculator.estimateTokensForParts(parts, 0);
return {
originalTokens: tokens,
currentTokens: tokens,
transformations: [],
};
};
const finalizeEpisode = () => {
if (currentEpisode && currentEpisode.trigger) {
episodes.push(currentEpisode as unknown as Episode); // eslint-disable-line @typescript-eslint/no-unsafe-type-assertion
@@ -67,15 +61,15 @@ export function toIr(history: readonly Content[]): Episode[] {
);
if (hasToolResponses) {
currentEpisode = parseToolResponses(msg, currentEpisode, pendingCallParts);
currentEpisode = parseToolResponses(msg, currentEpisode, pendingCallParts, tokenCalculator, createMetadata);
}
if (hasUserParts) {
finalizeEpisode();
currentEpisode = parseUserParts(msg);
currentEpisode = parseUserParts(msg, createMetadata);
}
} else if (msg.role === 'model') {
currentEpisode = parseModelParts(msg, currentEpisode, pendingCallParts);
currentEpisode = parseModelParts(msg, currentEpisode, pendingCallParts, createMetadata);
}
}
@@ -91,6 +85,8 @@ function parseToolResponses(
msg: Content,
currentEpisode: Partial<Episode> | null,
pendingCallParts: Map<string, Part>,
tokenCalculator: ContextTokenCalculator,
createMetadata: (parts: Part[]) => IrMetadata
): Partial<Episode> {
if (!currentEpisode) {
currentEpisode = {
@@ -113,9 +109,9 @@ function parseToolResponses(
const matchingCall = pendingCallParts.get(callId);
const intentTokens = matchingCall
? estimateContextTokenCountSync([matchingCall])
? tokenCalculator.estimateTokensForParts([matchingCall])
: 0;
const obsTokens = estimateContextTokenCountSync([part]);
const obsTokens = tokenCalculator.estimateTokensForParts([part]);
const step: ToolExecution = {
id: getStableId(part),
@@ -150,7 +146,7 @@ function parseToolResponses(
return currentEpisode;
}
function parseUserParts(msg: Content): Partial<Episode> {
function parseUserParts(msg: Content, createMetadata: (parts: Part[]) => IrMetadata): Partial<Episode> {
const semanticParts: SemanticPart[] = [];
for (const p of msg.parts!) {
if (p.text !== undefined)
@@ -192,6 +188,7 @@ function parseModelParts(
msg: Content,
currentEpisode: Partial<Episode> | null,
pendingCallParts: Map<string, Part>,
createMetadata: (parts: Part[]) => IrMetadata
): Partial<Episode> {
if (!currentEpisode) {
currentEpisode = {
packages/core/src/context/processors/semanticCompressionProcessor.ts
+9 -17
View File
@@ -10,7 +10,7 @@ import type { ContextEnvironment } from '../sidecar/environment.js';
import { debugLogger } from '../../utils/debugLogger.js';
import { LlmRole } from '../../telemetry/types.js';
import { getResponseText } from '../../utils/partUtils.js';
import { estimateTokenCountSync } from '../../utils/tokenCalculation.js';
import { estimateContextTokenCountSync } from '../utils/contextTokenCalculator.js';
export class SemanticCompressionProcessor implements ContextProcessor {
readonly name = 'SemanticCompression';
@@ -30,22 +30,14 @@ export class SemanticCompressionProcessor implements ContextProcessor {
episodes: Episode[],
state: ContextAccountingState,
): Promise<Episode[]> {
require('fs').appendFileSync(
'/tmp/debug2.json',
'SEMANTIC PROCESS: First episode ID: ' +
episodes[0]?.id +
'\nProtected IDs: ' +
Array.from(state.protectedEpisodeIds).join(', ') +
'\n',
);
// If the budget is satisfied, or semantic compression isn't enabled
if (state.isBudgetSatisfied) {
return episodes;
}
const semanticConfig = this.options;
// We estimate 4 chars per token for truncation logic
const thresholdChars = semanticConfig.nodeThresholdTokens * 4;
const limitTokens = semanticConfig.nodeThresholdTokens;
const thresholdChars = limitTokens * this.env.charsPerToken;
this.modelToUse = 'gemini-2.5-flash';
let currentDeficit = state.deficitTokens;
@@ -70,8 +62,8 @@ export class SemanticCompressionProcessor implements ContextProcessor {
part.text,
'User Prompt',
);
const newTokens = estimateTokenCountSync([{ text: summary }]);
const oldTokens = estimateTokenCountSync([{ text: part.text }]);
const newTokens = estimateContextTokenCountSync([{ text: summary }], 0, { charsPerToken: this.env.charsPerToken });
const oldTokens = estimateContextTokenCountSync([{ text: part.text }], 0, { charsPerToken: this.env.charsPerToken });
if (newTokens < oldTokens) {
part.presentation = { text: summary, tokens: newTokens };
@@ -96,8 +88,8 @@ export class SemanticCompressionProcessor implements ContextProcessor {
step.text,
'Agent Thought',
);
const newTokens = estimateTokenCountSync([{ text: summary }]);
const oldTokens = estimateTokenCountSync([{ text: step.text }]);
const newTokens = estimateContextTokenCountSync([{ text: summary }], 0, { charsPerToken: this.env.charsPerToken });
const oldTokens = estimateContextTokenCountSync([{ text: step.text }], 0, { charsPerToken: this.env.charsPerToken });
if (newTokens < oldTokens) {
step.presentation = { text: summary, tokens: newTokens };
@@ -138,7 +130,7 @@ export class SemanticCompressionProcessor implements ContextProcessor {
// Wrap the summary in an object so the Gemini API accepts it as a valid functionResponse.response
const newObsObject = { summary };
const newObsTokens = estimateTokenCountSync([
const newObsTokens = estimateContextTokenCountSync([
{
functionResponse: {
name: step.toolName,
@@ -146,7 +138,7 @@ export class SemanticCompressionProcessor implements ContextProcessor {
id: step.id,
},
},
]);
], 0, { charsPerToken: this.env.charsPerToken });
const oldObsTokens =
step.presentation?.tokens.observation ?? step.tokens.observation;
packages/core/src/context/processors/toolMaskingProcessor.ts
+2 -19
View File
@@ -6,7 +6,7 @@
import type { ContextAccountingState, ContextProcessor } from '../pipeline.js';
import type { ContextEnvironment } from '../sidecar/environment.js';
import { estimateTokenCountSync } from '../../utils/tokenCalculation.js';
import { estimateContextTokenCountSync } from '../utils/contextTokenCalculator.js';
import { sanitizeFilenamePart } from '../../utils/fileUtils.js';
import * as fsPromises from 'node:fs/promises';
import path from 'node:path';
@@ -50,7 +50,7 @@ export class ToolMaskingProcessor implements ContextProcessor {
const newEpisodes = [...episodes];
let currentDeficit = state.deficitTokens;
const limitChars = maskingConfig.stringLengthThresholdTokens * 4;
const limitChars = maskingConfig.stringLengthThresholdTokens * this.env.charsPerToken;
let toolOutputsDir = path.join(
this.env.projectTempDir,
@@ -125,23 +125,6 @@ export class ToolMaskingProcessor implements ContextProcessor {
nodeType: string,
): Promise<{ masked: any; changed: boolean }> => {
if (typeof obj === 'string') {
require('fs').appendFileSync(
'/tmp/debug.json',
'STRING FOUND. length: ' +
obj.length +
' limitChars: ' +
limitChars +
'\n',
);
if (obj.length > 1000)
console.log(
'Found string of length:',
obj.length,
'limitChars is:',
limitChars,
'isAlreadyMasked:',
this.isAlreadyMasked(obj),
);
if (obj.length > limitChars && !this.isAlreadyMasked(obj)) {
const newString = await handleMasking(
obj,
packages/core/src/context/sidecar/environment.ts
+3 -1
View File
@@ -6,6 +6,7 @@
import type { BaseLlmClient } from '../../core/baseLlmClient.js';
import type { ContextTracer } from '../tracer.js';
import type { ContextEventBus } from '../eventBus.js';
import type { ContextTokenCalculator } from '../utils/contextTokenCalculator.js';
export type { ContextTracer, ContextEventBus };
export interface ContextEnvironment {
@@ -16,6 +17,7 @@
readonly projectTempDir: string;
readonly tracer: ContextTracer;
readonly charsPerToken: number;
readonly tokenCalculator: ContextTokenCalculator;
readonly eventBus: ContextEventBus;
eventBus: ContextEventBus;
}
packages/core/src/context/sidecar/environmentImpl.ts
+7 -1
View File
@@ -10,7 +10,11 @@ import type { ContextEnvironment } from './environment.js';
import type { ContextEventBus } from '../eventBus.js';
import { ContextTokenCalculator } from '../utils/contextTokenCalculator.js';
export class ContextEnvironmentImpl implements ContextEnvironment {
public readonly tokenCalculator: ContextTokenCalculator;
constructor(
public readonly llmClient: BaseLlmClient,
public readonly sessionId: string,
@@ -20,5 +24,7 @@ export class ContextEnvironmentImpl implements ContextEnvironment {
public readonly tracer: ContextTracer,
public readonly charsPerToken: number,
public readonly eventBus: ContextEventBus,
) {}
) {
this.tokenCalculator = new ContextTokenCalculator(this.charsPerToken);
}
}
packages/core/src/context/testing/contextTestUtils.ts
+2
View File
@@ -26,6 +26,7 @@ export function createMockEnvironment(): ContextEnvironment {
eventBus: new ContextEventBus(),
tracer: new ContextTracer('/tmp', 'mock-session'),
charsPerToken: 1,
tokenCalculator: new ContextTokenCalculator(1),
};
}
@@ -89,6 +90,7 @@ import { ContextTracer } from '../tracer.js';
import { ContextEnvironmentImpl } from '../sidecar/environmentImpl.js';
import { SidecarLoader } from '../sidecar/SidecarLoader.js';
import { ContextEventBus } from '../eventBus.js';
import { ContextTokenCalculator } from '../utils/contextTokenCalculator.js';
import type { BaseLlmClient } from 'src/core/baseLlmClient.js';
export function setupContextComponentTest(config: Config) {
packages/core/src/context/utils/contextTokenCalculator.ts
+51 -28
View File
@@ -8,37 +8,60 @@ import type { Part } from '@google/genai';
import { estimateTokenCountSync as baseEstimate } from '../../utils/tokenCalculation.js';
import type { Episode } from '../ir/types.js';
export function calculateEpisodeListTokens(episodes: Episode[]): number {
let tokens = 0;
for (const ep of episodes) {
if (ep.trigger) tokens += ep.trigger.metadata.currentTokens;
for (const step of ep.steps) {
tokens += step.metadata.currentTokens;
}
if (ep.yield) tokens += ep.yield.metadata.currentTokens;
}
return tokens;
}
export class ContextTokenCalculator {
constructor(private readonly charsPerToken: number) {}
export function estimateContextTokenCountSync(
parts: Part[],
depth: number = 0,
config?: { charsPerToken?: number },
): number {
if (config?.charsPerToken !== undefined && config.charsPerToken !== 4) {
let totalTokens = 0;
for (const part of parts) {
if (typeof part.text === 'string') {
totalTokens += Math.ceil(part.text.length / config.charsPerToken);
} else {
totalTokens += Math.ceil(
JSON.stringify(part).length / config.charsPerToken,
);
/**
* Fast, simple heuristic estimation for a raw string.
*/
estimateTokensForString(text: string): number {
return Math.ceil(text.length / this.charsPerToken);
}
/**
* Fast, simple heuristic conversion from tokens to expected character length.
* Useful for calculating truncation thresholds.
*/
tokensToChars(tokens: number): number {
return tokens * this.charsPerToken;
}
/**
* Calculates the total token count for a complete Episodic IR graph.
* This is fast because it relies on pre-computed metadata where available.
*/
calculateEpisodeListTokens(episodes: Episode[]): number {
let tokens = 0;
for (const ep of episodes) {
if (ep.trigger) tokens += ep.trigger.metadata.currentTokens;
for (const step of ep.steps) {
tokens += step.metadata.currentTokens;
}
if (ep.yield) tokens += ep.yield.metadata.currentTokens;
}
return totalTokens;
return tokens;
}
// The baseEstimate no longer accepts config because we forked it!
return baseEstimate(parts, depth);
/**
* Slower, precise estimation for a Gemini Content/Part graph.
* Deeply inspects the nested structure and uses the base tokenization math.
*/
estimateTokensForParts(parts: Part[], depth: number = 0): number {
if (this.charsPerToken !== 4) {
let totalTokens = 0;
for (const part of parts) {
if (typeof part.text === 'string') {
totalTokens += Math.ceil(part.text.length / this.charsPerToken);
} else {
totalTokens += Math.ceil(
JSON.stringify(part).length / this.charsPerToken,
);
}
}
return totalTokens;
}
// The baseEstimate no longer accepts config because we forked it!
return baseEstimate(parts, depth);
}
}