CERBERUS/LibreChat
1
0
Fork 0
mirror of https://github.com/danny-avila/LibreChat.git synced 2026-10-21 01:22:54 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 6

refactor(agents): token counting for media content in messages

Browse source 
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.
This commit is contained in:
Danny Avila 2026-03-15 14:26:44 -04:00
parent d2a4d9fbee
commit 9ab29a9b29
No known key found for this signature in database
GPG key ID: BF31EEB2C5CA0956
2 changed files with 169 additions and 4 deletions
Download patch file Download diff file Expand all files Collapse all files

9
api/server/controllers/agents/client.js
View file

@ -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;

164
packages/api/src/agents/client.ts
View file

@ -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<typeof Tokenizer.getTokenCount>[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;
};
}