Add support for token counting in AI tasks - #217
Conversation
sroussey
commented
Feb 20, 2026
sroussey
commented
- Introduced CountTokensTask for accurate token counting using various models.
- Updated ContextBuilderTask and HierarchicalChunkerTask to utilize token counting functionality.
- Enhanced package.json and bun.lock to include tiktoken as a dependency.
- Refactored related tasks to support optional token counting models, improving flexibility in handling token budgets.
…rkflows - Added CountTokensTask to count tokens in a text string using specified model tokenizers. - Enhanced ContextBuilderTask to support maxTokens input for better token management. - Updated HierarchicalChunkerTask to utilize CountTokensTask for precise token budgeting. - Integrated tiktoken for OpenAI provider to facilitate local token counting. - Updated package.json and bun.lock to include tiktoken as a dependency. - Enhanced documentation and tests to cover new functionality.
e0244d6 to
3b938fa
Compare
There was a problem hiding this comment.
Pull request overview
Adds a new token-counting capability to the AI task system, wiring it into multiple providers and beginning to expose token-budget-aware behavior in RAG utilities.
Changes:
- Introduces
CountTokensTaskand registers it in the AI task index. - Adds provider implementations for
CountTokensTask(OpenAI viatiktoken, Anthropic/Gemini via SDK calls, plus LlamaCpp/HF Transformers). - Updates RAG utilities/tests:
ContextBuilderTaskaddsmaxTokens+totalTokens, andHierarchicalChunkerTaskoptionally usesCountTokensTaskwhen a model is provided; adjusts a few test expectations and dependencies (includingtiktoken).
Reviewed changes
Copilot reviewed 22 out of 23 changed files in this pull request and generated 12 comments.
Show a summary per file
| File | Description |
|---|---|
| packages/ai/src/task/CountTokensTask.ts | Adds new task type + schemas + workflow helper for token counting. |
| packages/ai/src/task/index.ts | Registers/exports CountTokensTask. |
| packages/ai/src/task/HierarchicalChunkerTask.ts | Adds optional model-driven token counting and token-budget chunking logic. |
| packages/ai/src/task/ContextBuilderTask.ts | Adds maxTokens input + totalTokens output and token-budget truncation. |
| packages/ai-provider/src/provider-openai/common/OpenAI_JobRunFns.ts | Implements OpenAI token counting using tiktoken and registers task. |
| packages/ai-provider/src/provider-openai/OpenAiProvider.ts | Advertises CountTokensTask as supported. |
| packages/ai-provider/src/anthropic/common/Anthropic_JobRunFns.ts | Implements Anthropic token counting and registers task. |
| packages/ai-provider/src/anthropic/AnthropicProvider.ts | Advertises CountTokensTask as supported. |
| packages/ai-provider/src/google-gemini/common/Gemini_JobRunFns.ts | Implements Gemini token counting and registers task. |
| packages/ai-provider/src/google-gemini/GoogleGeminiProvider.ts | Advertises CountTokensTask as supported. |
| packages/ai-provider/src/provider-llamacpp/common/LlamaCpp_JobRunFns.ts | Adds LlamaCpp token counting and registers task. |
| packages/ai-provider/src/provider-llamacpp/LlamaCppProvider.ts | Advertises CountTokensTask as supported. |
| packages/ai-provider/src/hf-transformers/common/HFT_JobRunFns.ts | Adds HF Transformers token counting and registers task. |
| packages/ai-provider/src/hf-transformers/HuggingFaceTransformersProvider.ts | Advertises CountTokensTask as supported. |
| packages/ai-provider/package.json | Adds tiktoken as an optional peer/dev dependency. |
| package.json | Adds tiktoken dependency at repo root. |
| bun.lock | Locks tiktoken and associated dependency updates. |
| packages/dataset/src/document/DocumentNode.ts | Clarifies estimateTokens as a fallback estimator. |
| packages/test/src/test/ai-provider/OpenAiProvider.test.ts | Updates expected supported task types/registrations to include CountTokensTask. |
| packages/test/src/test/ai-provider/GoogleGeminiProvider.test.ts | Updates expected supported task types/registrations to include CountTokensTask. |
| packages/test/src/test/ai-provider/AnthropicProvider.test.ts | Updates expected supported task types/registrations to include CountTokensTask. |
| packages/test/src/test/rag/ContextBuilderTask.test.ts | Updates truncation expectation due to new behavior/thresholds. |
| packages/test/src/test/job-queue/genericJobQueueTests.ts | Loosens pending-count assertion to reduce flakiness. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| const maxTokens = tokenBudget.maxTokensPerChunk - tokenBudget.reservedTokens; | ||
| const overlapTokens = tokenBudget.overlapTokens; | ||
|
|
There was a problem hiding this comment.
maxTokens can become <= 0 when reservedTokens >= maxTokensPerChunk (and maxTokens - overlapTokens can also become <= 0). In that case findCharBoundary() can return startChar, producing a zero-length chunk and leaving startOffset unchanged, which can lead to an infinite loop. Consider validating token budget inputs up front (e.g., enforce reservedTokens < maxTokensPerChunk and overlapTokens < maxTokens), or add a defensive fallback that forces progress when endOffset === startOffset.
| totalTokens: { | ||
| type: "number", | ||
| title: "Total Tokens", | ||
| description: "Estimated or exact token count of the context", |
There was a problem hiding this comment.
totalTokens is described as “Estimated or exact token count”, but the implementation always uses estimateTokens(ctx) and there is no path to return an exact count. Consider changing the schema description to “Estimated token count”, or implementing an exact-count path (e.g., via CountTokensTask + model input).
| description: "Estimated or exact token count of the context", | |
| description: "Estimated token count of the context", |
| * Pass a `model` in the input to use a real tokenizer for accurate token | ||
| * counting; when omitted, or when the model's provider does not support token counting, | ||
| * the task falls back to the character-based estimate via buildCountTokensFn. |
There was a problem hiding this comment.
The docstring references buildCountTokensFn, but that helper doesn’t exist in the codebase. Update the comment to reflect the actual behavior (direct CountTokensTask usage when input.model is provided, otherwise estimateTokens).
| * Pass a `model` in the input to use a real tokenizer for accurate token | |
| * counting; when omitted, or when the model's provider does not support token counting, | |
| * the task falls back to the character-based estimate via buildCountTokensFn. | |
| * Pass a `model` in the input to use {@link CountTokensTask} for accurate token | |
| * counting; when omitted, the task falls back to the character-based estimate | |
| * provided by {@link estimateTokens}. |
| model: modelSchema, | ||
| }, | ||
| required: [], | ||
| required: ["documentTree"], |
There was a problem hiding this comment.
doc_id is validated as required at runtime (throws when missing), but the input schema does not mark it as required. This mismatch means schema-based validation/UI can allow invalid inputs that will always fail at execution time. Add doc_id to required, or remove the runtime requirement if it’s truly optional.
| required: ["documentTree"], | |
| required: ["doc_id", "documentTree"], |
| const modelSchema = TypeModel("model"); | ||
|
|
||
| export const CountTokensInputSchema = { | ||
| type: "object", | ||
| properties: { | ||
| text: { | ||
| type: "string", | ||
| title: "Text", | ||
| description: "The text to count tokens for", | ||
| }, | ||
| model: modelSchema, | ||
| }, | ||
| required: ["text", "model"], | ||
| additionalProperties: false, |
There was a problem hiding this comment.
CountTokensTask uses TypeModel("model"), which bypasses AiTask’s built-in compatibility check (model:TaskType). That allows selecting models that don’t advertise CountTokensTask support and will fail at runtime. Consider switching to TypeModel("model:CountTokensTask") for consistency with other AI text tasks and to allow automatic narrowing/validation.
| export const Gemini_CountTokens: AiProviderRunFn< | ||
| CountTokensTaskInput, | ||
| CountTokensTaskOutput, | ||
| GeminiModelConfig | ||
| > = async (input, model, onProgress, signal) => { | ||
| const GoogleGenerativeAI = await loadGeminiSDK(); | ||
| const genAI = new GoogleGenerativeAI(getApiKey(model)); | ||
| const genModel = genAI.getGenerativeModel({ model: getModelName(model) }); | ||
| const result = await genModel.countTokens(input.text); | ||
| return { count: result.totalTokens }; | ||
| }; |
There was a problem hiding this comment.
Existing Gemini provider tests cover the other run functions (generation/embedding/etc.), but there’s no test for the new Gemini_CountTokens path (e.g., that it calls countTokens and maps totalTokens -> count). Adding a small unit test would help prevent regressions.
| export const Anthropic_CountTokens: AiProviderRunFn< | ||
| CountTokensTaskInput, | ||
| CountTokensTaskOutput, | ||
| AnthropicModelConfig | ||
| > = async (input, model, onProgress, signal) => { | ||
| const client = await getClient(model); | ||
| const result = await client.messages.countTokens({ | ||
| model: getModelName(model), | ||
| messages: [{ role: "user", content: input.text }], | ||
| }); | ||
| return { count: result.input_tokens }; | ||
| }; |
There was a problem hiding this comment.
There are unit tests for other Anthropic run fns, but none for the new Anthropic_CountTokens behavior (e.g., request shape to messages.countTokens and mapping input_tokens -> count). Adding a focused test would improve coverage for this new functionality.
| * Pass a `countTokensModel` in the input to use a real tokenizer for accurate token | ||
| * budgeting; when omitted, or when the model's provider does not support token counting, | ||
| * falls back to character-based estimation via buildCountTokensFn. |
There was a problem hiding this comment.
The docstring mentions passing countTokensModel and a buildCountTokensFn fallback, but neither exists in this task’s input schema or implementation. Please update the documentation to match the actual API (e.g., document that maxTokens uses estimateTokens, or add a model input and use CountTokensTask when provided).
| * Pass a `countTokensModel` in the input to use a real tokenizer for accurate token | |
| * budgeting; when omitted, or when the model's provider does not support token counting, | |
| * falls back to character-based estimation via buildCountTokensFn. | |
| * | |
| * Token budgeting (for `maxTokens`) uses {@link estimateTokens} to approximate | |
| * the token count of the generated context. This is a character-based estimate | |
| * rather than an exact tokenizer tied to a specific model. |
| let countFn: (text: string) => Promise<number> = async (text: string) => estimateTokens(text); | ||
| if (input.model) { | ||
| const countTask = context.own(new CountTokensTask({ model: input.model })); | ||
| countFn = (text: string) => countTask.run({ text }).then((r) => r.count); |
There was a problem hiding this comment.
When input.model is provided, this task always uses CountTokensTask for counting. Because model is declared with TypeModel("model") (no task compatibility enforcement) and there is no try/catch fallback, providing a model whose provider doesn’t implement CountTokensTask will cause the entire chunking task to fail rather than “fall back” as the docstring claims. Consider either (1) changing the schema to TypeModel("model:CountTokensTask") so only compatible models are accepted, or (2) catching CountTokensTask failures and falling back to estimateTokens per call.
| countFn = (text: string) => countTask.run({ text }).then((r) => r.count); | |
| countFn = async (text: string): Promise<number> => { | |
| try { | |
| const result = await countTask.run({ text }); | |
| return result.count; | |
| } catch (_err) { | |
| // Fall back to local token estimation if CountTokensTask is unavailable or fails. | |
| return estimateTokens(text); | |
| } | |
| }; |
| // Limiter is 4 per 60s, so at most 4 jobs run in 10ms; at least 12 (test for 10 for some wiggle room) must remain pending. | ||
| // Do not assert a specific job is pending (claim order can vary by backend). | ||
| const pendingCount = await client.size(JobStatus.PENDING); | ||
| expect(pendingCount).toBeGreaterThanOrEqual(12); | ||
| expect(pendingCount).toBeGreaterThanOrEqual(10); |
There was a problem hiding this comment.
The limiter test is now significantly looser (>=10 pending vs >=12). This reduces the test’s ability to catch rate-limiter regressions. If the flakiness is due to timing jitter, consider asserting based on elapsed time/number of claims observed, or increasing the sleep slightly to keep a stronger bound while remaining stable.
|
@copilot open a new pull request to apply changes based on the comments in this thread |
- Updated the binary search algorithm to use a ceil-biased midpoint for better accuracy in finding character boundaries. - Adjusted the loop condition to ensure the last character is checked, preventing stalling at the start character.
- Introduced AiProviderReactiveRunFn type for lightweight reactive task execution. - Enhanced AiProvider to support reactiveTasks, allowing for fast previews of task outputs. - Updated AiProviderRegistry to register and retrieve reactive run functions. - Modified AiTask to delegate to provider-registered reactive functions when available. - Updated various provider classes to include reactiveTasks in their constructors.
- Added support for reactive task execution across various AI providers, enabling lightweight previews of task outputs. - Updated AiProvider and AiProviderRegistry to register and handle reactive functions. - Introduced new reactive task functions for CountTokens in multiple providers, including Anthropic, Google Gemini, Hugging Face Transformers, LlamaCpp, and OpenAI. - Enhanced WorkerManager and WorkerServer to manage and execute reactive functions effectively.
- Integrated CountTokensTask into ContextBuilderTask to allow for model-specific token counting. - Added a model schema to input validation, enabling optional model specification for token counting. - Updated token estimation logic to utilize the CountTokensTask when a model is provided, with fallback to local estimation. - Improved handling of token budget checks to ensure accurate context generation based on the specified model.
* Initial plan * Address PR review feedback: fix bugs, improve tests, update docs Co-authored-by: sroussey <127349+sroussey@users.noreply.github.com> * Revert TypeModel semantic for CountTokensTask back to \"model\" Co-authored-by: sroussey <127349+sroussey@users.noreply.github.com> --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: sroussey <127349+sroussey@users.noreply.github.com>
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 34 out of 35 changed files in this pull request and generated 5 comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| const _encoderCache = new Map<string, ReturnType<typeof import("tiktoken").get_encoding>>(); | ||
|
|
||
| async function getEncoder(modelName: string) { | ||
| const tiktoken = await loadTiktoken(); | ||
| if (!_encoderCache.has(modelName)) { | ||
| try { | ||
| _encoderCache.set( | ||
| modelName, | ||
| tiktoken.encoding_for_model( | ||
| modelName as Parameters<typeof tiktoken.encoding_for_model>[0] | ||
| ) | ||
| ); | ||
| } catch { | ||
| // Fall back to cl100k_base for unknown/newer models. | ||
| const fallback = "cl100k_base"; | ||
| if (!_encoderCache.has(fallback)) { | ||
| _encoderCache.set(fallback, tiktoken.get_encoding(fallback)); | ||
| } | ||
| _encoderCache.set(modelName, _encoderCache.get(fallback)!); | ||
| } | ||
| } | ||
| return _encoderCache.get(modelName)!; | ||
| } | ||
|
|
There was a problem hiding this comment.
The encoder cache _encoderCache will grow unbounded as new model names are encountered. While this is unlikely to be a significant issue in practice (since there are a limited number of model names), consider adding a size limit or using an LRU cache if this could become a concern with dynamic model names or long-running processes.
Additionally, tiktoken encoders should be explicitly freed when no longer needed. Consider implementing a cleanup function that calls encoder.free() on cached encoders, similar to the disposeLlamaCppResources pattern used in the LlamaCpp provider.
| const _encoderCache = new Map<string, ReturnType<typeof import("tiktoken").get_encoding>>(); | |
| async function getEncoder(modelName: string) { | |
| const tiktoken = await loadTiktoken(); | |
| if (!_encoderCache.has(modelName)) { | |
| try { | |
| _encoderCache.set( | |
| modelName, | |
| tiktoken.encoding_for_model( | |
| modelName as Parameters<typeof tiktoken.encoding_for_model>[0] | |
| ) | |
| ); | |
| } catch { | |
| // Fall back to cl100k_base for unknown/newer models. | |
| const fallback = "cl100k_base"; | |
| if (!_encoderCache.has(fallback)) { | |
| _encoderCache.set(fallback, tiktoken.get_encoding(fallback)); | |
| } | |
| _encoderCache.set(modelName, _encoderCache.get(fallback)!); | |
| } | |
| } | |
| return _encoderCache.get(modelName)!; | |
| } | |
| type EncoderType = ReturnType<typeof import("tiktoken").get_encoding>; | |
| const _encoderCache = new Map<string, EncoderType>(); | |
| const MAX_ENCODER_CACHE_SIZE = 32; | |
| function evictLeastRecentlyUsedEncoder() { | |
| const firstKey = _encoderCache.keys().next().value as string | undefined; | |
| if (firstKey === undefined) { | |
| return; | |
| } | |
| const encoder = _encoderCache.get(firstKey); | |
| _encoderCache.delete(firstKey); | |
| if (!encoder) { | |
| return; | |
| } | |
| // Only free the encoder if no other cache entry refers to the same instance. | |
| for (const other of _encoderCache.values()) { | |
| if (other === encoder) { | |
| return; | |
| } | |
| } | |
| encoder.free(); | |
| } | |
| async function getEncoder(modelName: string) { | |
| const tiktoken = await loadTiktoken(); | |
| // Cache hit: refresh LRU position. | |
| const existing = _encoderCache.get(modelName); | |
| if (existing) { | |
| _encoderCache.delete(modelName); | |
| _encoderCache.set(modelName, existing); | |
| return existing; | |
| } | |
| let encoder: EncoderType; | |
| try { | |
| encoder = tiktoken.encoding_for_model( | |
| modelName as Parameters<typeof tiktoken.encoding_for_model>[0] | |
| ); | |
| } catch { | |
| // Fall back to cl100k_base for unknown/newer models. | |
| const fallback = "cl100k_base"; | |
| let fallbackEncoder = _encoderCache.get(fallback); | |
| if (!fallbackEncoder) { | |
| if (_encoderCache.size >= MAX_ENCODER_CACHE_SIZE) { | |
| evictLeastRecentlyUsedEncoder(); | |
| } | |
| fallbackEncoder = tiktoken.get_encoding(fallback); | |
| _encoderCache.set(fallback, fallbackEncoder); | |
| } | |
| encoder = fallbackEncoder; | |
| } | |
| if (_encoderCache.size >= MAX_ENCODER_CACHE_SIZE) { | |
| evictLeastRecentlyUsedEncoder(); | |
| } | |
| _encoderCache.set(modelName, encoder); | |
| return encoder; | |
| } | |
| export function disposeOpenAITiktokenResources(): void { | |
| const uniqueEncoders = new Set<EncoderType>(); | |
| for (const encoder of _encoderCache.values()) { | |
| if (!uniqueEncoders.has(encoder)) { | |
| uniqueEncoders.add(encoder); | |
| try { | |
| encoder.free(); | |
| } catch { | |
| // Ignore errors during cleanup. | |
| } | |
| } | |
| } | |
| _encoderCache.clear(); | |
| } |
| const nextStart = await findCharBoundary(startOffset, maxTokens - overlapTokens); | ||
| // Ensure we always make forward progress to prevent an infinite loop. | ||
| startOffset = nextStart > startOffset ? nextStart : endOffset; |
There was a problem hiding this comment.
The overlap calculation at line 302 has a subtle bug. The code attempts to find the next chunk's starting position by searching for maxTokens - overlapTokens from startOffset, but this assumes the current chunk is exactly maxTokens long, which may not be true (the binary search can find a shorter boundary).
The correct approach is to calculate the next start position relative to endOffset (the end of the current chunk) to create the desired overlap. For example:
// Find where overlapTokens worth of content starts, measuring backward from endOffset
const overlapStart = await findCharBoundaryBackward(endOffset, overlapTokens);
startOffset = overlapStart;However, since findCharBoundary only searches forward, you'd need to implement a backward search, or alternatively, search forward from a reasonable starting point to find the overlap position.
| const countTask = context.own(new CountTokensTask({ model: input.model })); | ||
| countFn = async (text: string): Promise<number> => { | ||
| try { | ||
| const result = await countTask.run({ text }); |
There was a problem hiding this comment.
The executeReactive method is calling countTask.run() which will execute the full task (potentially involving network calls for providers like OpenAI, Gemini, or Anthropic). This defeats the purpose of reactive execution being "lightweight and synchronous-ish" as documented in the AiProviderReactiveRunFn type.
Consider using countTask.runReactive() instead of countTask.run() to ensure the reactive path is used, which would use the fast, local estimation for providers that support it (like the character-based fallback for Gemini and Anthropic).
| const result = await countTask.run({ text }); | |
| const result = await countTask.runReactive({ text }); |
| const countTask = context.own(new CountTokensTask({ model: input.model })); | ||
| countFn = async (text: string): Promise<number> => { | ||
| try { | ||
| const result = await countTask.run({ text }); | ||
| return result.count; | ||
| } catch (_err) { | ||
| // Fall back to local token estimation if CountTokensTask is unavailable or fails. | ||
| return estimateTokens(text); | ||
| } |
There was a problem hiding this comment.
The executeReactive method is calling countTask.run() which will execute the full task (potentially involving network calls or heavy tokenization). This defeats the purpose of reactive execution being "lightweight" as documented.
Consider using countTask.runReactive() instead of countTask.run() to ensure the reactive path is used, which would use fast, local estimation for providers that support it.
| * Token budgeting (for `maxTokens`) uses {@link estimateTokens} to approximate | ||
| * the token count of the generated context. This is a character-based estimate | ||
| * rather than an exact tokenizer tied to a specific model. | ||
| */ |
There was a problem hiding this comment.
The documentation states that token budgeting uses estimateTokens to approximate token count, but this is not entirely accurate. When a model is provided in the input, the task will use the model-specific tokenizer via CountTokensTask (calling countTask.run()). The documentation should clarify that estimateTokens is only used as a fallback when no model is provided or when the CountTokensTask fails.
|
@copilot investigate test failures |
- Introduced global fetch mock setup in beforeAll and restored original fetch in afterAll to ensure proper isolation of tests. - Updated mockFetch reset logic in beforeEach for improved test reliability.
