diff --git a/api/package.json b/api/package.json index 043fcdd7ee..44bf9127c0 100644 --- a/api/package.json +++ b/api/package.json @@ -46,7 +46,7 @@ "@azure/storage-blob": "^12.30.0", "@google/genai": "^2.8.0", "@keyv/redis": "5.1.6", - "@librechat/agents": "^3.6.12", + "@librechat/agents": "^3.6.13", "@librechat/api": "*", "@librechat/data-schemas": "*", "@microsoft/microsoft-graph-client": "^3.0.7", diff --git a/api/server/controllers/agents/client.js b/api/server/controllers/agents/client.js index cb37cd6a31..0332bdddc1 100644 --- a/api/server/controllers/agents/client.js +++ b/api/server/controllers/agents/client.js @@ -17,7 +17,7 @@ const { omitTitleOptions, getProviderConfig, memoryInstructions, - createTokenCounter, + createCachedTokenCounter, applyContextToAgent, isMemoryAgentEnabled, recordCollectedUsage, @@ -3249,7 +3249,7 @@ class AgentClient extends BaseClient { }; const toolSet = buildToolSet(this.options.agent); - const tokenCounter = createTokenCounter(this.getEncoding()); + const tokenCounter = await createCachedTokenCounter(this.getEncoding()); /** Pre-resolve invoked skill bodies + re-prime files before formatting messages */ const skillPrimeResult = this.options.primeInvokedSkills @@ -3889,7 +3889,7 @@ class AgentClient extends BaseClient { this.contentParts.push(...seedContent); } - const tokenCounter = createTokenCounter(this.getEncoding()); + const tokenCounter = await createCachedTokenCounter(this.getEncoding()); const agents = collectReachableAgents([ this.options.agent, ...(this.agentConfigs?.size > 0 ? this.agentConfigs.values() : []), diff --git a/api/server/controllers/agents/client.test.js b/api/server/controllers/agents/client.test.js index 37d4b54094..027feffe12 100644 --- a/api/server/controllers/agents/client.test.js +++ b/api/server/controllers/agents/client.test.js @@ -47,7 +47,7 @@ jest.mock('@librechat/api', () => ({ createRun: (...args) => mockCreateRun(...args), countFormattedMessageTokens: jest.fn(() => 42), countTokens: jest.fn((text) => Math.ceil(String(text ?? '').length / 4)), - createTokenCounter: jest.fn(() => jest.fn(() => 0)), + createCachedTokenCounter: jest.fn(async () => jest.fn(() => 0)), createDetachedSubagentUsageRecorder: (...args) => mockCreateDetachedSubagentUsageRecorder(...args), captureAgentCheckpointGeneration: (...args) => mockCaptureAgentCheckpointGeneration(...args), diff --git a/package-lock.json b/package-lock.json index c0c365f79d..28e06d1b14 100644 --- a/package-lock.json +++ b/package-lock.json @@ -63,7 +63,7 @@ "@azure/storage-blob": "^12.30.0", "@google/genai": "^2.8.0", "@keyv/redis": "5.1.6", - "@librechat/agents": "^3.6.12", + "@librechat/agents": "^3.6.13", "@librechat/api": "*", "@librechat/data-schemas": "*", "@microsoft/microsoft-graph-client": "^3.0.7", @@ -10630,9 +10630,9 @@ } }, "node_modules/@librechat/agents": { - "version": "3.6.12", - "resolved": "https://registry.npmjs.org/@librechat/agents/-/agents-3.6.12.tgz", - "integrity": "sha512-S73bRfCJ4bklASObib7gpSn9Vm0IkJ9ZmVgtzcKzUYju8lkb18vU9FcK1ZFcyZl7/FZAQbHIHRYBDKHMqEn7xg==", + "version": "3.6.13", + "resolved": "https://registry.npmjs.org/@librechat/agents/-/agents-3.6.13.tgz", + "integrity": "sha512-nrcxEKwJDmIDsGf7txLqTgqqSKkJKl90xS5vqTtAoYSi4/EVPZ6ZFpJHLT3WMn2nWL5jjxdOmstet3JbSF1B+A==", "license": "MIT", "dependencies": { "@anthropic-ai/sdk": "^0.115.0", @@ -42822,7 +42822,7 @@ "@azure/storage-blob": "^12.30.0", "@google/genai": "^2.8.0", "@keyv/redis": "5.1.6", - "@librechat/agents": "^3.6.12", + "@librechat/agents": "^3.6.13", "@librechat/data-schemas": "*", "@modelcontextprotocol/sdk": "^1.30.0", "@opentelemetry/api": "^1.9.0", diff --git a/packages/api/package.json b/packages/api/package.json index b6006d72af..1b18249938 100644 --- a/packages/api/package.json +++ b/packages/api/package.json @@ -113,7 +113,7 @@ "@azure/storage-blob": "^12.30.0", "@google/genai": "^2.8.0", "@keyv/redis": "5.1.6", - "@librechat/agents": "^3.6.12", + "@librechat/agents": "^3.6.13", "@librechat/data-schemas": "*", "@modelcontextprotocol/sdk": "^1.30.0", "@opentelemetry/api": "^1.9.0", diff --git a/packages/api/src/agents/client.spec.ts b/packages/api/src/agents/client.spec.ts index c62cb09662..fa62cb4bb1 100644 --- a/packages/api/src/agents/client.spec.ts +++ b/packages/api/src/agents/client.spec.ts @@ -1,6 +1,10 @@ import { ContentTypes } from 'librechat-data-provider'; +import { Tokenizer as AiTokenizer } from 'ai-tokenizer'; +import { Providers, StandardGraph } from '@librechat/agents'; +import { HumanMessage } from '@librechat/agents/langchain/messages'; import type { TMessage } from 'librechat-data-provider'; import { + createCachedTokenCounter, prependQuotes, prependFileContext, applyAttachmentOnlyText, @@ -8,6 +12,39 @@ import { } from './client'; import { ATTACHMENT_ONLY_TEXT } from '~/files/context'; +describe('createCachedTokenCounter', () => { + it('enables stable-message reuse in the agents runtime', async () => { + const getTokenCount = jest.spyOn(AiTokenizer.prototype, 'count'); + try { + const tokenCounter = await createCachedTokenCounter('o200k_base'); + const graph = new StandardGraph({ + runId: 'token-cache-integration', + agents: [ + { + agentId: 'primary', + provider: Providers.OPENAI, + instructions: 'Test instructions', + }, + ], + tokenCounter, + }); + const agentContext = graph.agentContexts.get('primary'); + await agentContext?.tokenCalculationPromise; + getTokenCount.mockClear(); + const message = new HumanMessage('Stable retained context'); + + agentContext?.contextPressureTokenCounts?.count(message); + const callsAfterFirstCount = getTokenCount.mock.calls.length; + agentContext?.contextPressureTokenCounts?.count(message); + + expect(callsAfterFirstCount).toBeGreaterThan(0); + expect(getTokenCount).toHaveBeenCalledTimes(callsAfterFirstCount); + } finally { + getTokenCount.mockRestore(); + } + }); +}); + describe('prependFileContext', () => { it('prepends file context to string content', () => { const message: FormattedMessageWithContent = { content: 'Answer this question.' }; diff --git a/packages/api/src/agents/client.ts b/packages/api/src/agents/client.ts index bc7ca9c404..fe00e5ec2a 100644 --- a/packages/api/src/agents/client.ts +++ b/packages/api/src/agents/client.ts @@ -6,9 +6,10 @@ import { getTokenCountForMessage, estimateOpenAIImageTokens, estimateAnthropicImageTokens, + markTokenCounterCacheCompatible, } from '@librechat/agents'; +import type { MessageContentComplex, TokenCounter } from '@librechat/agents'; import type { BaseMessage } from '@librechat/agents/langchain/messages'; -import type { MessageContentComplex } from '@librechat/agents'; import type { Agent, TMessage } from 'librechat-data-provider'; import type { ServerRequest } from '~/types'; import { getSafeErrorMetadata, mergeQuotedText, formatQuotesAsMarkdown } from '~/utils'; @@ -378,9 +379,19 @@ export function countFormattedMessageTokens( return isClaude ? Math.ceil(numTokens * CLAUDE_TOKEN_CORRECTION) : numTokens; } -export function createTokenCounter(encoding: Parameters[1]) { +export function createTokenCounter( + encoding: Parameters[1], +): TokenCounter { + return createMessageTokenCounter(encoding, (text: string) => + Tokenizer.getTokenCount(text, encoding), + ); +} + +function createMessageTokenCounter( + encoding: Parameters[1], + countTokens: (text: string) => number, +): TokenCounter { const isClaude = encoding === 'claude'; - const countTokens = (text: string) => Tokenizer.getTokenCount(text, encoding); return function (message: BaseMessage): number { const count = getTokenCountForMessage( message, @@ -391,6 +402,13 @@ export function createTokenCounter(encoding: Parameters[1], +): Promise { + const countTokens = await Tokenizer.createExactTokenCounter(encoding ?? 'o200k_base'); + return markTokenCounterCacheCompatible(createMessageTokenCounter(encoding, countTokens)); +} + export function logToolError(_graph: unknown, error: unknown, toolId: string): void { /** * A GraphInterrupt unwinding out of a tool body is the HITL pause working as diff --git a/packages/api/src/utils/tokenizer.spec.ts b/packages/api/src/utils/tokenizer.spec.ts index b8c1bd8d98..d2a4ccb6a5 100644 --- a/packages/api/src/utils/tokenizer.spec.ts +++ b/packages/api/src/utils/tokenizer.spec.ts @@ -1,3 +1,4 @@ +import { Tokenizer as AiTokenizer } from 'ai-tokenizer'; import Tokenizer from './tokenizer'; jest.mock('@librechat/data-schemas', () => ({ @@ -53,4 +54,19 @@ describe('Tokenizer', () => { expect(count).toBeGreaterThan(0); }); }); + + describe('createExactTokenCounter', () => { + it('throws instead of returning a fallback estimate after a tokenizer failure', async () => { + const counter = await Tokenizer.createExactTokenCounter('o200k_base'); + const count = jest.spyOn(AiTokenizer.prototype, 'count').mockImplementationOnce(() => { + throw new Error('tokenizer failed'); + }); + try { + expect(() => counter('Stable cached context')).toThrow('tokenizer failed'); + } finally { + count.mockRestore(); + await Tokenizer.initEncoding('o200k_base'); + } + }); + }); }); diff --git a/packages/api/src/utils/tokenizer.ts b/packages/api/src/utils/tokenizer.ts index d66908712a..83a3c2691f 100644 --- a/packages/api/src/utils/tokenizer.ts +++ b/packages/api/src/utils/tokenizer.ts @@ -27,6 +27,23 @@ class Tokenizer { return this.loadingPromises[encoding]; } + /** Returns a counter that never substitutes an estimate for tokenizer errors. */ + async createExactTokenCounter(encoding: EncodingName): Promise<(text: string) => number> { + await this.initEncoding(encoding); + const tokenizer = this.tokenizersCache[encoding]; + if (!tokenizer) { + throw new Error(`Tokenizer encoding failed to initialize: ${encoding}`); + } + return (text: string): number => { + try { + return tokenizer.count(text); + } catch (error) { + this.handleCountError(encoding, tokenizer, error); + throw error; + } + }; + } + getTokenCount(text: string, encoding: EncodingName = 'o200k_base'): number { const tokenizer = this.tokenizersCache[encoding]; if (!tokenizer) { @@ -36,13 +53,22 @@ class Tokenizer { try { return tokenizer.count(text); } catch (error) { - logger.error('[Tokenizer] Error getting token count:', error); - delete this.tokenizersCache[encoding]; - delete this.loadingPromises[encoding]; - this.initEncoding(encoding); + this.handleCountError(encoding, tokenizer, error); return Math.ceil(text.length / 4); } } + + private handleCountError(encoding: EncodingName, tokenizer: AiTokenizer, error: unknown): void { + logger.error('[Tokenizer] Error getting token count:', error); + if (this.tokenizersCache[encoding] !== tokenizer) { + return; + } + delete this.tokenizersCache[encoding]; + delete this.loadingPromises[encoding]; + void this.initEncoding(encoding).catch((initError: unknown) => { + logger.error(`[Tokenizer] Error reloading ${encoding} encoding:`, initError); + }); + } } const TokenizerSingleton: Tokenizer = new Tokenizer();