refactor(core): shrink the content generator interface (#9676)

* refactor(core): shrink content generator interface

* refactor(core): remove orphaned request-tokenizer estimator cluster

Removing countTokens from both providers deleted the last production
consumers of RequestTokenEstimator. Delete the orphaned cluster:
requestTokenizer.ts (330), imageTokenizer.ts (534), types.ts (36), the
directory barrel (11), and both test files (608 lines). Also drop the
inert vi.mock of requestTokenizer.js left in client.test.ts and the
stale dimension-extractor cross-reference in review/lib/assets.ts.

textTokenizer.ts and supportedImageFormats.ts stay: converter.ts, pdf.ts,
and fileUtils.ts still consume them and the core barrel re-exports them.

* docs(design): sync lazy-google-genai-loading record with shrunk interface

countTokens and useSummarizedThinking no longer exist on ContentGenerator,
so the design record for the lazy-wrapper architecture must not keep
advertising them: list the three remaining shared async operations, drop
the useSummarizedThinking sentence and the summarized-thinking item from
the consumer audit and Verification section, and add a dated note
recording the interface shrink from PR #9676.

* ci: record cd-cua-driver.yml size growth in .size-baseline

Same latent main-side violation as fixed in #9682: #9587 grew the
workflow without a baseline update; record the new size as the check
message directs (precedent #9747).

* docs: finish scrubbing tokenizer references after estimator-cluster removal

Follow-up to 0ee17632c7/1871bb5b81 (review round 2):
- supportedImageFormats.ts header and getSupportedImageFormatsString doc
  no longer describe a tokenizer decode/metadata-extraction stage; the
  list is now documented as the vision-input acceptance list, with token
  accounting noted as the flat DEFAULT_IMAGE_TOKEN_ESTIMATE.
- web-shell-image-drag-and-drop.md's BMP rationale no longer claims
  ImageTokenizer parses BMP dimensions; dated sync note added stating
  BMP support rests on SUPPORTED_IMAGE_MIME_TYPES plus converter
  passthrough since PR #9676.

* docs: drop tokenizer from the BMP test-plan line

Follow-up to 18f08c0924: the test plan still required converter/tokenizer
focused tests for image paths; the image-tokenizer estimator cluster was
removed in PR #9676 (text tokenizer is unaffected and out of scope here).
This commit is contained in:
Dragon authored and GitHub committed 2026-08-24 08:30:18 +00:00
1 parent 6420d69d14
commit 43d46be912
33 files changed
+81 -1972

No files matched your search

@@ -22,14 +22,21 @@ Provider implementations continue to use the official SDK classes. In particular
`createContentGenerator()` still validates configuration, preloads the runtime fetch implementation, and performs Qwen OAuth credential acquisition at its current point in the session lifecycle. It returns a private lazy `ContentGenerator` whose memoized loader constructs the selected provider and wraps it in `LoggingContentGenerator` on the first asynchronous content-generator operation.
All four asynchronous operations share the same loader promise:
All three asynchronous operations share the same loader promise:
- `generateContent`
- `generateContentStream`
- `countTokens`
- `embedContent`
Concurrent first calls therefore import and construct the provider once. `useSummarizedThinking()` remains synchronous and is supplied from the selected provider's known behavior: true for Gemini/Vertex and false for OpenAI, Qwen OAuth, and Anthropic.
Concurrent first calls therefore import and construct the provider once.
> **Update (2026-08, PR #9676)**: `countTokens` and `useSummarizedThinking`
> were removed from the `ContentGenerator` interface — no production caller
> used either, and the removal narrows every provider and test double to the
> operations something actually calls. This doc originally listed four shared
> asynchronous operations (including `countTokens`) and a synchronous
> `useSummarizedThinking()` supplied from each provider's known behavior; both
> are gone from the interface, the lazy wrapper, and the four providers.
Qwen OAuth credential acquisition remains eager within `createContentGenerator()`. An expired or missing cached credential therefore continues to reject ACP session creation rather than producing an apparently usable session that fails only on its first prompt.
@@ -49,7 +56,7 @@ The serve fast-path metafile guard adds `@google/genai` to the ACP forbidden-pac
There are three direct production creation paths. `Config.refreshAuth()` owns the main-session generator. `BaseLlmClient` owns cached per-model generators for routed side requests. `createRuntimeContentGeneratorView()` owns dedicated generators used by the in-process agent backend, subagent manager, and forked agents. Each path stores and consumes only the `ContentGenerator` interface, so the private lazy wrapper preserves its ownership and routing boundary.
The interface consumers call only `generateContent`, `generateContentStream`, `countTokens`, `embedContent`, and `useSummarizedThinking`. The main chat path, prompt hooks, memory/goal/side queries, vision routing, subagents, and session resume do not inspect the concrete provider or unwrap `LoggingContentGenerator`; a repository-wide search found no production `instanceof` or `getWrapped()` caller. MCP tool discovery is separate from generator ownership and keeps the SDK-provided `mcpToTool` adapter behind its own first-use import.
The interface consumers call only `generateContent`, `generateContentStream`, and `embedContent`. The main chat path, prompt hooks, memory/goal/side queries, vision routing, subagents, and session resume do not inspect the concrete provider or unwrap `LoggingContentGenerator`; a repository-wide search found no production `instanceof` or `getWrapped()` caller. MCP tool discovery is separate from generator ownership and keeps the SDK-provided `mcpToTool` adapter behind its own first-use import.
## Alternatives rejected
@@ -71,7 +78,7 @@ The interface consumers call only `generateContent`, `generateContentStream`, `c
## Verification
Unit tests cover helper parity, deferred construction, Qwen credential timing, single-flight behavior, provider-specific summarized-thinking values, deferred module failures, and MCP discovery behavior. The bundled metafile must show `@google/genai` absent from the ACP static closure while retaining it in dynamic provider/MCP chunks.
Unit tests cover helper parity, deferred construction, Qwen credential timing, single-flight behavior, deferred module failures, and MCP discovery behavior. The bundled metafile must show `@google/genai` absent from the ACP static closure while retaining it in dynamic provider/MCP chunks.
The 2C4G acceptance run follows #7264: 30 paired serial cold starts, `channel.initialize` P50/P95, process-to-first-session, preheated/warm behavior, concurrent first sessions, telemetry on/off, and peak RSS. Because this change moves work later, it additionally records session-response-to-first-token and process-to-first-token for an immediate first prompt. A startup win that is fully repaid as a first-token regression is reported rather than treated as a successful optimization.
@@ -502,9 +502,14 @@ oversized placeholder。它们验证既有服务端契约;Web Shell helper 测
候选顺序和 encoded-data 剩余预算。
BMP 以 `image/bmp` 进入缩略图 data URL 和 daemon image block。Core 的
`SUPPORTED_IMAGE_MIME_TYPES` 明确包含 `image/bmp`,`ImageTokenizer` 解析 BMP 尺寸,
OpenAI converter 把启用 image modality 的 `inlineData` 原样构造成
`data:image/bmp;base64,...`;Gemini 路径保留相同 `inlineData`。因此 V1 不在浏览器转码。
`SUPPORTED_IMAGE_MIME_TYPES` 明确包含 `image/bmp`,OpenAI converter 把启用 image modality 的
`inlineData` 原样构造成 `data:image/bmp;base64,...`;Gemini 路径保留相同 `inlineData`。
因此 V1 不在浏览器转码。
> **2026-08-24 同步注记(PR #9676)**:request-tokenizer 估计器簇(含 `ImageTokenizer`
> 及其 BMP 尺寸解析)已作为孤儿代码删除。BMP 支持现在仅依赖 `SUPPORTED_IMAGE_MIME_TYPES`
> 接受清单与 converter 透传;token 计数使用 `compactionInputSlimming.ts` 中的固定
> `DEFAULT_IMAGE_TOKEN_ESTIMATE`。下文对 BMP 路径的 E2E/人工验收要求不变。
浏览器若不能解码缩略图,不影响附件数据传输,但 E2E 必须覆盖 Chromium 解码,
Firefox/Linux 必须完成人工验收。
@@ -634,8 +639,8 @@ Web Shell 输入层静默改变格式。
BMP 的下游回归不只停在 mock HTTP 入参:在既有 ACP session prompt 转换测试中加入
`image/bmp`,验证最终 Core canonical content 保持
`inlineData.mimeType === 'image/bmp'`;OpenAI/Gemini converter/tokenizer 聚焦测试验证
各自既有图片路径,Anthropic 聚焦测试明确断言 BMP 转为 unsupported-media 文本。daemon
`inlineData.mimeType === 'image/bmp'`;OpenAI/Gemini converter 聚焦测试验证
各自既有图片路径(tokenizer 估计器簇已随 PR #9676 删除,见上文同步注记),Anthropic 聚焦测试明确断言 BMP 转为 unsupported-media 文本。daemon
已有结构化 `413` 测试,Core 已有 inline-media within/over limit 测试;本功能不复制
production 限制,只确认 Web Shell 对这些既有失败语义的状态保留。
@@ -70,9 +70,8 @@ export const ASSET_HEADER_BYTES = 16;
*
* A sibling signature table lives in core: `sniffFileKind` in
* `packages/core/src/utils/binary-content.ts` (best-effort kind detection for
* fetched web content, deliberately looser) and the dimension extractors in
* `packages/core/src/utils/request-tokenizer/imageTokenizer.ts`. Admitting or
* correcting a format here means checking those sites too.
* fetched web content, deliberately looser). Admitting or correcting a format
* here means checking that site too.
*/
export function sniffImageFormat(header: Uint8Array): ImageFormat | null {
const at = (i: number): number => header[i] ?? -1;
+3 -3
View File
@@ -5080,9 +5080,9 @@ describe('AppContainer State Management', () => {
lastCompletedUuid: 'a1',
};
vi.spyOn(mockConfig, 'getContentGenerator').mockReturnValue({
useSummarizedThinking: vi.fn(() => false),
} as unknown as ReturnType<typeof mockConfig.getContentGenerator>);
vi.spyOn(mockConfig, 'getContentGenerator').mockReturnValue(
{} as unknown as ReturnType<typeof mockConfig.getContentGenerator>,
);
vi.spyOn(mockConfig, 'initialize').mockResolvedValue(undefined);
vi.spyOn(mockConfig, 'getResumedSessionData').mockReturnValue(
resumeSessionData as ReturnType<
@@ -701,7 +701,6 @@ describe('SessionPicker', () => {
getIdeMode: () => false,
isTrustedFolder: () => false,
getToolRegistry: () => ({ getTool: () => undefined }),
getContentGenerator: () => ({ useSummarizedThinking: () => false }),
} as unknown as Config;
const PREVIEW_SETTINGS_STUB = {
merged: { ui: {} },
@@ -37,7 +37,6 @@ const PREVIEW_CONFIG_STUB = {
getIdeMode: () => false,
isTrustedFolder: () => false,
getToolRegistry: () => ({ getTool: () => undefined }),
getContentGenerator: () => ({ useSummarizedThinking: () => false }),
} as unknown as Config;
const PREVIEW_SETTINGS_STUB = {
@@ -74,9 +74,7 @@ vi.mock('../../core/contentGenerator.js', async (importOriginal) => {
createContentGenerator: vi.fn().mockResolvedValue({
generateContent: vi.fn(),
generateContentStream: vi.fn(),
countTokens: vi.fn().mockResolvedValue({ totalTokens: 100 }),
embedContent: vi.fn(),
useSummarizedThinking: vi.fn().mockReturnValue(false),
}),
createContentGeneratorConfig: vi.fn().mockReturnValue({
model: DEFAULT_QWEN_MODEL,
-4
View File
@@ -5409,7 +5409,6 @@ describe('Server Config (config.ts)', () => {
vi.mocked(createContentGenerator).mockResolvedValue({
generateContent: vi.fn(),
generateContentStream: vi.fn(),
countTokens: vi.fn(),
embedContent: vi.fn(),
} as unknown as ContentGenerator);
@@ -5460,7 +5459,6 @@ describe('Server Config (config.ts)', () => {
vi.mocked(createContentGenerator).mockResolvedValue({
generateContent: vi.fn(),
generateContentStream: vi.fn(),
countTokens: vi.fn(),
embedContent: vi.fn(),
} as unknown as ContentGenerator);
@@ -5494,7 +5492,6 @@ describe('Server Config (config.ts)', () => {
vi.mocked(createContentGenerator).mockResolvedValue({
generateContent: vi.fn(),
generateContentStream: vi.fn(),
countTokens: vi.fn(),
embedContent: vi.fn(),
} as unknown as ContentGenerator);
@@ -6051,7 +6048,6 @@ describe('Server Config (config.ts)', () => {
vi.mocked(createContentGenerator).mockResolvedValue({
generateContent: vi.fn(),
generateContentStream: vi.fn(),
countTokens: vi.fn(),
embedContent: vi.fn(),
} as unknown as ContentGenerator);
@@ -363,24 +363,5 @@ describe('OpenAIContentGenerator Timeout Handling', () => {
generator.generateContent(request, 'test-prompt-id'),
).rejects.toThrow(/Request timeout after \d+s/);
});
it('should fall back to character-based estimation if countTokens fails', async () => {
const timeoutError = new Error('Request timeout');
mockOpenAIClient.chat.completions.create.mockRejectedValue(timeoutError);
// Mock countTokens to throw error
const mockCountTokens = vi.spyOn(generator, 'countTokens');
mockCountTokens.mockRejectedValue(new Error('Count tokens failed'));
const request = {
contents: [{ role: 'user' as const, parts: [{ text: 'Hello world' }] }],
model: 'gpt-4',
};
// Should not throw due to token counting failure
await expect(
generator.generateContent(request, 'test-prompt-id'),
).rejects.toThrow(/Request timeout after \d+s/);
});
});
});
@@ -6,10 +6,7 @@
import { getEventListeners } from 'node:events';
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
import type {
CountTokensParameters,
GenerateContentParameters,
} from '@google/genai';
import type { GenerateContentParameters } from '@google/genai';
import { FinishReason, GenerateContentResponse } from '@google/genai';
import type { ContentGeneratorConfig } from '../contentGenerator.js';
import {
@@ -17,19 +14,11 @@ import {
DISABLED_REQUEST_TIMEOUT_MS,
} from '../openaiContentGenerator/constants.js';
// Mock the request tokenizer module BEFORE importing the class that uses it.
const mockTokenizer = {
calculateTokens: vi.fn(),
dispose: vi.fn(),
};
const mockReportAnthropicRequest = vi.hoisted(() => vi.fn());
const mockReportAnthropicFollowingRequest = vi.hoisted(() => vi.fn());
const mockReportAnthropicResponse = vi.hoisted(() => vi.fn());
const mockReportAnthropicEvent = vi.hoisted(() => vi.fn());
vi.mock('../../utils/request-tokenizer/index.js', () => ({
RequestTokenEstimator: vi.fn(() => mockTokenizer),
}));
vi.mock('../../telemetry/gen-ai-request.js', () => ({
reportAnthropicRequest: mockReportAnthropicRequest,
reportAnthropicFollowingRequest: mockReportAnthropicFollowingRequest,
@@ -107,16 +96,6 @@ describe('AnthropicContentGenerator', () => {
savedMaxOutputTokensEnv = process.env[MAX_OUTPUT_TOKENS_ENV];
delete process.env[MAX_OUTPUT_TOKENS_ENV];
mockTokenizer.calculateTokens.mockResolvedValue({
totalTokens: 50,
breakdown: {
textTokens: 50,
imageTokens: 0,
audioTokens: 0,
otherTokens: 0,
},
processingTime: 1,
});
anthropicState = anthropicMockState;
anthropicState.createImpl.mockReset();
@@ -3305,58 +3284,6 @@ describe('AnthropicContentGenerator', () => {
});
});
describe('countTokens', () => {
it('counts tokens using the request tokenizer', async () => {
const { AnthropicContentGenerator } = await importGenerator();
const generator = new AnthropicContentGenerator(
{
model: 'claude-test',
apiKey: 'test-key',
timeout: 10_000,
maxRetries: 2,
samplingParams: {},
schemaCompliance: 'auto',
},
mockConfig,
);
const request: CountTokensParameters = {
contents: [{ role: 'user', parts: [{ text: 'Hello world' }] }],
model: 'claude-test',
};
const result = await generator.countTokens(request);
expect(mockTokenizer.calculateTokens).toHaveBeenCalledWith(request);
expect(result.totalTokens).toBe(50);
});
it('falls back to character approximation when tokenizer throws', async () => {
const { AnthropicContentGenerator } = await importGenerator();
mockTokenizer.calculateTokens.mockRejectedValueOnce(new Error('boom'));
const generator = new AnthropicContentGenerator(
{
model: 'claude-test',
apiKey: 'test-key',
timeout: 10_000,
maxRetries: 2,
samplingParams: {},
schemaCompliance: 'auto',
},
mockConfig,
);
const request: CountTokensParameters = {
contents: [{ role: 'user', parts: [{ text: 'Hello' }] }],
model: 'claude-test',
};
const content = JSON.stringify(request.contents);
const expected = Math.ceil(content.length / 4);
const result = await generator.countTokens(request);
expect(result.totalTokens).toBe(expected);
});
});
describe('generateContentStream', () => {
it('emits tool preparation metadata before the complete function call', async () => {
const { AnthropicContentGenerator } = await importGenerator();
@@ -6,8 +6,6 @@
import Anthropic from '@anthropic-ai/sdk';
import type {
CountTokensParameters,
CountTokensResponse,
EmbedContentParameters,
EmbedContentResponse,
GenerateContentParameters,
@@ -29,7 +27,6 @@ type MessageCreateParamsNonStreaming =
Anthropic.MessageCreateParamsNonStreaming;
type MessageCreateParamsStreaming = Anthropic.MessageCreateParamsStreaming;
type RawMessageStreamEvent = Anthropic.RawMessageStreamEvent;
import { RequestTokenEstimator } from '../../utils/request-tokenizer/index.js';
import { safeJsonParse } from '../../utils/safeJsonParse.js';
import { AnthropicContentConverter } from './converter.js';
import { buildAnthropicUsageMetadata } from './usage.js';
@@ -445,41 +442,11 @@ export class AnthropicContentGenerator implements ContentGenerator {
return drainThenCleanup();
}
async countTokens(
request: CountTokensParameters,
): Promise<CountTokensResponse> {
try {
const estimator = new RequestTokenEstimator();
const result = await estimator.calculateTokens(request);
return {
totalTokens: result.totalTokens,
};
} catch (error) {
debugLogger.warn(
'Failed to calculate tokens with tokenizer, ' +
'falling back to simple method:',
error,
);
const content = JSON.stringify(request.contents);
const totalTokens = Math.ceil(content.length / 4);
return {
totalTokens,
};
}
}
async embedContent(
_request: EmbedContentParameters,
): Promise<EmbedContentResponse> {
throw new Error('Anthropic does not support embeddings.');
}
useSummarizedThinking(): boolean {
return false;
}
private buildHeaders(useProxyIdentity: boolean): Record<string, string> {
// Beta headers are computed per-request in buildPerRequestHeaders so they
// stay in sync with what the request body actually carries — see #3788
-21
View File
@@ -425,26 +425,6 @@ vi.mock(
);
import { microcompactHistory } from '../services/microcompaction/microcompact.js';
// Mock RequestTokenizer to use simple character-based estimation
vi.mock('../utils/request-tokenizer/requestTokenizer.js', () => ({
RequestTokenizer: class {
async calculateTokens(request: { contents: unknown }) {
// Simple estimation: count characters in JSON and divide by 4
const totalChars = JSON.stringify(request.contents).length;
return {
totalTokens: Math.floor(totalChars / 4),
breakdown: {
textTokens: Math.floor(totalChars / 4),
imageTokens: 0,
audioTokens: 0,
otherTokens: 0,
},
processingTime: 0,
};
}
},
}));
/**
* Array.fromAsync ponyfill, which will be available in es 2024.
*
@@ -563,7 +543,6 @@ describe('Gemini Client (client.ts)', () => {
generateContent: mockGenerateContentFn,
generateContentStream: vi.fn(),
batchEmbedContents: vi.fn(),
countTokens: vi.fn().mockResolvedValue({ totalTokens: 100 }),
} as unknown as ContentGenerator;
// Because the GeminiClient constructor kicks off an async process (startChat)
+24 -37
View File
@@ -56,9 +56,7 @@ vi.mock('./openaiContentGenerator/index.js', () => ({
(async function* () {
yield {};
})(),
countTokens: async () => ({ totalTokens: 1 }),
embedContent: async () => ({ embeddings: [] }),
useSummarizedThinking: () => false,
});
const gate = openaiMockState.constructionGates[attempt];
if (!gate) {
@@ -89,12 +87,8 @@ vi.mock('../qwen/qwenContentGenerator.js', () => ({
qwenMockState.constructorModels.push(generatorConfig.model);
}
async countTokens() {
return { totalTokens: 1 };
}
useSummarizedThinking() {
return false;
async embedContent() {
return { embeddings: [] };
}
},
}));
@@ -132,7 +126,7 @@ describe('createContentGenerator', () => {
const mockGenerator = {
models: {
countTokens: vi.fn().mockResolvedValue({ totalTokens: 1 }),
embedContent: vi.fn().mockResolvedValue({ embeddings: [] }),
},
} as unknown as GoogleGenAI;
vi.mocked(GoogleGenAI).mockImplementation(() => mockGenerator as never);
@@ -145,9 +139,8 @@ describe('createContentGenerator', () => {
mockConfig,
);
expect(GoogleGenAI).not.toHaveBeenCalled();
expect(generator.useSummarizedThinking()).toBe(true);
await generator.countTokens({
await generator.embedContent({
model: 'test-model',
contents: 'hello',
});
@@ -174,7 +167,7 @@ describe('createContentGenerator', () => {
} as unknown as Config;
const mockGenerator = {
models: {
countTokens: vi.fn().mockResolvedValue({ totalTokens: 1 }),
embedContent: vi.fn().mockResolvedValue({ embeddings: [] }),
},
} as unknown as GoogleGenAI;
vi.mocked(GoogleGenAI).mockImplementation(() => mockGenerator as never);
@@ -187,7 +180,7 @@ describe('createContentGenerator', () => {
mockConfig,
);
expect(GoogleGenAI).not.toHaveBeenCalled();
await generator.countTokens({
await generator.embedContent({
model: 'test-model',
contents: 'hello',
});
@@ -220,10 +213,9 @@ describe('createContentGenerator', () => {
);
expect(openaiMockState.createCount).toBe(0);
expect(generator.useSummarizedThinking()).toBe(false);
await Promise.all([
generator.countTokens({ model: 'test-model', contents: 'one' }),
generator.countTokens({ model: 'test-model', contents: 'two' }),
generator.embedContent({ model: 'test-model', contents: 'one' }),
generator.embedContent({ model: 'test-model', contents: 'two' }),
]);
expect(openaiMockState.createCount).toBe(1);
});
@@ -232,18 +224,14 @@ describe('createContentGenerator', () => {
const generator: ContentGenerator = {
generateContent: vi.fn(),
generateContentStream: vi.fn(),
countTokens: vi.fn(),
embedContent: vi.fn(),
useSummarizedThinking: vi.fn(),
};
await expect(preloadContentGenerator(generator)).resolves.toBeUndefined();
resetPreloadedContentGenerator(generator);
expect(generator.generateContent).not.toHaveBeenCalled();
expect(generator.generateContentStream).not.toHaveBeenCalled();
expect(generator.countTokens).not.toHaveBeenCalled();
expect(generator.embedContent).not.toHaveBeenCalled();
expect(generator.useSummarizedThinking).not.toHaveBeenCalled();
});
it('loads a provider once across concurrent preload and first use', async () => {
@@ -265,7 +253,7 @@ describe('createContentGenerator', () => {
await Promise.all([
preloadContentGenerator(generator),
generator.countTokens({ model: 'test-model', contents: 'hello' }),
generator.embedContent({ model: 'test-model', contents: 'hello' }),
]);
expect(openaiMockState.createCount).toBe(1);
@@ -290,7 +278,7 @@ describe('createContentGenerator', () => {
await preloadContentGenerator(generator);
resetPreloadedContentGenerator(generator);
await generator.countTokens({
await generator.embedContent({
model: 'test-model',
contents: 'hello',
});
@@ -316,12 +304,12 @@ describe('createContentGenerator', () => {
);
await preloadContentGenerator(generator);
await generator.countTokens({
await generator.embedContent({
model: 'test-model',
contents: 'first',
});
resetPreloadedContentGenerator(generator);
await generator.countTokens({
await generator.embedContent({
model: 'test-model',
contents: 'second',
});
@@ -354,7 +342,7 @@ describe('createContentGenerator', () => {
const preload = preloadContentGenerator(generator);
await vi.waitFor(() => expect(openaiMockState.createCount).toBe(1));
const firstUse = generator.countTokens({
const firstUse = generator.embedContent({
model: 'test-model',
contents: 'hello',
});
@@ -362,7 +350,7 @@ describe('createContentGenerator', () => {
releaseConstruction();
await expect(preload).resolves.toBeUndefined();
await expect(firstUse).resolves.toEqual({ totalTokens: 1 });
await expect(firstUse).resolves.toEqual({ embeddings: [] });
expect(openaiMockState.createCount).toBe(1);
});
@@ -400,22 +388,22 @@ describe('createContentGenerator', () => {
);
await vi.waitFor(() => expect(openaiMockState.createCount).toBe(1));
resetPreloadedContentGenerator(generator);
const firstUse = generator.countTokens({
const firstUse = generator.embedContent({
model: 'test-model',
contents: 'hello',
});
await vi.waitFor(() => expect(openaiMockState.createCount).toBe(2));
releaseFirstUse();
await expect(firstUse).resolves.toEqual({ totalTokens: 1 });
await expect(firstUse).resolves.toEqual({ embeddings: [] });
releasePreload();
await expect(discardedPreload).resolves.toBe(preloadError);
await expect(
generator.countTokens({
generator.embedContent({
model: 'test-model',
contents: 'still uses the replacement',
}),
).resolves.toEqual({ totalTokens: 1 });
).resolves.toEqual({ embeddings: [] });
expect(openaiMockState.createCount).toBe(2);
});
@@ -436,7 +424,7 @@ describe('createContentGenerator', () => {
await preloadContentGenerator(generator);
generatorConfig.model = 'coder-model';
resetPreloadedContentGenerator(generator);
await generator.countTokens({
await generator.embedContent({
model: 'coder-model',
contents: 'hello',
});
@@ -471,7 +459,7 @@ describe('createContentGenerator', () => {
await preloadContentGenerator(generator);
workingDir = '/workspace/after';
resetPreloadedContentGenerator(generator);
await generator.countTokens({
await generator.embedContent({
model: 'test-model',
contents: 'hello',
});
@@ -508,7 +496,7 @@ describe('createContentGenerator', () => {
(error: unknown) => error,
);
const firstUseError = await generator
.countTokens({ model: 'test-model', contents: 'hello' })
.embedContent({ model: 'test-model', contents: 'hello' })
.catch((error: unknown) => error);
expect(preloadError).toBeInstanceOf(Error);
@@ -535,8 +523,7 @@ describe('createContentGenerator', () => {
expect(qwenMockState.oauthCount).toBe(1);
expect(qwenMockState.constructorCount).toBe(0);
expect(generator.useSummarizedThinking()).toBe(false);
await generator.countTokens({ model: 'test-model', contents: 'hello' });
await generator.embedContent({ model: 'test-model', contents: 'hello' });
expect(qwenMockState.constructorCount).toBe(1);
});
@@ -639,7 +626,7 @@ describe('createContentGenerator - ERR_MODULE_NOT_FOUND handling', () => {
mockConfig,
);
await expect(
generator.countTokens({ model: 'test-model', contents: 'hello' }),
generator.embedContent({ model: 'test-model', contents: 'hello' }),
).rejects.toThrow('network timeout');
});
@@ -689,7 +676,7 @@ describe('createContentGenerator - ERR_MODULE_NOT_FOUND handling', () => {
},
mockConfig,
);
await generator.countTokens({ model: 'test-model', contents: 'hello' });
await generator.embedContent({ model: 'test-model', contents: 'hello' });
expect.unreachable('should have thrown');
} catch (error) {
expect(error).toBeInstanceOf(Error);
+17 -39
View File
@@ -5,8 +5,6 @@
*/
import type {
CountTokensParameters,
CountTokensResponse,
EmbedContentParameters,
EmbedContentResponse,
GenerateContentParameters,
@@ -33,7 +31,7 @@ import { preloadRuntimeFetchModule } from '../utils/runtimeFetchOptions.js';
import type { ReasoningEffort } from './reasoning-effort.js';
/**
* Interface abstracting the core functionalities for generating content and counting tokens.
* Interface abstracting the core content generation functionality.
*/
export interface ContentGenerator {
generateContent(
@@ -46,11 +44,7 @@ export interface ContentGenerator {
userPromptId: string,
): Promise<AsyncGenerator<GenerateContentResponse>>;
countTokens(request: CountTokensParameters): Promise<CountTokensResponse>;
embedContent(request: EmbedContentParameters): Promise<EmbedContentResponse>;
useSummarizedThinking(): boolean;
}
import { AuthType } from '../utils/auth-type.js';
@@ -434,10 +428,7 @@ class LazyContentGenerator implements ContentGenerator {
private generatorPromise?: Promise<ContentGenerator>;
private preloadedOnly = false;
constructor(
private readonly loader: () => Promise<ContentGenerator>,
private readonly summarizedThinking: boolean,
) {}
constructor(private readonly loader: () => Promise<ContentGenerator>) {}
private getGenerator(): Promise<ContentGenerator> {
this.generatorPromise ??= this.loader();
@@ -482,21 +473,11 @@ class LazyContentGenerator implements ContentGenerator {
);
}
async countTokens(
request: CountTokensParameters,
): Promise<CountTokensResponse> {
return (await this.getGeneratorForUse()).countTokens(request);
}
async embedContent(
request: EmbedContentParameters,
): Promise<EmbedContentResponse> {
return (await this.getGeneratorForUse()).embedContent(request);
}
useSummarizedThinking(): boolean {
return this.summarizedThinking;
}
}
/** @internal */
@@ -595,22 +576,19 @@ export async function createContentGenerator(
throw wrapProviderLoadError(error, authType);
}
return new LazyContentGenerator(
async () => {
try {
const [baseGenerator, { LoggingContentGenerator }] = await Promise.all([
loadBaseGenerator(),
import('./loggingContentGenerator/index.js'),
]);
return new LoggingContentGenerator(
baseGenerator,
config,
generatorConfig,
);
} catch (error) {
throw wrapProviderLoadError(error, authType);
}
},
authType === AuthType.USE_GEMINI || authType === AuthType.USE_VERTEX_AI,
);
return new LazyContentGenerator(async () => {
try {
const [baseGenerator, { LoggingContentGenerator }] = await Promise.all([
loadBaseGenerator(),
import('./loggingContentGenerator/index.js'),
]);
return new LoggingContentGenerator(
baseGenerator,
config,
generatorConfig,
);
} catch (error) {
throw wrapProviderLoadError(error, authType);
}
});
}
-14
View File
@@ -163,10 +163,8 @@ describe('GeminiChat', async () => {
mockContentGenerator = {
generateContent: vi.fn(),
generateContentStream: vi.fn(),
countTokens: vi.fn(),
embedContent: vi.fn(),
batchEmbedContents: vi.fn(),
useSummarizedThinking: vi.fn().mockReturnValue(false),
} as unknown as ContentGenerator;
// Default mock implementation for tests that don't care about retry logic
@@ -7202,10 +7200,8 @@ describe('GeminiChat', async () => {
({
generateContent: vi.fn(),
generateContentStream,
countTokens: vi.fn(),
embedContent: vi.fn(),
batchEmbedContents: vi.fn(),
useSummarizedThinking: vi.fn().mockReturnValue(false),
}) as unknown as ContentGenerator;
const resolveForModel = vi
.fn()
@@ -7449,10 +7445,8 @@ describe('GeminiChat', async () => {
({
generateContent: vi.fn(),
generateContentStream,
countTokens: vi.fn(),
embedContent: vi.fn(),
batchEmbedContents: vi.fn(),
useSummarizedThinking: vi.fn().mockReturnValue(false),
}) as unknown as ContentGenerator;
const resolveForModel = vi
.fn()
@@ -7541,10 +7535,8 @@ describe('GeminiChat', async () => {
const fallbackBGenerator = {
generateContent: vi.fn(),
generateContentStream: fallbackBGenerateContentStream,
countTokens: vi.fn(),
embedContent: vi.fn(),
batchEmbedContents: vi.fn(),
useSummarizedThinking: vi.fn().mockReturnValue(false),
} as unknown as ContentGenerator;
const resolveError = new Error('unknown fallback alias');
const resolveForModel = vi
@@ -7630,10 +7622,8 @@ describe('GeminiChat', async () => {
({
generateContent: vi.fn(),
generateContentStream,
countTokens: vi.fn(),
embedContent: vi.fn(),
batchEmbedContents: vi.fn(),
useSummarizedThinking: vi.fn().mockReturnValue(false),
}) as unknown as ContentGenerator;
const resolveForModel = vi
.fn()
@@ -7794,10 +7784,8 @@ describe('GeminiChat', async () => {
({
generateContent: vi.fn(),
generateContentStream,
countTokens: vi.fn(),
embedContent: vi.fn(),
batchEmbedContents: vi.fn(),
useSummarizedThinking: vi.fn().mockReturnValue(false),
}) as unknown as ContentGenerator;
const resolveForModel = vi
.fn()
@@ -9582,10 +9570,8 @@ describe('GeminiChat', async () => {
contentGenerator: {
generateContent: vi.fn(),
generateContentStream: fallbackGenerateContentStream,
countTokens: vi.fn(),
embedContent: vi.fn(),
batchEmbedContents: vi.fn(),
useSummarizedThinking: vi.fn().mockReturnValue(false),
} as unknown as ContentGenerator,
retryAuthType: AuthType.USE_GEMINI,
retryErrorCodes: undefined,
@@ -15,7 +15,6 @@ const mockReportGeminiChunk = vi.hoisted(() => vi.fn());
vi.mock('@google/genai', () => {
const mockGenerateContent = vi.fn();
const mockGenerateContentStream = vi.fn();
const mockCountTokens = vi.fn();
const mockEmbedContent = vi.fn();
return {
@@ -23,7 +22,6 @@ vi.mock('@google/genai', () => {
models: {
generateContent: mockGenerateContent,
generateContentStream: mockGenerateContentStream,
countTokens: mockCountTokens,
embedContent: mockEmbedContent,
},
})),
@@ -211,17 +209,6 @@ describe('GeminiContentGenerator', () => {
expect(mockReportGeminiChunk).not.toHaveBeenCalled();
});
it('should call countTokens on the underlying model', async () => {
const request = { model: 'gemini-1.5-flash', contents: [] };
const expectedResponse = { totalTokens: 10 };
mockGoogleGenAI.models.countTokens.mockResolvedValue(expectedResponse);
const response = await generator.countTokens(request);
expect(mockGoogleGenAI.models.countTokens).toHaveBeenCalledWith(request);
expect(response).toBe(expectedResponse);
});
it('should call embedContent on the underlying model', async () => {
const request = { model: 'embedding-model', contents: [] };
const expectedResponse = { embeddings: [] };
@@ -5,8 +5,6 @@
*/
import type {
CountTokensParameters,
CountTokensResponse,
EmbedContentParameters,
EmbedContentResponse,
GenerateContentParameters,
@@ -366,19 +364,9 @@ export class GeminiContentGenerator implements ContentGenerator {
return part;
}
async countTokens(
request: CountTokensParameters,
): Promise<CountTokensResponse> {
return this.googleGenAI.models.countTokens(request);
}
async embedContent(
request: EmbedContentParameters,
): Promise<EmbedContentResponse> {
return this.googleGenAI.models.embedContent(request);
}
useSummarizedThinking(): boolean {
return true;
}
}
@@ -339,9 +339,7 @@ const createWrappedGenerator = (
({
generateContent,
generateContentStream,
countTokens: vi.fn(),
embedContent: vi.fn(),
useSummarizedThinking: vi.fn().mockReturnValue(false),
}) as ContentGenerator;
const createResponse = (
@@ -7,8 +7,6 @@
import {
GenerateContentResponse,
type Content,
type CountTokensParameters,
type CountTokensResponse,
type EmbedContentParameters,
type EmbedContentResponse,
type GenerateContentParameters,
@@ -1245,20 +1243,12 @@ export class LoggingContentGenerator implements ContentGenerator {
return areSensitiveSpanAttributesEnabled(this.config);
}
async countTokens(req: CountTokensParameters): Promise<CountTokensResponse> {
return this.wrapped.countTokens(req);
}
async embedContent(
req: EmbedContentParameters,
): Promise<EmbedContentResponse> {
return this.wrapped.embedContent(req);
}
useSummarizedThinking(): boolean {
return this.wrapped.useSummarizedThinking();
}
private toContents(contents: ContentListUnion): Content[] {
if (Array.isArray(contents)) {
// it's a Content[] or a PartsUnion[]
@@ -5,34 +5,10 @@
*/
import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
// Mock the request tokenizer module BEFORE importing the class that uses it
const mockTokenizer = {
calculateTokens: vi.fn().mockResolvedValue({
totalTokens: 50,
breakdown: {
textTokens: 50,
imageTokens: 0,
audioTokens: 0,
otherTokens: 0,
},
processingTime: 1,
}),
dispose: vi.fn(),
};
vi.mock('../../../utils/request-tokenizer/index.js', () => ({
RequestTokenEstimator: vi.fn(() => mockTokenizer),
}));
// Now import the modules that depend on the mocked modules
import { OpenAIContentGenerator } from './openaiContentGenerator.js';
import type { Config } from '../../config/config.js';
import { AuthType } from '../contentGenerator.js';
import type {
GenerateContentParameters,
CountTokensParameters,
} from '@google/genai';
import type { GenerateContentParameters } from '@google/genai';
import type { OpenAICompatibleProvider } from './provider/index.js';
import type OpenAI from 'openai';
@@ -123,39 +99,6 @@ describe('OpenAIContentGenerator (Refactored)', () => {
});
});
describe('countTokens', () => {
it('should count tokens using character-based estimation', async () => {
const request: CountTokensParameters = {
contents: [{ role: 'user', parts: [{ text: 'Hello world' }] }],
model: 'gpt-4',
};
const result = await generator.countTokens(request);
// 'Hello world' = 11 ASCII chars
// 11 / 4 = 2.75 -> ceil = 3 tokens
expect(result.totalTokens).toBe(3);
});
it('should handle multimodal content', async () => {
const request: CountTokensParameters = {
contents: [
{
role: 'user',
parts: [{ text: 'Hello' }, { text: ' world' }],
},
],
model: 'gpt-4',
};
const result = await generator.countTokens(request);
// Parts are combined for estimation:
// 'Hello world' = 11 ASCII chars -> 11/4 = 2.75 -> ceil = 3 tokens
expect(result.totalTokens).toBe(3);
});
});
describe('embedContent', () => {
it('should delegate to pipeline.client.embeddings.create', async () => {
// This test verifies the method exists and can be called
@@ -2,8 +2,6 @@ import type { ContentGenerator } from '../contentGenerator.js';
import type { Config } from '../../config/config.js';
import { type OpenAICompatibleProvider } from './provider/index.js';
import type {
CountTokensParameters,
CountTokensResponse,
EmbedContentParameters,
EmbedContentResponse,
GenerateContentParameters,
@@ -12,7 +10,6 @@ import type {
import type { PipelineConfig } from './types.js';
import { ContentGenerationPipeline } from './pipeline.js';
import { EnhancedErrorHandler } from './errorHandler.js';
import { RequestTokenEstimator } from '../../utils/request-tokenizer/index.js';
import type { ContentGeneratorConfig } from '../contentGenerator.js';
import { isAbortError } from '../../utils/errors.js';
import { createDebugLogger } from '../../utils/debugLogger.js';
@@ -80,33 +77,6 @@ export class OpenAIContentGenerator implements ContentGenerator {
return this.pipeline.executeStream(request, userPromptId);
}
async countTokens(
request: CountTokensParameters,
): Promise<CountTokensResponse> {
try {
// Use the request token estimator (character-based).
const estimator = new RequestTokenEstimator();
const result = await estimator.calculateTokens(request);
return {
totalTokens: result.totalTokens,
};
} catch (error) {
debugLogger.warn(
'Failed to calculate tokens with new tokenizer, falling back to simple method:',
error,
);
// Fallback to original simple method
const content = JSON.stringify(request.contents);
const totalTokens = Math.ceil(content.length / 4); // Rough estimate: 1 token ≈ 4 characters
return {
totalTokens,
};
}
}
async embedContent(
request: EmbedContentParameters,
): Promise<EmbedContentResponse> {
@@ -163,8 +133,4 @@ export class OpenAIContentGenerator implements ContentGenerator {
);
}
}
useSummarizedThinking(): boolean {
return false;
}
}
@@ -39,9 +39,7 @@ describe('Prompt Hook Integration', () => {
getContentGenerator: vi.fn().mockReturnValue({
generateContent: mockGenerateContent,
generateContentStream: vi.fn(),
countTokens: vi.fn(),
embedContent: vi.fn(),
useSummarizedThinking: vi.fn().mockReturnValue(false),
}),
getBaseLlmClient: vi.fn().mockReturnValue({
resolveForModel: vi.fn().mockResolvedValue({
@@ -40,9 +40,7 @@ describe('PromptHookRunner', () => {
getContentGenerator: vi.fn().mockReturnValue({
generateContent: mockGenerateContent,
generateContentStream: vi.fn(),
countTokens: vi.fn(),
embedContent: vi.fn(),
useSummarizedThinking: vi.fn().mockReturnValue(false),
}),
getBaseLlmClient: vi.fn().mockReturnValue({
resolveForModel: mockResolveForModel,
@@ -10,8 +10,6 @@ import { type QwenCredentials, type ErrorData } from './qwenOAuth2.js';
import type {
GenerateContentParameters,
GenerateContentResponse,
CountTokensParameters,
CountTokensResponse,
EmbedContentParameters,
EmbedContentResponse,
} from '@google/genai';
@@ -96,12 +94,6 @@ vi.mock('../core/openaiContentGenerator/pipeline.js', () => ({
})();
}
async countTokens(
_request: CountTokensParameters,
): Promise<CountTokensResponse> {
return { totalTokens: 15 };
}
async embedContent(
_request: EmbedContentParameters,
): Promise<EmbedContentResponse> {
@@ -251,12 +243,6 @@ vi.mock('../core/openaiContentGenerator/index.js', () => ({
})();
}
async countTokens(
_request: CountTokensParameters,
): Promise<CountTokensResponse> {
return { totalTokens: 15 };
}
async embedContent(
_request: EmbedContentParameters,
): Promise<EmbedContentResponse> {
@@ -401,22 +387,6 @@ describe('QwenContentGenerator', () => {
expect(mockQwenClient.getAccessToken).toHaveBeenCalled();
});
it('should count tokens without requiring authentication', async () => {
// Clear any previous mock calls
vi.clearAllMocks();
const request: CountTokensParameters = {
model: 'qwen-turbo',
contents: [{ role: 'user', parts: [{ text: 'Count me' }] }],
};
const result = await qwenContentGenerator.countTokens(request);
expect(result.totalTokens).toBe(15);
// countTokens is a local operation and should not require OAuth credentials
expect(mockQwenClient.getAccessToken).not.toHaveBeenCalled();
});
it('should embed content with valid token', async () => {
vi.mocked(mockQwenClient.getAccessToken).mockResolvedValue({
token: 'valid-token',
@@ -1650,7 +1620,7 @@ describe('QwenContentGenerator', () => {
SharedTokenManager.getInstance = originalGetInstance;
});
it('should handle method types with token failure (except countTokens)', async () => {
it('should handle method types with token failure', async () => {
const mockTokenManager = {
getValidCredentials: vi
.fn()
@@ -1673,11 +1643,6 @@ describe('QwenContentGenerator', () => {
contents: [{ role: 'user', parts: [{ text: 'Hello' }] }],
};
const countRequest: CountTokensParameters = {
model: 'qwen-turbo',
contents: [{ role: 'user', parts: [{ text: 'Count' }] }],
};
const embedRequest: EmbedContentParameters = {
model: 'qwen-turbo',
contents: [{ parts: [{ text: 'Embed' }] }],
@@ -1696,10 +1661,6 @@ describe('QwenContentGenerator', () => {
'Failed to obtain valid Qwen access token',
);
// countTokens should succeed as it's a local operation
const countResult = await newGenerator.countTokens(countRequest);
expect(countResult.totalTokens).toBe(15);
SharedTokenManager.getInstance = originalGetInstance;
});
});
@@ -12,8 +12,6 @@ import { type Config } from '../config/config.js';
import type {
GenerateContentParameters,
GenerateContentResponse,
CountTokensParameters,
CountTokensResponse,
EmbedContentParameters,
EmbedContentResponse,
} from '@google/genai';
@@ -173,15 +171,6 @@ export class QwenContentGenerator extends OpenAIContentGenerator {
);
}
/**
* Override to use dynamic token and endpoint with automatic retry
*/
override async countTokens(
request: CountTokensParameters,
): Promise<CountTokensResponse> {
return super.countTokens(request);
}
/**
* Override to use dynamic token and endpoint with automatic retry
*/
@@ -59,9 +59,7 @@ describe('checkNextSpeaker', () => {
{
generateContent: vi.fn(),
generateContentStream: vi.fn(),
countTokens: vi.fn(),
embedContent: vi.fn(),
useSummarizedThinking: vi.fn().mockReturnValue(false),
} as ContentGenerator,
{} as Config,
);
@@ -1,337 +0,0 @@
/**
* @license
* Copyright 2025 Qwen
* SPDX-License-Identifier: Apache-2.0
*/
import { describe, it, expect } from 'vitest';
import { ImageTokenizer } from './imageTokenizer.js';
describe('ImageTokenizer', () => {
const tokenizer = new ImageTokenizer();
describe('token calculation', () => {
it('should calculate tokens based on image dimensions with reference logic', () => {
const metadata = {
width: 28,
height: 28,
mimeType: 'image/png',
dataSize: 1000,
};
const tokens = tokenizer.calculateTokens(metadata);
// 28x28 = 784 pixels = 1 image token + 2 special tokens = 3 total
// But minimum scaling may apply for small images
expect(tokens).toBeGreaterThanOrEqual(6); // Minimum after scaling + special tokens
});
it('should calculate tokens for larger images', () => {
const metadata = {
width: 512,
height: 512,
mimeType: 'image/png',
dataSize: 10000,
};
const tokens = tokenizer.calculateTokens(metadata);
// 512x512 with reference logic: rounded dimensions + scaling + special tokens
expect(tokens).toBeGreaterThan(300);
expect(tokens).toBeLessThan(400); // Should be reasonable for 512x512
});
it('should enforce minimum tokens per image with scaling', () => {
const metadata = {
width: 1,
height: 1,
mimeType: 'image/png',
dataSize: 100,
};
const tokens = tokenizer.calculateTokens(metadata);
// Tiny images get scaled up to minimum pixels + special tokens
expect(tokens).toBeGreaterThanOrEqual(6); // 4 image tokens + 2 special tokens
});
it('should handle very large images with scaling', () => {
const metadata = {
width: 8192,
height: 8192,
mimeType: 'image/png',
dataSize: 100000,
};
const tokens = tokenizer.calculateTokens(metadata);
// Very large images should be scaled down to max limit + special tokens
expect(tokens).toBeLessThanOrEqual(16386); // 16384 max + 2 special tokens
expect(tokens).toBeGreaterThan(16000); // Should be close to the limit
});
});
describe('PNG dimension extraction', () => {
it('should extract dimensions from valid PNG', async () => {
// 1x1 PNG image in base64
const pngBase64 =
'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChAI9jU77yQAAAABJRU5ErkJggg==';
const metadata = await tokenizer.extractImageMetadata(
pngBase64,
'image/png',
);
expect(metadata.width).toBe(1);
expect(metadata.height).toBe(1);
expect(metadata.mimeType).toBe('image/png');
});
it('should handle invalid PNG gracefully', async () => {
const invalidBase64 = 'invalid-png-data';
const metadata = await tokenizer.extractImageMetadata(
invalidBase64,
'image/png',
);
// Should return default dimensions
expect(metadata.width).toBe(512);
expect(metadata.height).toBe(512);
expect(metadata.mimeType).toBe('image/png');
});
});
describe('WebP dimension extraction', () => {
function buildWebp(format: string, totalLength: number): Buffer {
const buf = Buffer.alloc(totalLength);
buf.write('RIFF', 0, 'ascii');
buf.writeUInt32LE(totalLength - 8, 4);
buf.write('WEBP', 8, 'ascii');
buf.write(format, 12, 'ascii');
buf.writeUInt32LE(Math.max(totalLength - 20, 0), 16);
return buf;
}
function buildVp8lWebp(
width: number,
height: number,
signature = 0x2f,
): Buffer {
const buf = buildWebp('VP8L', 26);
buf.writeUInt32LE(5, 16);
buf.writeUInt8(signature, 20);
buf.writeUInt32LE((width - 1) | ((height - 1) << 14), 21);
return buf;
}
it('should extract dimensions from a short VP8L lossless WebP', async () => {
const width = 17;
const height = 13;
const buf = buildVp8lWebp(width, height);
const metadata = await tokenizer.extractImageMetadata(
buf.toString('base64'),
'image/webp',
);
expect(metadata.width).toBe(width);
expect(metadata.height).toBe(height);
});
it('should fall back for VP8L with an invalid lossless signature', async () => {
const buf = buildVp8lWebp(17, 13, 0x00);
const metadata = await tokenizer.extractImageMetadata(
buf.toString('base64'),
'image/webp',
);
expect(metadata.width).toBe(512);
expect(metadata.height).toBe(512);
});
it('should still reject short VP8 and VP8X WebP files', async () => {
const formats = ['VP8 ', 'VP8X'];
for (const format of formats) {
const metadata = await tokenizer.extractImageMetadata(
buildWebp(format, 26).toString('base64'),
'image/webp',
);
expect(metadata.width).toBe(512);
expect(metadata.height).toBe(512);
}
});
it('should extract canvas dimensions from VP8X', async () => {
const width = 100;
const height = 80;
const buf = Buffer.alloc(30);
buf.write('RIFF', 0, 'ascii');
buf.writeUInt32LE(22, 4);
buf.write('WEBP', 8, 'ascii');
buf.write('VP8X', 12, 'ascii');
buf.writeUInt32LE(10, 16); // VP8X chunk size
buf.writeUInt8(0, 20); // flags
buf.writeUIntLE(width - 1, 24, 3); // canvas width minus one (24-bit LE)
buf.writeUIntLE(height - 1, 27, 3); // canvas height minus one (24-bit LE)
const metadata = await tokenizer.extractImageMetadata(
buf.toString('base64'),
'image/webp',
);
expect(metadata.width).toBe(width);
expect(metadata.height).toBe(height);
});
});
describe('batch processing', () => {
it('should process multiple images serially', async () => {
const pngBase64 =
'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChAI9jU77yQAAAABJRU5ErkJggg==';
const images = [
{ data: pngBase64, mimeType: 'image/png' },
{ data: pngBase64, mimeType: 'image/png' },
{ data: pngBase64, mimeType: 'image/png' },
];
const tokens = await tokenizer.calculateTokensBatch(images);
expect(tokens).toHaveLength(3);
expect(tokens.every((t) => t >= 4)).toBe(true); // All should have at least 4 tokens
});
it('should handle mixed valid and invalid images', async () => {
const validPng =
'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChAI9jU77yQAAAABJRU5ErkJggg==';
const invalidPng = 'invalid-data';
const images = [
{ data: validPng, mimeType: 'image/png' },
{ data: invalidPng, mimeType: 'image/png' },
];
const tokens = await tokenizer.calculateTokensBatch(images);
expect(tokens).toHaveLength(2);
expect(tokens.every((t) => t >= 4)).toBe(true); // All should have at least minimum tokens
});
});
describe('different image formats', () => {
it('should handle different MIME types', async () => {
const pngBase64 =
'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChAI9jU77yQAAAABJRU5ErkJggg==';
const formats = ['image/png', 'image/jpeg', 'image/webp', 'image/gif'];
for (const mimeType of formats) {
const metadata = await tokenizer.extractImageMetadata(
pngBase64,
mimeType,
);
expect(metadata.mimeType).toBe(mimeType);
expect(metadata.width).toBeGreaterThan(0);
expect(metadata.height).toBeGreaterThan(0);
}
});
it('should extract dimensions from GIF images', async () => {
const buf = Buffer.alloc(10);
buf.write('GIF89a', 0, 'ascii');
buf.writeUInt16LE(2, 6);
buf.writeUInt16LE(3, 8);
const metadata = await tokenizer.extractImageMetadata(
buf.toString('base64'),
'image/gif',
);
expect(metadata.width).toBe(2);
expect(metadata.height).toBe(3);
expect(metadata.mimeType).toBe('image/gif');
});
});
describe('TIFF dimension extraction', () => {
// Build a minimal single-IFD TIFF whose ImageWidth/ImageLength are stored
// as SHORT (type 3) -- the most common TIFF layout for small dimensions.
function buildShortTiff(
byteOrder: 'II' | 'MM',
width: number,
height: number,
): string {
const buf = Buffer.alloc(38);
const le = byteOrder === 'II';
buf.write(byteOrder, 0, 'ascii');
const w16 = (off: number, v: number) =>
le ? buf.writeUInt16LE(v, off) : buf.writeUInt16BE(v, off);
const w32 = (off: number, v: number) =>
le ? buf.writeUInt32LE(v, off) : buf.writeUInt32BE(v, off);
w16(2, 42); // magic
w32(4, 8); // IFD starts at offset 8
w16(8, 2); // two directory entries
// entry 0: ImageWidth (0x0100), SHORT, count 1, value left-justified
w16(10, 0x0100);
w16(12, 3);
w32(14, 1);
w16(18, width);
// entry 1: ImageLength (0x0101), SHORT, count 1
w16(22, 0x0101);
w16(24, 3);
w32(26, 1);
w16(30, height);
w32(34, 0); // next-IFD offset
return buf.toString('base64');
}
it('reads SHORT dimensions from a big-endian (MM) TIFF', async () => {
const tiff = buildShortTiff('MM', 800, 600);
const metadata = await tokenizer.extractImageMetadata(tiff, 'image/tiff');
expect(metadata.width).toBe(800);
expect(metadata.height).toBe(600);
});
it('reads SHORT dimensions from a little-endian (II) TIFF', async () => {
const tiff = buildShortTiff('II', 1024, 768);
const metadata = await tokenizer.extractImageMetadata(tiff, 'image/tiff');
expect(metadata.width).toBe(1024);
expect(metadata.height).toBe(768);
});
});
describe('BMP dimension extraction', () => {
function buildBmp(width: number, height: number): string {
// A negative height encodes a top-down BMP per the BITMAPINFOHEADER spec.
const buf = Buffer.alloc(26);
buf.write('BM', 0, 'ascii');
buf.writeUInt32LE(26, 2); // file size
buf.writeUInt32LE(0, 6); // reserved
buf.writeUInt32LE(26, 10); // pixel data offset
buf.writeUInt32LE(40, 14); // BITMAPINFOHEADER size
buf.writeInt32LE(width, 18);
buf.writeInt32LE(height, 22);
return buf.toString('base64');
}
it('reads a bottom-up (positive height) BMP', async () => {
const bmp = buildBmp(120, 80);
const metadata = await tokenizer.extractImageMetadata(bmp, 'image/bmp');
expect(metadata.width).toBe(120);
expect(metadata.height).toBe(80);
});
it('reads a top-down (negative height) BMP as its absolute height', async () => {
const bmp = buildBmp(100, -50);
const metadata = await tokenizer.extractImageMetadata(bmp, 'image/bmp');
expect(metadata.width).toBe(100);
expect(metadata.height).toBe(50);
});
});
});
@@ -1,534 +0,0 @@
/**
* @license
* Copyright 2025 Qwen
* SPDX-License-Identifier: Apache-2.0
*/
import type { ImageMetadata } from './types.js';
import { isSupportedImageMimeType } from './supportedImageFormats.js';
import { createDebugLogger } from '../debugLogger.js';
const debugLogger = createDebugLogger('IMAGE_TOKENIZER');
/**
* Image tokenizer for calculating image tokens based on dimensions
*
* Key rules:
* - 28x28 pixels = 1 token
* - Minimum: 4 tokens per image
* - Maximum: 16384 tokens per image
* - Additional: 2 special tokens (vision_bos + vision_eos)
* - Supports: PNG, JPEG, WebP, GIF, BMP, TIFF, HEIC formats
*/
export class ImageTokenizer {
/** 28x28 pixels = 1 token */
private static readonly PIXELS_PER_TOKEN = 28 * 28;
/** Minimum tokens per image */
private static readonly MIN_TOKENS_PER_IMAGE = 4;
/** Maximum tokens per image */
private static readonly MAX_TOKENS_PER_IMAGE = 16384;
/** Special tokens for vision markers */
private static readonly VISION_SPECIAL_TOKENS = 2;
/**
* Extract image metadata from base64 data
*
* @param base64Data Base64-encoded image data (with or without data URL prefix)
* @param mimeType MIME type of the image
* @returns Promise resolving to ImageMetadata with dimensions and format info
*/
async extractImageMetadata(
base64Data: string,
mimeType: string,
): Promise<ImageMetadata> {
try {
// Check if the MIME type is supported
if (!isSupportedImageMimeType(mimeType)) {
debugLogger.warn(`Unsupported image format: ${mimeType}`);
// Return default metadata for unsupported formats
return {
width: 512,
height: 512,
mimeType,
dataSize: Math.floor(base64Data.length * 0.75),
};
}
const cleanBase64 = base64Data.replace(/^data:[^;]+;base64,/, '');
const buffer = Buffer.from(cleanBase64, 'base64');
const dimensions = await this.extractDimensions(buffer, mimeType);
return {
width: dimensions.width,
height: dimensions.height,
mimeType,
dataSize: buffer.length,
};
} catch (error) {
debugLogger.warn('Failed to extract image metadata:', error);
// Return default metadata for fallback
return {
width: 512,
height: 512,
mimeType,
dataSize: Math.floor(base64Data.length * 0.75),
};
}
}
/**
* Extract image dimensions from buffer based on format
*
* @param buffer Binary image data buffer
* @param mimeType MIME type to determine parsing strategy
* @returns Promise resolving to width and height dimensions
*/
private async extractDimensions(
buffer: Buffer,
mimeType: string,
): Promise<{ width: number; height: number }> {
if (mimeType.includes('png')) {
return this.extractPngDimensions(buffer);
}
if (mimeType.includes('jpeg') || mimeType.includes('jpg')) {
return this.extractJpegDimensions(buffer);
}
if (mimeType.includes('webp')) {
return this.extractWebpDimensions(buffer);
}
if (mimeType.includes('gif')) {
return this.extractGifDimensions(buffer);
}
if (mimeType.includes('bmp')) {
return this.extractBmpDimensions(buffer);
}
if (mimeType.includes('tiff')) {
return this.extractTiffDimensions(buffer);
}
if (mimeType.includes('heic')) {
return this.extractHeicDimensions(buffer);
}
return { width: 512, height: 512 };
}
/**
* Extract PNG dimensions from IHDR chunk
* PNG signature: 89 50 4E 47 0D 0A 1A 0A
* Width/height at bytes 16-19 and 20-23 (big-endian)
*/
private extractPngDimensions(buffer: Buffer): {
width: number;
height: number;
} {
if (buffer.length < 24) {
throw new Error('Invalid PNG: buffer too short');
}
// Verify PNG signature
const signature = buffer.subarray(0, 8);
const expectedSignature = Buffer.from([
0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a,
]);
if (!signature.equals(expectedSignature)) {
throw new Error('Invalid PNG signature');
}
const width = buffer.readUInt32BE(16);
const height = buffer.readUInt32BE(20);
return { width, height };
}
/**
* Extract JPEG dimensions from SOF (Start of Frame) markers
* JPEG starts with FF D8, SOF markers: 0xC0-0xC3, 0xC5-0xC7, 0xC9-0xCB, 0xCD-0xCF
* Dimensions at offset +5 (height) and +7 (width) from SOF marker
*/
private extractJpegDimensions(buffer: Buffer): {
width: number;
height: number;
} {
if (buffer.length < 4 || buffer[0] !== 0xff || buffer[1] !== 0xd8) {
throw new Error('Invalid JPEG signature');
}
let offset = 2;
while (offset < buffer.length - 8) {
if (buffer[offset] !== 0xff) {
offset++;
continue;
}
const marker = buffer[offset + 1];
// SOF markers
if (
(marker >= 0xc0 && marker <= 0xc3) ||
(marker >= 0xc5 && marker <= 0xc7) ||
(marker >= 0xc9 && marker <= 0xcb) ||
(marker >= 0xcd && marker <= 0xcf)
) {
const height = buffer.readUInt16BE(offset + 5);
const width = buffer.readUInt16BE(offset + 7);
return { width, height };
}
const segmentLength = buffer.readUInt16BE(offset + 2);
offset += 2 + segmentLength;
}
throw new Error('Could not find JPEG dimensions');
}
/**
* Extract WebP dimensions from RIFF container
* Supports VP8, VP8L, and VP8X formats
*/
private extractWebpDimensions(buffer: Buffer): {
width: number;
height: number;
} {
if (buffer.length < 16) {
throw new Error('Invalid WebP: too short');
}
const riffSignature = buffer.subarray(0, 4).toString('ascii');
const webpSignature = buffer.subarray(8, 12).toString('ascii');
if (riffSignature !== 'RIFF' || webpSignature !== 'WEBP') {
throw new Error('Invalid WebP signature');
}
const format = buffer.subarray(12, 16).toString('ascii');
if (format === 'VP8 ') {
if (buffer.length < 30) {
throw new Error('Invalid VP8 WebP: too short');
}
// Lossy: 14-bit width/height at bytes 26-27 and 28-29 (little-endian)
const width = buffer.readUInt16LE(26) & 0x3fff;
const height = buffer.readUInt16LE(28) & 0x3fff;
return { width, height };
} else if (format === 'VP8L') {
if (buffer.length < 25) {
throw new Error('Invalid VP8L WebP: too short');
}
if (buffer[20] !== 0x2f) {
throw new Error('Invalid VP8L WebP signature');
}
// Lossless: 14-bit (width-1) then (height-1) packed from byte 21 (little-endian)
const bits = buffer.readUInt32LE(21);
const width = (bits & 0x3fff) + 1;
const height = ((bits >> 14) & 0x3fff) + 1;
return { width, height };
} else if (format === 'VP8X') {
if (buffer.length < 30) {
throw new Error('Invalid VP8X WebP: too short');
}
// Extended: 24-bit (canvas width-1) at bytes 24-26 and (height-1) at bytes 27-29 (little-endian)
const width = buffer.readUIntLE(24, 3) + 1;
const height = buffer.readUIntLE(27, 3) + 1;
return { width, height };
}
throw new Error('Unsupported WebP format');
}
/**
* Extract GIF dimensions from header
* Supports GIF87a and GIF89a formats
*/
private extractGifDimensions(buffer: Buffer): {
width: number;
height: number;
} {
if (buffer.length < 10) {
throw new Error('Invalid GIF: too short');
}
const signature = buffer.subarray(0, 6).toString('ascii');
if (signature !== 'GIF87a' && signature !== 'GIF89a') {
throw new Error('Invalid GIF signature');
}
const width = buffer.readUInt16LE(6);
const height = buffer.readUInt16LE(8);
return { width, height };
}
/**
* Calculate tokens for an image based on its metadata
*
* @param metadata Image metadata containing width, height, and format info
* @returns Total token count including base image tokens and special tokens
*/
calculateTokens(metadata: ImageMetadata): number {
return this.calculateTokensWithScaling(metadata.width, metadata.height);
}
/**
* Calculate tokens with scaling logic
*
* Steps:
* 1. Normalize to 28-pixel multiples
* 2. Scale large images down, small images up
* 3. Calculate tokens: pixels / 784 + 2 special tokens
*
* @param width Original image width in pixels
* @param height Original image height in pixels
* @returns Total token count for the image
*/
private calculateTokensWithScaling(width: number, height: number): number {
// Normalize to 28-pixel multiples
let hBar = Math.round(height / 28) * 28;
let wBar = Math.round(width / 28) * 28;
// Define pixel boundaries
const minPixels =
ImageTokenizer.MIN_TOKENS_PER_IMAGE * ImageTokenizer.PIXELS_PER_TOKEN;
const maxPixels =
ImageTokenizer.MAX_TOKENS_PER_IMAGE * ImageTokenizer.PIXELS_PER_TOKEN;
// Apply scaling
if (hBar * wBar > maxPixels) {
// Scale down large images
const beta = Math.sqrt((height * width) / maxPixels);
hBar = Math.floor(height / beta / 28) * 28;
wBar = Math.floor(width / beta / 28) * 28;
} else if (hBar * wBar < minPixels) {
// Scale up small images
const beta = Math.sqrt(minPixels / (height * width));
hBar = Math.ceil((height * beta) / 28) * 28;
wBar = Math.ceil((width * beta) / 28) * 28;
}
// Calculate tokens
const imageTokens = Math.floor(
(hBar * wBar) / ImageTokenizer.PIXELS_PER_TOKEN,
);
return imageTokens + ImageTokenizer.VISION_SPECIAL_TOKENS;
}
/**
* Calculate tokens for multiple images serially
*
* @param base64DataArray Array of image data with MIME type information
* @returns Promise resolving to array of token counts in same order as input
*/
async calculateTokensBatch(
base64DataArray: Array<{ data: string; mimeType: string }>,
): Promise<number[]> {
const results: number[] = [];
for (const { data, mimeType } of base64DataArray) {
try {
const metadata = await this.extractImageMetadata(data, mimeType);
results.push(this.calculateTokens(metadata));
} catch (error) {
debugLogger.warn('Error calculating tokens for image:', error);
// Return minimum tokens as fallback
results.push(
ImageTokenizer.MIN_TOKENS_PER_IMAGE +
ImageTokenizer.VISION_SPECIAL_TOKENS,
);
}
}
return results;
}
/**
* Extract BMP dimensions from header
* BMP signature: 42 4D (BM)
* Width/height at bytes 18-21 and 22-25 (little-endian)
*/
private extractBmpDimensions(buffer: Buffer): {
width: number;
height: number;
} {
if (buffer.length < 26) {
throw new Error('Invalid BMP: buffer too short');
}
// Verify BMP signature
if (buffer[0] !== 0x42 || buffer[1] !== 0x4d) {
throw new Error('Invalid BMP signature');
}
const width = buffer.readUInt32LE(18);
// Height is a signed int32: a negative value means a top-down BMP. Read it
// signed so Math.abs recovers the real height instead of a ~4-billion value.
const height = buffer.readInt32LE(22);
return { width, height: Math.abs(height) }; // Height can be negative for top-down BMPs
}
/**
* Extract TIFF dimensions from IFD (Image File Directory)
* TIFF can be little-endian (II) or big-endian (MM)
* Width/height are stored in IFD entries with tags 0x0100 and 0x0101
*/
private extractTiffDimensions(buffer: Buffer): {
width: number;
height: number;
} {
if (buffer.length < 8) {
throw new Error('Invalid TIFF: buffer too short');
}
// Check byte order
const byteOrder = buffer.subarray(0, 2).toString('ascii');
const isLittleEndian = byteOrder === 'II';
const isBigEndian = byteOrder === 'MM';
if (!isLittleEndian && !isBigEndian) {
throw new Error('Invalid TIFF byte order');
}
// Read magic number (should be 42)
const magic = isLittleEndian
? buffer.readUInt16LE(2)
: buffer.readUInt16BE(2);
if (magic !== 42) {
throw new Error('Invalid TIFF magic number');
}
// Read IFD offset
const ifdOffset = isLittleEndian
? buffer.readUInt32LE(4)
: buffer.readUInt32BE(4);
if (ifdOffset >= buffer.length) {
throw new Error('Invalid TIFF IFD offset');
}
// Read number of directory entries
const numEntries = isLittleEndian
? buffer.readUInt16LE(ifdOffset)
: buffer.readUInt16BE(ifdOffset);
let width = 0;
let height = 0;
// Parse IFD entries
for (let i = 0; i < numEntries; i++) {
const entryOffset = ifdOffset + 2 + i * 12;
if (entryOffset + 12 > buffer.length) break;
const tag = isLittleEndian
? buffer.readUInt16LE(entryOffset)
: buffer.readUInt16BE(entryOffset);
const type = isLittleEndian
? buffer.readUInt16LE(entryOffset + 2)
: buffer.readUInt16BE(entryOffset + 2);
// A SHORT (type 3) value occupies only the first two bytes of the
// 4-byte value field. Reading it as a 32-bit int happens to work on
// little-endian (II) files but yields value << 16 on big-endian (MM)
// ones, so read it according to the field type.
const value =
type === 3
? isLittleEndian
? buffer.readUInt16LE(entryOffset + 8)
: buffer.readUInt16BE(entryOffset + 8)
: isLittleEndian
? buffer.readUInt32LE(entryOffset + 8)
: buffer.readUInt32BE(entryOffset + 8);
if (tag === 0x0100) {
// ImageWidth
width = value;
} else if (tag === 0x0101) {
// ImageLength (height)
height = value;
}
if (width > 0 && height > 0) break;
}
if (width === 0 || height === 0) {
throw new Error('Could not find TIFF dimensions');
}
return { width, height };
}
/**
* Extract HEIC dimensions from meta box
* HEIC is based on ISO Base Media File Format
* This is a simplified implementation that looks for 'ispe' (Image Spatial Extents) box
*/
private extractHeicDimensions(buffer: Buffer): {
width: number;
height: number;
} {
if (buffer.length < 12) {
throw new Error('Invalid HEIC: buffer too short');
}
// Check for ftyp box with HEIC brand
const ftypBox = buffer.subarray(4, 8).toString('ascii');
if (ftypBox !== 'ftyp') {
throw new Error('Invalid HEIC: missing ftyp box');
}
const brand = buffer.subarray(8, 12).toString('ascii');
if (!['heic', 'heix', 'hevc', 'hevx'].includes(brand)) {
throw new Error('Invalid HEIC brand');
}
// Look for meta box and then ispe box
let offset = 0;
while (offset < buffer.length - 8) {
const boxSize = buffer.readUInt32BE(offset);
const boxType = buffer.subarray(offset + 4, offset + 8).toString('ascii');
if (boxType === 'meta') {
// Look for ispe box inside meta box
const metaOffset = offset + 8;
let innerOffset = metaOffset + 4; // Skip version and flags
while (innerOffset < offset + boxSize - 8) {
const innerBoxSize = buffer.readUInt32BE(innerOffset);
const innerBoxType = buffer
.subarray(innerOffset + 4, innerOffset + 8)
.toString('ascii');
if (innerBoxType === 'ispe') {
// Found Image Spatial Extents box
if (innerOffset + 20 <= buffer.length) {
const width = buffer.readUInt32BE(innerOffset + 12);
const height = buffer.readUInt32BE(innerOffset + 16);
return { width, height };
}
}
if (innerBoxSize === 0) break;
innerOffset += innerBoxSize;
}
}
if (boxSize === 0) break;
offset += boxSize;
}
// Fallback: return default dimensions if we can't parse the structure
debugLogger.warn('Could not extract HEIC dimensions, using default');
return { width: 512, height: 512 };
}
}
@@ -1,11 +0,0 @@
/**
* @license
* Copyright 2025 Qwen
* SPDX-License-Identifier: Apache-2.0
*/
export { RequestTokenizer as RequestTokenEstimator } from './requestTokenizer.js';
export { TextTokenizer } from './textTokenizer.js';
export { ImageTokenizer } from './imageTokenizer.js';
export type { TokenCalculationResult, ImageMetadata } from './types.js';
@@ -1,271 +0,0 @@
/**
* @license
* Copyright 2025 Qwen
* SPDX-License-Identifier: Apache-2.0
*/
import { describe, it, expect, beforeEach } from 'vitest';
import { RequestTokenizer } from './requestTokenizer.js';
import type { CountTokensParameters } from '@google/genai';
describe('RequestTokenEstimator', () => {
let tokenizer: RequestTokenizer;
beforeEach(() => {
tokenizer = new RequestTokenizer();
});
describe('text token calculation', () => {
it('should calculate tokens for simple text content', async () => {
const request: CountTokensParameters = {
model: 'test-model',
contents: [
{
role: 'user',
parts: [{ text: 'Hello, world!' }],
},
],
};
const result = await tokenizer.calculateTokens(request);
expect(result.totalTokens).toBeGreaterThan(0);
expect(result.breakdown.textTokens).toBeGreaterThan(0);
expect(result.breakdown.imageTokens).toBe(0);
expect(result.processingTime).toBeGreaterThan(0);
});
it('should handle multiple text parts', async () => {
const request: CountTokensParameters = {
model: 'test-model',
contents: [
{
role: 'user',
parts: [
{ text: 'First part' },
{ text: 'Second part' },
{ text: 'Third part' },
],
},
],
};
const result = await tokenizer.calculateTokens(request);
expect(result.totalTokens).toBeGreaterThan(0);
expect(result.breakdown.textTokens).toBeGreaterThan(0);
});
it('should handle string content', async () => {
const request: CountTokensParameters = {
model: 'test-model',
contents: ['Simple string content'],
};
const result = await tokenizer.calculateTokens(request);
expect(result.totalTokens).toBeGreaterThan(0);
expect(result.breakdown.textTokens).toBeGreaterThan(0);
});
});
describe('image token calculation', () => {
it('should calculate tokens for image content', async () => {
// Create a simple 1x1 PNG image in base64
const pngBase64 =
'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChAI9jU77yQAAAABJRU5ErkJggg==';
const request: CountTokensParameters = {
model: 'test-model',
contents: [
{
role: 'user',
parts: [
{
inlineData: {
mimeType: 'image/png',
data: pngBase64,
},
},
],
},
],
};
const result = await tokenizer.calculateTokens(request);
expect(result.totalTokens).toBeGreaterThanOrEqual(4); // Minimum 4 tokens per image
expect(result.breakdown.imageTokens).toBeGreaterThanOrEqual(4);
expect(result.breakdown.textTokens).toBe(0);
});
it('should handle multiple images', async () => {
const pngBase64 =
'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChAI9jU77yQAAAABJRU5ErkJggg==';
const request: CountTokensParameters = {
model: 'test-model',
contents: [
{
role: 'user',
parts: [
{
inlineData: {
mimeType: 'image/png',
data: pngBase64,
},
},
{
inlineData: {
mimeType: 'image/png',
data: pngBase64,
},
},
],
},
],
};
const result = await tokenizer.calculateTokens(request);
expect(result.totalTokens).toBeGreaterThanOrEqual(8); // At least 4 tokens per image
expect(result.breakdown.imageTokens).toBeGreaterThanOrEqual(8);
});
});
describe('mixed content', () => {
it('should handle text and image content together', async () => {
const pngBase64 =
'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChAI9jU77yQAAAABJRU5ErkJggg==';
const request: CountTokensParameters = {
model: 'test-model',
contents: [
{
role: 'user',
parts: [
{ text: 'Here is an image:' },
{
inlineData: {
mimeType: 'image/png',
data: pngBase64,
},
},
{ text: 'What do you see?' },
],
},
],
};
const result = await tokenizer.calculateTokens(request);
expect(result.totalTokens).toBeGreaterThan(4);
expect(result.breakdown.textTokens).toBeGreaterThan(0);
expect(result.breakdown.imageTokens).toBeGreaterThanOrEqual(4);
});
});
describe('function content', () => {
it('should handle function calls', async () => {
const request: CountTokensParameters = {
model: 'test-model',
contents: [
{
role: 'user',
parts: [
{
functionCall: {
name: 'test_function',
args: { param1: 'value1', param2: 42 },
},
},
],
},
],
};
const result = await tokenizer.calculateTokens(request);
expect(result.totalTokens).toBeGreaterThan(0);
expect(result.breakdown.otherTokens).toBeGreaterThan(0);
});
});
describe('empty content', () => {
it('should handle empty request', async () => {
const request: CountTokensParameters = {
model: 'test-model',
contents: [],
};
const result = await tokenizer.calculateTokens(request);
expect(result.totalTokens).toBe(0);
expect(result.breakdown.textTokens).toBe(0);
expect(result.breakdown.imageTokens).toBe(0);
});
it('should handle undefined contents', async () => {
const request: CountTokensParameters = {
model: 'test-model',
contents: [],
};
const result = await tokenizer.calculateTokens(request);
expect(result.totalTokens).toBe(0);
});
});
describe('images', () => {
it('should process multiple images serially', async () => {
const pngBase64 =
'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChAI9jU77yQAAAABJRU5ErkJggg==';
const request: CountTokensParameters = {
model: 'test-model',
contents: [
{
role: 'user',
parts: Array(10).fill({
inlineData: {
mimeType: 'image/png',
data: pngBase64,
},
}),
},
],
};
const result = await tokenizer.calculateTokens(request);
expect(result.totalTokens).toBeGreaterThanOrEqual(60); // At least 6 tokens per image * 10 images
});
});
describe('error handling', () => {
it('should handle malformed image data gracefully', async () => {
const request: CountTokensParameters = {
model: 'test-model',
contents: [
{
role: 'user',
parts: [
{
inlineData: {
mimeType: 'image/png',
data: 'invalid-base64-data',
},
},
],
},
],
};
const result = await tokenizer.calculateTokens(request);
// Should still return some tokens (fallback to minimum)
expect(result.totalTokens).toBeGreaterThanOrEqual(4);
});
});
});
@@ -1,330 +0,0 @@
/**
* @license
* Copyright 2025 Qwen
* SPDX-License-Identifier: Apache-2.0
*/
import type {
CountTokensParameters,
Content,
Part,
PartUnion,
} from '@google/genai';
import type { TokenCalculationResult } from './types.js';
import { TextTokenizer } from './textTokenizer.js';
import { ImageTokenizer } from './imageTokenizer.js';
import { createDebugLogger } from '../debugLogger.js';
const debugLogger = createDebugLogger('TOKENIZER');
/**
* Simple request token estimator that handles text and image content serially
*/
export class RequestTokenizer {
private textTokenizer: TextTokenizer;
private imageTokenizer: ImageTokenizer;
constructor() {
this.textTokenizer = new TextTokenizer();
this.imageTokenizer = new ImageTokenizer();
}
/**
* Calculate tokens for a request using serial processing
*/
async calculateTokens(
request: CountTokensParameters,
): Promise<TokenCalculationResult> {
const startTime = performance.now();
try {
// Process request content and group by type
const { textContents, imageContents, audioContents, otherContents } =
this.processAndGroupContents(request);
if (
textContents.length === 0 &&
imageContents.length === 0 &&
audioContents.length === 0 &&
otherContents.length === 0
) {
return {
totalTokens: 0,
breakdown: {
textTokens: 0,
imageTokens: 0,
audioTokens: 0,
otherTokens: 0,
},
processingTime: performance.now() - startTime,
};
}
// Calculate tokens for each content type serially
const textTokens = await this.calculateTextTokens(textContents);
const imageTokens = await this.calculateImageTokens(imageContents);
const audioTokens = await this.calculateAudioTokens(audioContents);
const otherTokens = await this.calculateOtherTokens(otherContents);
const totalTokens = textTokens + imageTokens + audioTokens + otherTokens;
const processingTime = performance.now() - startTime;
return {
totalTokens,
breakdown: {
textTokens,
imageTokens,
audioTokens,
otherTokens,
},
processingTime,
};
} catch (error) {
debugLogger.error('Error calculating tokens:', error);
// Fallback calculation
const fallbackTokens = this.calculateFallbackTokens(request);
return {
totalTokens: fallbackTokens,
breakdown: {
textTokens: fallbackTokens,
imageTokens: 0,
audioTokens: 0,
otherTokens: 0,
},
processingTime: performance.now() - startTime,
};
}
}
/**
* Calculate tokens for text contents
*/
private async calculateTextTokens(textContents: string[]): Promise<number> {
if (textContents.length === 0) return 0;
try {
// Avoid per-part rounding inflation by estimating once on the combined text.
return await this.textTokenizer.calculateTokens(textContents.join(''));
} catch (error) {
debugLogger.warn('Error calculating text tokens:', error);
// Fallback: character-based estimation
const totalChars = textContents.join('').length;
return Math.ceil(totalChars / 4);
}
}
/**
* Calculate tokens for image contents using serial processing
*/
private async calculateImageTokens(
imageContents: Array<{ data: string; mimeType: string }>,
): Promise<number> {
if (imageContents.length === 0) return 0;
try {
const tokenCounts =
await this.imageTokenizer.calculateTokensBatch(imageContents);
return tokenCounts.reduce((sum, count) => sum + count, 0);
} catch (error) {
debugLogger.warn('Error calculating image tokens:', error);
// Fallback: minimum tokens per image
return imageContents.length * 6; // 4 image tokens + 2 special tokens as minimum
}
}
/**
* Calculate tokens for audio contents
* TODO: Implement proper audio token calculation
*/
private async calculateAudioTokens(
audioContents: Array<{ data: string; mimeType: string }>,
): Promise<number> {
if (audioContents.length === 0) return 0;
// Placeholder implementation - audio token calculation would depend on
// the specific model's audio processing capabilities
// For now, estimate based on data size
let totalTokens = 0;
for (const audioContent of audioContents) {
try {
const dataSize = Math.floor(audioContent.data.length * 0.75); // Approximate binary size
// Rough estimate: 1 token per 100 bytes of audio data
totalTokens += Math.max(Math.ceil(dataSize / 100), 10); // Minimum 10 tokens per audio
} catch (error) {
debugLogger.warn('Error calculating audio tokens:', error);
totalTokens += 10; // Fallback minimum
}
}
return totalTokens;
}
/**
* Calculate tokens for other content types (functions, files, etc.)
*/
private async calculateOtherTokens(otherContents: string[]): Promise<number> {
if (otherContents.length === 0) return 0;
try {
// Treat other content as text, and avoid per-item rounding inflation.
return await this.textTokenizer.calculateTokens(otherContents.join(''));
} catch (error) {
debugLogger.warn('Error calculating other content tokens:', error);
// Fallback: character-based estimation
const totalChars = otherContents.join('').length;
return Math.ceil(totalChars / 4);
}
}
/**
* Fallback token calculation using simple string serialization
*/
private calculateFallbackTokens(request: CountTokensParameters): number {
try {
const content = JSON.stringify(request.contents);
return Math.ceil(content.length / 4); // Rough estimate: 1 token ≈ 4 characters
} catch (error) {
debugLogger.warn('Error in fallback token calculation:', error);
return 100; // Conservative fallback
}
}
/**
* Process request contents and group by type
*/
private processAndGroupContents(request: CountTokensParameters): {
textContents: string[];
imageContents: Array<{ data: string; mimeType: string }>;
audioContents: Array<{ data: string; mimeType: string }>;
otherContents: string[];
} {
const textContents: string[] = [];
const imageContents: Array<{ data: string; mimeType: string }> = [];
const audioContents: Array<{ data: string; mimeType: string }> = [];
const otherContents: string[] = [];
if (!request.contents) {
return { textContents, imageContents, audioContents, otherContents };
}
const contents = Array.isArray(request.contents)
? request.contents
: [request.contents];
for (const content of contents) {
this.processContent(
content,
textContents,
imageContents,
audioContents,
otherContents,
);
}
return { textContents, imageContents, audioContents, otherContents };
}
/**
* Process a single content item and add to appropriate arrays
*/
private processContent(
content: Content | string | PartUnion,
textContents: string[],
imageContents: Array<{ data: string; mimeType: string }>,
audioContents: Array<{ data: string; mimeType: string }>,
otherContents: string[],
): void {
if (typeof content === 'string') {
if (content.trim()) {
textContents.push(content);
}
return;
}
if ('parts' in content && content.parts) {
for (const part of content.parts) {
this.processPart(
part,
textContents,
imageContents,
audioContents,
otherContents,
);
}
return;
}
// Some request shapes (e.g. CountTokensParameters) allow passing parts directly
// instead of wrapping them in a { parts: [...] } Content object.
this.processPart(
content as Part | string,
textContents,
imageContents,
audioContents,
otherContents,
);
}
/**
* Process a single part and add to appropriate arrays
*/
private processPart(
part: Part | string,
textContents: string[],
imageContents: Array<{ data: string; mimeType: string }>,
audioContents: Array<{ data: string; mimeType: string }>,
otherContents: string[],
): void {
if (typeof part === 'string') {
if (part.trim()) {
textContents.push(part);
}
return;
}
if ('text' in part && part.text) {
textContents.push(part.text);
return;
}
if ('inlineData' in part && part.inlineData) {
const { data, mimeType } = part.inlineData;
if (mimeType && mimeType.startsWith('image/')) {
imageContents.push({ data: data || '', mimeType });
return;
}
if (mimeType && mimeType.startsWith('audio/')) {
audioContents.push({ data: data || '', mimeType });
return;
}
}
if ('fileData' in part && part.fileData) {
otherContents.push(JSON.stringify(part.fileData));
return;
}
if ('functionCall' in part && part.functionCall) {
otherContents.push(JSON.stringify(part.functionCall));
return;
}
if ('functionResponse' in part && part.functionResponse) {
otherContents.push(JSON.stringify(part.functionResponse));
return;
}
// Unknown part type - try to serialize
try {
const serialized = JSON.stringify(part);
if (serialized && serialized !== '{}') {
otherContents.push(serialized);
}
} catch (error) {
debugLogger.warn('Failed to serialize unknown part type:', error);
}
}
}
@@ -5,11 +5,15 @@
*/
/**
* Image MIME types the image tokenizer can decode for metadata extraction.
* This is a capability list, not an acceptance contract: the file-read path
* only forwards a narrower set to model endpoints (see
* PROVIDER_SAFE_IMAGE_MIME_TYPES in fileUtils.ts and #9291), so some types
* here are omitted from requests before they ever reach the tokenizer.
* Image MIME types accepted for vision input (attachment/thumbnail paths).
* This is an acceptance list for inputs, not a decode-capability list:
* token accounting uses the flat DEFAULT_IMAGE_TOKEN_ESTIMATE in
* compactionInputSlimming.ts — the former request-tokenizer estimator
* cluster, including ImageTokenizer and its dimension parsing, was removed
* as orphaned in PR #9676. The file-read path forwards only the narrower
* PIPELINE_IMAGE_MIME_TYPES subset to model endpoints (see
* PROVIDER_SAFE_IMAGE_MIME_TYPES in fileUtils.ts and #9291); anything else
* is omitted from requests with an in-band notice.
*/
export const SUPPORTED_IMAGE_MIME_TYPES = [
'image/bmp',
@@ -58,7 +62,7 @@ export function isSupportedImageMimeType(
/**
* Get a human-readable list of image formats the pipeline forwards to the
* model (not the tokenizer's wider decode capability).
* model (the narrower pipeline subset, not the full acceptance list above).
* @returns Comma-separated string of forwarded formats
*/
export function getSupportedImageFormatsString(): string {
@@ -1,36 +0,0 @@
/**
* @license
* Copyright 2025 Qwen
* SPDX-License-Identifier: Apache-2.0
*/
/**
* Token calculation result for different content types
*/
export interface TokenCalculationResult {
/** Total tokens calculated */
totalTokens: number;
/** Breakdown by content type */
breakdown: {
textTokens: number;
imageTokens: number;
audioTokens: number;
otherTokens: number;
};
/** Processing time in milliseconds */
processingTime: number;
}
/**
* Image metadata extracted from base64 data
*/
export interface ImageMetadata {
/** Image width in pixels */
width: number;
/** Image height in pixels */
height: number;
/** MIME type of the image */
mimeType: string;
/** Size of the base64 data in bytes */
dataSize: number;
}