mirror of
https://github.com/QwenLM/qwen-code.git
synced 2026-11-07 06:19:43 -10:00
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:
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;
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
});
|
||||
}
|
||||
@@ -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;
|
||||
}
|
||||
Reference in new issue
Block a user