From 9ab29a9b2997401e2933b82aa06a300346cd74c1 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Sun, 15 Mar 2026 14:26:44 -0400 Subject: [PATCH] refactor(agents): token counting for media content in messages Introduced a new method to estimate token costs for image and document blocks in messages, improving the accuracy of token counting. This enhancement ensures that media content is properly accounted for, particularly for the Claude model, by integrating additional token estimation logic for various content types. Updated the token counting function to utilize this new method, enhancing overall reliability and functionality. --- api/server/controllers/agents/client.js | 9 +- packages/api/src/agents/client.ts | 164 +++++++++++++++++++++++- 2 files changed, 169 insertions(+), 4 deletions(-) diff --git a/api/server/controllers/agents/client.js b/api/server/controllers/agents/client.js index 763ebce8f0..98ac4d700d 100644 --- a/api/server/controllers/agents/client.js +++ b/api/server/controllers/agents/client.js @@ -25,6 +25,7 @@ const { loadAgent: loadAgentFn, createMultiAgentMapper, filterMalformedContentParts, + estimateMediaTokensForMessage, hydrateMissingIndexTokenCounts, } = require('@librechat/api'); const { @@ -1237,8 +1238,12 @@ class AgentClient extends BaseClient { * @returns {number} */ getTokenCountForMessage(message) { - const count = super.getTokenCountForMessage(message); - if (this.getEncoding() === 'claude') { + const isClaude = this.getEncoding() === 'claude'; + let count = super.getTokenCountForMessage(message); + count += estimateMediaTokensForMessage(message.content ?? message.text, isClaude, (text) => + this.getTokenCount(text), + ); + if (isClaude) { return Math.ceil(count * AgentClient.CLAUDE_TOKEN_CORRECTION); } return count; diff --git a/packages/api/src/agents/client.ts b/packages/api/src/agents/client.ts index 8d23df4c46..919737addf 100644 --- a/packages/api/src/agents/client.ts +++ b/packages/api/src/agents/client.ts @@ -1,6 +1,12 @@ import { logger } from '@librechat/data-schemas'; import { isAgentsEndpoint } from 'librechat-data-provider'; -import { labelContentByAgent, getTokenCountForMessage } from '@librechat/agents'; +import { + labelContentByAgent, + extractImageDimensions, + getTokenCountForMessage, + estimateOpenAIImageTokens, + estimateAnthropicImageTokens, +} from '@librechat/agents'; import type { MessageContentComplex } from '@librechat/agents'; import type { Agent, TMessage } from 'librechat-data-provider'; import type { BaseMessage } from '@langchain/core/messages'; @@ -33,12 +39,166 @@ export function payloadParser({ req, endpoint }: { req: ServerRequest; endpoint: * Verified empirically across content types via the count_tokens endpoint. */ const CLAUDE_TOKEN_CORRECTION = 1.1; +const IMAGE_TOKEN_SAFETY_MARGIN = 1.05; +const BASE64_BYTES_PER_PDF_PAGE = 75_000; +const PDF_TOKENS_PER_PAGE_CLAUDE = 2000; +const PDF_TOKENS_PER_PAGE_OPENAI = 1500; +const URL_DOCUMENT_FALLBACK_TOKENS = 2000; + +type ContentBlock = { + type?: string; + image_url?: string | { url?: string }; + source?: { type?: string; data?: string; media_type?: string; content?: unknown[] }; + source_type?: string; + mime_type?: string; + data?: string; + text?: string; +}; + +/** + * Estimates token cost for image and document blocks in a message's + * content array that BaseClient.getTokenCountForMessage skips. + * Covers: image_url, image, image_file, document, file. + */ +export function estimateMediaTokensForMessage( + content: unknown, + isClaude: boolean, + getTokenCount?: (text: string) => number, +): number { + if (!Array.isArray(content)) { + return 0; + } + let tokens = 0; + const pdfPerPage = isClaude ? PDF_TOKENS_PER_PAGE_CLAUDE : PDF_TOKENS_PER_PAGE_OPENAI; + + for (const block of content as ContentBlock[]) { + if (block == null || typeof block !== 'object' || typeof block.type !== 'string') { + continue; + } + const type = block.type; + + if (type === 'image_url' || type === 'image' || type === 'image_file') { + let base64Data: string | undefined; + if (type === 'image_url') { + const url = typeof block.image_url === 'string' ? block.image_url : block.image_url?.url; + if (typeof url === 'string' && url.startsWith('data:')) { + base64Data = url; + } + } else if (type === 'image') { + if (block.source?.type === 'base64' && typeof block.source.data === 'string') { + base64Data = block.source.data; + } + } + + if (base64Data == null) { + tokens += 1024; + continue; + } + const dims = extractImageDimensions(base64Data); + if (dims == null) { + tokens += 1024; + continue; + } + const imgTokens = isClaude + ? estimateAnthropicImageTokens(dims.width, dims.height) + : estimateOpenAIImageTokens(dims.width, dims.height); + tokens += Math.ceil(imgTokens * IMAGE_TOKEN_SAFETY_MARGIN); + continue; + } + + if (type === 'document' || type === 'file') { + // LangChain standard format + if (typeof block.source_type === 'string') { + if (block.source_type === 'text' && typeof block.text === 'string') { + tokens += + getTokenCount != null ? getTokenCount(block.text) : Math.ceil(block.text.length / 4); + continue; + } + if (block.source_type === 'base64' && typeof block.data === 'string') { + const mime = (block.mime_type ?? '').split(';')[0]; + if (mime === 'application/pdf' || mime === '') { + const pages = Math.max(1, Math.ceil(block.data.length / BASE64_BYTES_PER_PDF_PAGE)); + tokens += pages * pdfPerPage; + } else if (mime.startsWith('image/')) { + const dims = extractImageDimensions(block.data); + tokens += + dims != null + ? Math.ceil( + (isClaude + ? estimateAnthropicImageTokens(dims.width, dims.height) + : estimateOpenAIImageTokens(dims.width, dims.height)) * + IMAGE_TOKEN_SAFETY_MARGIN, + ) + : 1024; + } + continue; + } + tokens += URL_DOCUMENT_FALLBACK_TOKENS; + continue; + } + + // Anthropic format + if (block.source != null) { + if (block.source.type === 'text' && typeof block.source.data === 'string') { + tokens += + getTokenCount != null + ? getTokenCount(block.source.data) + : Math.ceil(block.source.data.length / 4); + continue; + } + if (block.source.type === 'base64' && typeof block.source.data === 'string') { + const mime = (block.source.media_type ?? '').split(';')[0]; + if (mime === 'application/pdf' || mime === '') { + const pages = Math.max( + 1, + Math.ceil(block.source.data.length / BASE64_BYTES_PER_PDF_PAGE), + ); + tokens += pages * pdfPerPage; + } + continue; + } + if (block.source.type === 'url') { + tokens += URL_DOCUMENT_FALLBACK_TOKENS; + continue; + } + if (block.source.type === 'content' && Array.isArray(block.source.content)) { + for (const inner of block.source.content) { + const innerBlock = inner as ContentBlock | null; + if ( + innerBlock?.type === 'image' && + innerBlock.source?.type === 'base64' && + typeof innerBlock.source.data === 'string' + ) { + const dims = extractImageDimensions(innerBlock.source.data); + tokens += + dims != null + ? Math.ceil( + (isClaude + ? estimateAnthropicImageTokens(dims.width, dims.height) + : estimateOpenAIImageTokens(dims.width, dims.height)) * + IMAGE_TOKEN_SAFETY_MARGIN, + ) + : 1024; + } + } + continue; + } + } + tokens += URL_DOCUMENT_FALLBACK_TOKENS; + } + } + return tokens; +} export function createTokenCounter(encoding: Parameters[1]) { const isClaude = encoding === 'claude'; return function (message: BaseMessage) { const countTokens = (text: string) => Tokenizer.getTokenCount(text, encoding); - const count = getTokenCountForMessage(message, countTokens); + const count = getTokenCountForMessage( + message, + countTokens, + encoding as 'claude' | 'o200k_base', + ); return isClaude ? Math.ceil(count * CLAUDE_TOKEN_CORRECTION) : count; }; }