mirror of
https://github.com/danny-avila/LibreChat.git
synced 2026-08-04 14:57:42 +00:00
📊 feat: Real-Time Context Window & Token Usage Tracking
This commit is contained in:
Danny Avila
parent
197a1dc4e2
commit
2c7f5c38af
44 changed files with 2271 additions and 50 deletions
|
|
@ -700,7 +700,6 @@ class BaseClient {
|
|||
user,
|
||||
);
|
||||
this.savedMessageIds.add(responseMessage.messageId);
|
||||
delete responseMessage.tokenCount;
|
||||
return responseMessage;
|
||||
}
|
||||
|
||||
|
|
|
|||
50
api/server/controllers/TokenConfigController.js
Normal file
50
api/server/controllers/TokenConfigController.js
Normal file
|
|
@ -0,0 +1,50 @@
|
|||
const { logger } = require('@librechat/data-schemas');
|
||||
const { EModelEndpoint } = require('librechat-data-provider');
|
||||
const { buildTokenConfigMap, getTokenConfigKey, tokenConfigCache } = require('@librechat/api');
|
||||
const { getModelsConfig } = require('~/server/controllers/ModelController');
|
||||
const { getValueKey, getMultiplier, getCacheMultiplier } = require('~/models');
|
||||
|
||||
/**
|
||||
* Returns server-resolved context windows (and pricing when
|
||||
* `interface.contextCost` is enabled) for every configured model.
|
||||
* @param {ServerRequest} req
|
||||
* @param {ServerResponse} res
|
||||
*/
|
||||
async function tokenConfigController(req, res) {
|
||||
try {
|
||||
const appConfig = req.config;
|
||||
const includePricing = appConfig?.interfaceConfig?.contextCost === true;
|
||||
const modelsConfig = await getModelsConfig(req);
|
||||
|
||||
/** @type {Record<string, import('@librechat/api').EndpointTokenConfig | undefined>} */
|
||||
const endpointTokenConfigs = {};
|
||||
const customEndpoints = appConfig?.endpoints?.[EModelEndpoint.custom] ?? [];
|
||||
const cache = tokenConfigCache();
|
||||
for (const endpointConfig of customEndpoints) {
|
||||
const name = endpointConfig?.name;
|
||||
if (!name) {
|
||||
continue;
|
||||
}
|
||||
if (endpointConfig.tokenConfig != null) {
|
||||
endpointTokenConfigs[name] = endpointConfig.tokenConfig;
|
||||
continue;
|
||||
}
|
||||
const tokenKey = getTokenConfigKey(endpointConfig, name, req.user.id);
|
||||
const cached = await cache.get(tokenKey);
|
||||
if (cached) {
|
||||
endpointTokenConfigs[name] = cached;
|
||||
}
|
||||
}
|
||||
|
||||
const tokenConfigMap = buildTokenConfigMap(
|
||||
{ modelsConfig, endpointTokenConfigs, includePricing },
|
||||
{ getValueKey, getMultiplier, getCacheMultiplier },
|
||||
);
|
||||
res.json(tokenConfigMap);
|
||||
} catch (error) {
|
||||
logger.error('[tokenConfigController]', error);
|
||||
res.status(500).json({ error: 'Failed to resolve token config' });
|
||||
}
|
||||
}
|
||||
|
||||
module.exports = tokenConfigController;
|
||||
|
|
@ -0,0 +1,252 @@
|
|||
const { z } = require('zod');
|
||||
const { tool } = require('@langchain/core/tools');
|
||||
const { ChatGenerationChunk } = require('@langchain/core/outputs');
|
||||
const { HumanMessage, AIMessageChunk } = require('@langchain/core/messages');
|
||||
const {
|
||||
Run,
|
||||
Providers,
|
||||
GraphEvents,
|
||||
FakeChatModel,
|
||||
createContentAggregator,
|
||||
} = require('@librechat/agents');
|
||||
const { GenerationJobManager } = require('@librechat/api');
|
||||
const { getDefaultHandlers } = require('~/server/controllers/agents/callbacks');
|
||||
|
||||
jest.mock('nanoid', () => ({
|
||||
nanoid: jest.fn(() => 'mock-nanoid'),
|
||||
}));
|
||||
|
||||
jest.mock('~/server/services/Files/Citations', () => ({
|
||||
processFileCitations: jest.fn(),
|
||||
}));
|
||||
|
||||
jest.mock('~/server/services/Files/Code/process', () => ({
|
||||
processCodeOutput: jest.fn(),
|
||||
runPreviewFinalize: jest.fn(),
|
||||
}));
|
||||
|
||||
jest.mock('~/server/services/Files/process', () => ({
|
||||
saveBase64Image: jest.fn(),
|
||||
}));
|
||||
|
||||
/** Real pipeline guard: published lib versions without the event skip its assertions */
|
||||
const hasContextUsageEvent = GraphEvents.ON_CONTEXT_USAGE != null;
|
||||
|
||||
/**
|
||||
* FakeChatModel that attaches provider-style usage_metadata on a final
|
||||
* empty chunk (the OpenAI streaming pattern), so CHAT_MODEL_END carries
|
||||
* aggregated usage through the real @librechat/agents pipeline.
|
||||
*/
|
||||
class UsageFakeModel extends FakeChatModel {
|
||||
constructor(options, usagePerCall) {
|
||||
super(options);
|
||||
this.usagePerCall = usagePerCall;
|
||||
this.usageCallIndex = 0;
|
||||
}
|
||||
|
||||
async *_streamResponseChunks(messages, options, runManager) {
|
||||
yield* super._streamResponseChunks(messages, options, runManager);
|
||||
const index = Math.min(this.usageCallIndex, this.usagePerCall.length - 1);
|
||||
this.usageCallIndex += 1;
|
||||
yield new ChatGenerationChunk({
|
||||
text: '',
|
||||
message: new AIMessageChunk({ content: '', usage_metadata: this.usagePerCall[index] }),
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
const addTool = tool(async ({ a, b }) => String(a + b), {
|
||||
name: 'add',
|
||||
description: 'Add two numbers',
|
||||
schema: z.object({ a: z.number(), b: z.number() }),
|
||||
});
|
||||
|
||||
const charCounter = (msg) => {
|
||||
const content = msg.content;
|
||||
if (typeof content === 'string') {
|
||||
return content.length + 3;
|
||||
}
|
||||
if (Array.isArray(content)) {
|
||||
let length = 3;
|
||||
for (const part of content) {
|
||||
if (typeof part === 'string') {
|
||||
length += part.length;
|
||||
} else if (typeof part?.text === 'string') {
|
||||
length += part.text.length;
|
||||
}
|
||||
}
|
||||
return length;
|
||||
}
|
||||
return 3;
|
||||
};
|
||||
|
||||
function createMockRes() {
|
||||
const events = [];
|
||||
return {
|
||||
events,
|
||||
headersSent: true,
|
||||
writableEnded: false,
|
||||
write(payload) {
|
||||
for (const line of String(payload).split('\n')) {
|
||||
if (line.startsWith('data: ')) {
|
||||
events.push(JSON.parse(line.slice(6)));
|
||||
}
|
||||
}
|
||||
return true;
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
const FIRST_CALL_USAGE = {
|
||||
input_tokens: 100,
|
||||
output_tokens: 20,
|
||||
total_tokens: 120,
|
||||
};
|
||||
|
||||
const SECOND_CALL_USAGE = {
|
||||
input_tokens: 150,
|
||||
output_tokens: 10,
|
||||
total_tokens: 160,
|
||||
input_token_details: { cache_creation: 30, cache_read: 50 },
|
||||
};
|
||||
|
||||
const MAX_CONTEXT_TOKENS = 8000;
|
||||
|
||||
async function runToolLoop({ res, streamId = null, collectedUsage }) {
|
||||
const { contentParts, aggregateContent } = createContentAggregator();
|
||||
const handlers = getDefaultHandlers({
|
||||
res,
|
||||
aggregateContent,
|
||||
toolEndCallback: () => {},
|
||||
collectedUsage,
|
||||
streamId,
|
||||
});
|
||||
|
||||
const run = await Run.create({
|
||||
runId: 'usage-e2e-response',
|
||||
graphConfig: {
|
||||
type: 'standard',
|
||||
llmConfig: {
|
||||
provider: Providers.OPENAI,
|
||||
model: 'gpt-4o-mini',
|
||||
streaming: true,
|
||||
streamUsage: false,
|
||||
},
|
||||
instructions: 'You are a helpful assistant.',
|
||||
maxContextTokens: MAX_CONTEXT_TOKENS,
|
||||
tools: [addTool],
|
||||
},
|
||||
returnContent: true,
|
||||
customHandlers: handlers,
|
||||
tokenCounter: charCounter,
|
||||
indexTokenCountMap: {},
|
||||
});
|
||||
|
||||
run.Graph.overrideModel = new UsageFakeModel(
|
||||
{
|
||||
responses: ['Let me calculate that.', 'The answer is 4.'],
|
||||
toolCalls: [{ name: 'add', args: { a: 2, b: 2 }, id: 'tc_1', type: 'tool_call' }],
|
||||
},
|
||||
[FIRST_CALL_USAGE, SECOND_CALL_USAGE],
|
||||
);
|
||||
|
||||
await run.processStream(
|
||||
{ messages: [new HumanMessage('What is 2+2?')] },
|
||||
{
|
||||
configurable: { thread_id: 'usage-e2e-thread', user_id: 'user-1' },
|
||||
streamMode: 'values',
|
||||
version: 'v2',
|
||||
},
|
||||
);
|
||||
|
||||
return { run, contentParts };
|
||||
}
|
||||
|
||||
describe('usage events through the real agents pipeline', () => {
|
||||
jest.setTimeout(30000);
|
||||
|
||||
afterAll(async () => {
|
||||
await GenerationJobManager.destroy();
|
||||
});
|
||||
|
||||
test('emits on_token_usage per model call with collectedUsage parity', async () => {
|
||||
const res = createMockRes();
|
||||
const collectedUsage = [];
|
||||
const { contentParts } = await runToolLoop({ res, collectedUsage });
|
||||
|
||||
const usageEvents = res.events.filter((e) => e.event === 'on_token_usage');
|
||||
expect(usageEvents).toHaveLength(2);
|
||||
|
||||
expect(usageEvents[0].data).toMatchObject(FIRST_CALL_USAGE);
|
||||
expect(usageEvents[1].data).toMatchObject(SECOND_CALL_USAGE);
|
||||
expect(usageEvents[0].data.provider).toBe(Providers.OPENAI);
|
||||
expect(usageEvents[0].data.model).toBeTruthy();
|
||||
expect(usageEvents[0].data.usage_type).toBeUndefined();
|
||||
|
||||
expect(collectedUsage).toHaveLength(2);
|
||||
expect(collectedUsage[0]).toMatchObject(FIRST_CALL_USAGE);
|
||||
expect(collectedUsage[1]).toMatchObject(SECOND_CALL_USAGE);
|
||||
|
||||
const text = contentParts
|
||||
.filter((part) => part?.type === 'text')
|
||||
.map((part) => part.text)
|
||||
.join('');
|
||||
expect(text).toContain('The answer is 4.');
|
||||
});
|
||||
|
||||
test('emits a context snapshot before each model call', async () => {
|
||||
if (!hasContextUsageEvent) {
|
||||
console.warn('Skipping: installed @librechat/agents predates ON_CONTEXT_USAGE');
|
||||
return;
|
||||
}
|
||||
const res = createMockRes();
|
||||
const { run } = await runToolLoop({ res, collectedUsage: [] });
|
||||
expect(run).toBeDefined();
|
||||
|
||||
const contextEvents = res.events.filter((e) => e.event === 'on_context_usage');
|
||||
expect(contextEvents).toHaveLength(2);
|
||||
|
||||
for (const event of contextEvents) {
|
||||
const { breakdown, contextBudget, remainingContextTokens, effectiveInstructionTokens } =
|
||||
event.data;
|
||||
expect(breakdown.maxContextTokens).toBe(MAX_CONTEXT_TOKENS);
|
||||
expect(contextBudget).toBeGreaterThan(0);
|
||||
expect(contextBudget).toBeLessThanOrEqual(MAX_CONTEXT_TOKENS);
|
||||
expect(effectiveInstructionTokens).toBeGreaterThan(0);
|
||||
expect(remainingContextTokens).toBeGreaterThan(0);
|
||||
expect(remainingContextTokens).toBeLessThan(contextBudget);
|
||||
}
|
||||
|
||||
/** Tool loop grows the context between calls */
|
||||
expect(contextEvents[1].data.prePruneContextTokens).toBeGreaterThan(
|
||||
contextEvents[0].data.prePruneContextTokens,
|
||||
);
|
||||
|
||||
/** Snapshot precedes the call's usage event */
|
||||
const firstContextIndex = res.events.findIndex((e) => e.event === 'on_context_usage');
|
||||
const firstUsageIndex = res.events.findIndex((e) => e.event === 'on_token_usage');
|
||||
expect(firstContextIndex).toBeGreaterThanOrEqual(0);
|
||||
expect(firstContextIndex).toBeLessThan(firstUsageIndex);
|
||||
});
|
||||
|
||||
test('persists usage and context snapshot for resume via GenerationJobManager', async () => {
|
||||
const streamId = `usage-e2e-stream-${Date.now()}`;
|
||||
await GenerationJobManager.createJob(streamId, 'user-1', 'convo-1');
|
||||
|
||||
const res = createMockRes();
|
||||
await runToolLoop({ res, streamId, collectedUsage: [] });
|
||||
|
||||
const resumeState = await GenerationJobManager.getResumeState(streamId);
|
||||
expect(resumeState).not.toBeNull();
|
||||
|
||||
expect(resumeState.collectedUsage).toHaveLength(2);
|
||||
expect(resumeState.collectedUsage[0]).toMatchObject(FIRST_CALL_USAGE);
|
||||
expect(resumeState.collectedUsage[1]).toMatchObject(SECOND_CALL_USAGE);
|
||||
|
||||
if (hasContextUsageEvent) {
|
||||
expect(resumeState.contextUsage.breakdown.maxContextTokens).toBe(MAX_CONTEXT_TOKENS);
|
||||
/** Latest-wins: the persisted snapshot is the second call's */
|
||||
expect(resumeState.contextUsage.prePruneContextTokens).toBeGreaterThan(0);
|
||||
}
|
||||
});
|
||||
});
|
||||
|
|
@ -1,6 +1,12 @@
|
|||
const { nanoid } = require('nanoid');
|
||||
const { logger } = require('@librechat/data-schemas');
|
||||
const { Tools, StepTypes, FileContext, ErrorTypes } = require('librechat-data-provider');
|
||||
const {
|
||||
Tools,
|
||||
StepTypes,
|
||||
FileContext,
|
||||
ErrorTypes,
|
||||
UsageEvents,
|
||||
} = require('librechat-data-provider');
|
||||
const {
|
||||
GraphEvents,
|
||||
GraphNodeKeys,
|
||||
|
|
@ -41,13 +47,16 @@ class ModelEndHandler {
|
|||
* Optional; when `null`, the handler is a no-op for signatures. Non-Vertex
|
||||
* providers don't emit `additional_kwargs.signatures`, so capture is also
|
||||
* a no-op for them even when the map is provided.
|
||||
* @param {(data: Record<string, unknown>) => Promise<void> | void} [emitUsage] Optional
|
||||
* callback to stream per-call token usage to the client.
|
||||
*/
|
||||
constructor(collectedUsage, collectedThoughtSignatures = null) {
|
||||
constructor(collectedUsage, collectedThoughtSignatures = null, emitUsage = null) {
|
||||
if (!Array.isArray(collectedUsage)) {
|
||||
throw new Error('collectedUsage must be an array');
|
||||
}
|
||||
this.collectedUsage = collectedUsage;
|
||||
this.collectedThoughtSignatures = collectedThoughtSignatures;
|
||||
this.emitUsage = emitUsage;
|
||||
}
|
||||
|
||||
finalize(errorMessage) {
|
||||
|
|
@ -104,6 +113,19 @@ class ModelEndHandler {
|
|||
|
||||
this.collectedUsage.push(taggedUsage);
|
||||
|
||||
if (this.emitUsage) {
|
||||
await this.emitUsage({
|
||||
input_tokens: taggedUsage.input_tokens,
|
||||
output_tokens: taggedUsage.output_tokens,
|
||||
total_tokens: taggedUsage.total_tokens,
|
||||
input_token_details: taggedUsage.input_token_details,
|
||||
model: taggedUsage.model,
|
||||
provider: taggedUsage.provider,
|
||||
usage_type: taggedUsage.usage_type,
|
||||
runId: metadata?.run_id,
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* `additional_kwargs.signatures` is a flat array indexed by response
|
||||
* part position (text + functionCall interleaved). `tool_calls` is
|
||||
|
|
@ -240,7 +262,11 @@ function getDefaultHandlers({
|
|||
);
|
||||
}
|
||||
const handlers = {
|
||||
[GraphEvents.CHAT_MODEL_END]: new ModelEndHandler(collectedUsage, collectedThoughtSignatures),
|
||||
[GraphEvents.CHAT_MODEL_END]: new ModelEndHandler(
|
||||
collectedUsage,
|
||||
collectedThoughtSignatures,
|
||||
(data) => emitEvent(res, streamId, { event: UsageEvents.ON_TOKEN_USAGE, data }),
|
||||
),
|
||||
[GraphEvents.TOOL_END]: new ToolEndHandler(toolEndCallback, logger),
|
||||
[GraphEvents.ON_RUN_STEP]: {
|
||||
/**
|
||||
|
|
@ -425,6 +451,20 @@ function getDefaultHandlers({
|
|||
|
||||
handlers[GraphEvents.ON_AGENT_LOG] = { handle: agentLogHandler };
|
||||
|
||||
/** Guarded: no-op when the installed @librechat/agents predates the event */
|
||||
if (GraphEvents.ON_CONTEXT_USAGE) {
|
||||
handlers[GraphEvents.ON_CONTEXT_USAGE] = {
|
||||
/**
|
||||
* Forward per-model-call context usage snapshots to the client.
|
||||
* @param {string} event - The event name.
|
||||
* @param {StreamEventData} data - The event data.
|
||||
*/
|
||||
handle: async (event, data) => {
|
||||
await emitEvent(res, streamId, { event, data });
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
return handlers;
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -1,9 +1,11 @@
|
|||
const express = require('express');
|
||||
const requireJwtAuth = require('~/server/middleware/requireJwtAuth');
|
||||
const endpointController = require('~/server/controllers/EndpointController');
|
||||
const tokenConfigController = require('~/server/controllers/TokenConfigController');
|
||||
|
||||
const router = express.Router();
|
||||
/** Auth required for role/tenant-scoped endpoint config resolution. */
|
||||
router.get('/', requireJwtAuth, endpointController);
|
||||
router.get('/token-config', requireJwtAuth, tokenConfigController);
|
||||
|
||||
module.exports = router;
|
||||
|
|
|
|||
|
|
@ -5,12 +5,6 @@ import { useRecoilState, useRecoilValue } from 'recoil';
|
|||
import { Constants, isAssistantsEndpoint, isAgentsEndpoint } from 'librechat-data-provider';
|
||||
import type { TConversation } from 'librechat-data-provider';
|
||||
import type { ExtendedFile, FileSetter, ConvoGenerator } from '~/common';
|
||||
import {
|
||||
useChatContext,
|
||||
useChatFormContext,
|
||||
useAddedChatContext,
|
||||
useAssistantsMapContext,
|
||||
} from '~/Providers';
|
||||
import {
|
||||
useTextarea,
|
||||
useAutoSave,
|
||||
|
|
@ -21,6 +15,12 @@ import {
|
|||
useSubmitMessage,
|
||||
useFocusChatEffect,
|
||||
} from '~/hooks';
|
||||
import {
|
||||
useChatContext,
|
||||
useChatFormContext,
|
||||
useAddedChatContext,
|
||||
useAssistantsMapContext,
|
||||
} from '~/Providers';
|
||||
import PendingManualSkillsChips from './PendingManualSkillsChips';
|
||||
import { cn, getModelSpec, removeFocusRings } from '~/utils';
|
||||
import { useGetStartupConfig } from '~/data-provider';
|
||||
|
|
@ -28,11 +28,12 @@ import { mainTextareaId, BadgeItem } from '~/common';
|
|||
import AttachFileChat from './Files/AttachFileChat';
|
||||
import FileFormChat from './Files/FileFormChat';
|
||||
import TextareaHeader from './TextareaHeader';
|
||||
import SkillsCommand from './SkillsCommand';
|
||||
import PromptsCommand from './PromptsCommand';
|
||||
import SkillsCommand from './SkillsCommand';
|
||||
import AudioRecorder from './AudioRecorder';
|
||||
import CollapseChat from './CollapseChat';
|
||||
import StreamAudio from './StreamAudio';
|
||||
import TokenUsage from './TokenUsage';
|
||||
import StopButton from './StopButton';
|
||||
import SendButton from './SendButton';
|
||||
import EditBadges from './EditBadges';
|
||||
|
|
@ -379,6 +380,7 @@ const ChatForm = memo(function ChatForm({
|
|||
}
|
||||
/>
|
||||
<div className="mx-auto flex" />
|
||||
<TokenUsage index={index} conversation={conversation} isSubmitting={isSubmitting} />
|
||||
{SpeechToText && (
|
||||
<AudioRecorder
|
||||
methods={methods}
|
||||
|
|
|
|||
139
client/src/components/Chat/Input/TokenUsage/Breakdown.tsx
Normal file
139
client/src/components/Chat/Input/TokenUsage/Breakdown.tsx
Normal file
|
|
@ -0,0 +1,139 @@
|
|||
import type { TokenUsageView } from '~/hooks/Chat/useTokenUsage';
|
||||
import { formatTokens, formatCost } from '~/utils';
|
||||
import { useLocalize } from '~/hooks';
|
||||
|
||||
interface RowProps {
|
||||
label: string;
|
||||
value: number;
|
||||
max?: number;
|
||||
}
|
||||
|
||||
function Row({ label, value, max }: RowProps) {
|
||||
const percent = max != null && max > 0 ? Math.min((value / max) * 100, 100) : null;
|
||||
return (
|
||||
<div className="flex items-center justify-between gap-4 text-sm">
|
||||
<span className="text-text-secondary">{label}</span>
|
||||
<span className="font-medium text-text-primary">
|
||||
{formatTokens(value)}
|
||||
{percent != null && (
|
||||
<span className="ml-1 text-xs text-text-secondary" aria-hidden="true">
|
||||
({Math.round(percent)}%)
|
||||
</span>
|
||||
)}
|
||||
</span>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
||||
interface BreakdownProps {
|
||||
view: TokenUsageView;
|
||||
showCost: boolean;
|
||||
}
|
||||
|
||||
export default function Breakdown({ view, showCost }: BreakdownProps) {
|
||||
const localize = useLocalize();
|
||||
const { usedTokens, maxTokens, percent, snapshot, snapshotActive, usageTotals } = view;
|
||||
|
||||
const breakdown = snapshotActive ? snapshot?.breakdown : undefined;
|
||||
const instructionTokens =
|
||||
snapshot?.effectiveInstructionTokens ?? breakdown?.instructionTokens ?? 0;
|
||||
const systemTokens =
|
||||
(breakdown?.systemMessageTokens ?? 0) + (breakdown?.dynamicInstructionTokens ?? 0);
|
||||
const messageTokens = Math.max(0, usedTokens - instructionTokens);
|
||||
const freeTokens = maxTokens != null ? Math.max(0, maxTokens - usedTokens) : null;
|
||||
|
||||
return (
|
||||
<div className="w-64 space-y-3" role="region" aria-label={localize('com_ui_context_usage')}>
|
||||
<div className="flex items-center justify-between">
|
||||
<span className="text-sm font-medium text-text-primary">
|
||||
{localize('com_ui_context_window')}
|
||||
</span>
|
||||
<span className="text-xs font-medium text-text-secondary">
|
||||
{maxTokens != null
|
||||
? `${formatTokens(usedTokens)} / ${formatTokens(maxTokens)} (${Math.round(percent)}%)`
|
||||
: formatTokens(usedTokens)}
|
||||
</span>
|
||||
</div>
|
||||
|
||||
<div
|
||||
role="progressbar"
|
||||
aria-valuemin={0}
|
||||
aria-valuemax={100}
|
||||
aria-valuenow={maxTokens != null ? Math.round(percent) : undefined}
|
||||
aria-label={localize('com_ui_context_usage')}
|
||||
className="h-2 w-full overflow-hidden rounded-full bg-surface-secondary"
|
||||
>
|
||||
<div
|
||||
className="h-full rounded-full bg-text-secondary transition-all duration-300"
|
||||
style={{ width: `${Math.min(percent, 100)}%` }}
|
||||
/>
|
||||
</div>
|
||||
|
||||
<div className="space-y-1.5">
|
||||
{breakdown ? (
|
||||
<>
|
||||
<Row
|
||||
label={localize('com_ui_context_messages')}
|
||||
value={messageTokens}
|
||||
max={maxTokens}
|
||||
/>
|
||||
<Row label={localize('com_ui_context_system')} value={systemTokens} max={maxTokens} />
|
||||
<Row
|
||||
label={localize('com_ui_context_tools')}
|
||||
value={breakdown.toolSchemaTokens}
|
||||
max={maxTokens}
|
||||
/>
|
||||
{breakdown.summaryTokens > 0 && (
|
||||
<Row
|
||||
label={localize('com_ui_context_summary')}
|
||||
value={breakdown.summaryTokens}
|
||||
max={maxTokens}
|
||||
/>
|
||||
)}
|
||||
{freeTokens != null && (
|
||||
<Row label={localize('com_ui_context_free')} value={freeTokens} max={maxTokens} />
|
||||
)}
|
||||
</>
|
||||
) : (
|
||||
<>
|
||||
<Row label={localize('com_ui_input')} value={view.branchTotals.input} />
|
||||
<Row
|
||||
label={localize('com_ui_output')}
|
||||
value={view.branchTotals.output + view.liveTokens}
|
||||
/>
|
||||
{maxTokens == null && (
|
||||
<p className="text-xs text-text-secondary">{localize('com_ui_context_unknown')}</p>
|
||||
)}
|
||||
<p className="text-xs italic text-text-secondary">{localize('com_ui_estimated')}</p>
|
||||
</>
|
||||
)}
|
||||
</div>
|
||||
|
||||
{usageTotals.eventCount > 0 && (
|
||||
<>
|
||||
<div className="border-t border-border-light" role="separator" />
|
||||
<div className="space-y-1.5">
|
||||
<Row label={localize('com_ui_input')} value={usageTotals.input} />
|
||||
<Row label={localize('com_ui_output')} value={usageTotals.output} />
|
||||
{usageTotals.cacheRead > 0 && (
|
||||
<Row label={localize('com_ui_cache_read')} value={usageTotals.cacheRead} />
|
||||
)}
|
||||
{usageTotals.cacheWrite > 0 && (
|
||||
<Row label={localize('com_ui_cache_write')} value={usageTotals.cacheWrite} />
|
||||
)}
|
||||
</div>
|
||||
</>
|
||||
)}
|
||||
|
||||
{showCost && view.costUSD != null && (
|
||||
<>
|
||||
<div className="border-t border-border-light" role="separator" />
|
||||
<div className="flex items-center justify-between text-sm">
|
||||
<span className="text-text-secondary">{localize('com_ui_session_cost')}</span>
|
||||
<span className="font-medium text-text-primary">{formatCost(view.costUSD)}</span>
|
||||
</div>
|
||||
</>
|
||||
)}
|
||||
</div>
|
||||
);
|
||||
}
|
||||
62
client/src/components/Chat/Input/TokenUsage/Gauge.tsx
Normal file
62
client/src/components/Chat/Input/TokenUsage/Gauge.tsx
Normal file
|
|
@ -0,0 +1,62 @@
|
|||
import { cn } from '~/utils';
|
||||
|
||||
const SIZE = 28;
|
||||
const STROKE_WIDTH = 3.5;
|
||||
const RADIUS = (SIZE - STROKE_WIDTH) / 2;
|
||||
const CIRCUMFERENCE = 2 * Math.PI * RADIUS;
|
||||
|
||||
interface GaugeProps {
|
||||
/** 0–100, clamped by the caller */
|
||||
percent: number;
|
||||
/** Max context unknown — render an empty track only */
|
||||
indeterminate: boolean;
|
||||
}
|
||||
|
||||
function getStrokeClass(percent: number, indeterminate: boolean): string {
|
||||
if (indeterminate) {
|
||||
return 'stroke-text-secondary';
|
||||
}
|
||||
if (percent > 90) {
|
||||
return 'stroke-red-500';
|
||||
}
|
||||
if (percent > 75) {
|
||||
return 'stroke-yellow-500';
|
||||
}
|
||||
return 'stroke-text-secondary';
|
||||
}
|
||||
|
||||
export default function Gauge({ percent, indeterminate }: GaugeProps) {
|
||||
const offset = CIRCUMFERENCE - (percent / 100) * CIRCUMFERENCE;
|
||||
|
||||
return (
|
||||
<svg
|
||||
width={SIZE}
|
||||
height={SIZE}
|
||||
viewBox={`0 0 ${SIZE} ${SIZE}`}
|
||||
className="-rotate-90"
|
||||
aria-hidden="true"
|
||||
focusable="false"
|
||||
>
|
||||
<circle
|
||||
cx={SIZE / 2}
|
||||
cy={SIZE / 2}
|
||||
r={RADIUS}
|
||||
fill="transparent"
|
||||
strokeWidth={STROKE_WIDTH}
|
||||
className="stroke-border-heavy"
|
||||
strokeDasharray={indeterminate ? '2 4' : undefined}
|
||||
/>
|
||||
<circle
|
||||
cx={SIZE / 2}
|
||||
cy={SIZE / 2}
|
||||
r={RADIUS}
|
||||
fill="transparent"
|
||||
strokeWidth={STROKE_WIDTH}
|
||||
strokeDasharray={CIRCUMFERENCE}
|
||||
strokeDashoffset={indeterminate ? CIRCUMFERENCE : offset}
|
||||
strokeLinecap="round"
|
||||
className={cn('transition-all duration-300', getStrokeClass(percent, indeterminate))}
|
||||
/>
|
||||
</svg>
|
||||
);
|
||||
}
|
||||
83
client/src/components/Chat/Input/TokenUsage/index.tsx
Normal file
83
client/src/components/Chat/Input/TokenUsage/index.tsx
Normal file
|
|
@ -0,0 +1,83 @@
|
|||
import { memo } from 'react';
|
||||
import { HoverCard, HoverCardTrigger, HoverCardContent, HoverCardPortal } from '@librechat/client';
|
||||
import type { TConversation } from 'librechat-data-provider';
|
||||
import useTokenUsage from '~/hooks/Chat/useTokenUsage';
|
||||
import { useGetStartupConfig } from '~/data-provider';
|
||||
import { formatTokens } from '~/utils';
|
||||
import { useLocalize } from '~/hooks';
|
||||
import Breakdown from './Breakdown';
|
||||
import Gauge from './Gauge';
|
||||
|
||||
interface TokenUsageProps {
|
||||
index: number;
|
||||
conversation: TConversation | null;
|
||||
isSubmitting: boolean;
|
||||
}
|
||||
|
||||
function TokenUsageIndicator({
|
||||
index,
|
||||
conversation,
|
||||
isSubmitting,
|
||||
showCost,
|
||||
}: TokenUsageProps & {
|
||||
showCost: boolean;
|
||||
}) {
|
||||
const localize = useLocalize();
|
||||
const view = useTokenUsage({ index, conversation, isSubmitting });
|
||||
|
||||
if (view.usedTokens <= 0 && view.maxTokens == null) {
|
||||
return null;
|
||||
}
|
||||
|
||||
const hasMax = view.maxTokens != null && view.maxTokens > 0;
|
||||
const ariaLabel = hasMax
|
||||
? localize('com_ui_context_usage_label', {
|
||||
0: formatTokens(view.usedTokens),
|
||||
1: formatTokens(view.maxTokens ?? 0),
|
||||
2: String(Math.round(view.percent)),
|
||||
})
|
||||
: localize('com_ui_context_usage_label_unknown', { 0: formatTokens(view.usedTokens) });
|
||||
|
||||
return (
|
||||
<HoverCard openDelay={150} closeDelay={100}>
|
||||
<HoverCardTrigger asChild>
|
||||
<button
|
||||
type="button"
|
||||
data-testid="token-usage"
|
||||
className="flex size-9 items-center justify-center rounded-full p-1 transition-colors hover:bg-surface-hover focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring"
|
||||
aria-label={ariaLabel}
|
||||
aria-haspopup="dialog"
|
||||
>
|
||||
<span
|
||||
role="meter"
|
||||
aria-valuemin={0}
|
||||
aria-valuemax={hasMax ? view.maxTokens : undefined}
|
||||
aria-valuenow={view.usedTokens}
|
||||
aria-label={localize('com_ui_context_usage')}
|
||||
className="flex items-center justify-center"
|
||||
>
|
||||
<Gauge percent={view.percent} indeterminate={!hasMax} />
|
||||
</span>
|
||||
</button>
|
||||
</HoverCardTrigger>
|
||||
<HoverCardPortal>
|
||||
<HoverCardContent side="top" align="end" className="w-auto p-3">
|
||||
<Breakdown view={view} showCost={showCost} />
|
||||
</HoverCardContent>
|
||||
</HoverCardPortal>
|
||||
</HoverCard>
|
||||
);
|
||||
}
|
||||
|
||||
/** Config gate kept outside the indicator so disabled deployments mount nothing */
|
||||
const TokenUsage = memo(function TokenUsage(props: TokenUsageProps) {
|
||||
const { data: startupConfig } = useGetStartupConfig();
|
||||
if (startupConfig?.interface?.contextUsage === false) {
|
||||
return null;
|
||||
}
|
||||
return (
|
||||
<TokenUsageIndicator {...props} showCost={startupConfig?.interface?.contextCost === true} />
|
||||
);
|
||||
});
|
||||
|
||||
export default TokenUsage;
|
||||
|
|
@ -1,6 +1,6 @@
|
|||
import { useRecoilValue } from 'recoil';
|
||||
import { QueryKeys, dataService } from 'librechat-data-provider';
|
||||
import { useQuery } from '@tanstack/react-query';
|
||||
import { QueryKeys, dataService } from 'librechat-data-provider';
|
||||
import type { QueryObserverResult, UseQueryOptions } from '@tanstack/react-query';
|
||||
import type t from 'librechat-data-provider';
|
||||
import store from '~/store';
|
||||
|
|
@ -23,6 +23,20 @@ export const useGetEndpointsQuery = <TData = t.TEndpointsConfig>(
|
|||
);
|
||||
};
|
||||
|
||||
export const useTokenConfigQuery = (
|
||||
config?: UseQueryOptions<t.TTokenConfigMap>,
|
||||
): QueryObserverResult<t.TTokenConfigMap> => {
|
||||
const queriesEnabled = useRecoilValue<boolean>(store.queriesEnabled);
|
||||
return useQuery<t.TTokenConfigMap>([QueryKeys.tokenConfig], () => dataService.getTokenConfig(), {
|
||||
staleTime: Infinity,
|
||||
refetchOnWindowFocus: false,
|
||||
refetchOnReconnect: false,
|
||||
refetchOnMount: false,
|
||||
...config,
|
||||
enabled: (config?.enabled ?? true) === true && queriesEnabled,
|
||||
});
|
||||
};
|
||||
|
||||
/**
|
||||
* Auth-aware query key so unauthenticated (login page) and authenticated
|
||||
* (chat page) configs are cached independently, preventing stale
|
||||
|
|
|
|||
|
|
@ -1,4 +1,6 @@
|
|||
export { default as useChatHelpers } from './useChatHelpers';
|
||||
export { default as useTokenLimits } from './useTokenLimits';
|
||||
export { default as useTokenUsage } from './useTokenUsage';
|
||||
export { default as useAddedResponse } from './useAddedResponse';
|
||||
export { default as useChatFunctions } from './useChatFunctions';
|
||||
export { default as useGetAddedConvo } from './useGetAddedConvo';
|
||||
|
|
|
|||
59
client/src/hooks/Chat/useTokenLimits.ts
Normal file
59
client/src/hooks/Chat/useTokenLimits.ts
Normal file
|
|
@ -0,0 +1,59 @@
|
|||
import { useMemo } from 'react';
|
||||
import { isAgentsEndpoint } from 'librechat-data-provider';
|
||||
import type { TConversation, TModelTokenomics } from 'librechat-data-provider';
|
||||
import { useGetStartupConfig, useTokenConfigQuery, useGetAgentByIdQuery } from '~/data-provider';
|
||||
import { getModelSpec } from '~/utils';
|
||||
|
||||
export interface TokenLimits {
|
||||
/** Statically resolved max context; live snapshots override this at run time */
|
||||
maxContextTokens?: number;
|
||||
rates?: TModelTokenomics;
|
||||
endpoint?: string;
|
||||
model?: string;
|
||||
}
|
||||
|
||||
function toNumber(value: unknown): number | undefined {
|
||||
const num = typeof value === 'string' ? parseFloat(value) : value;
|
||||
return typeof num === 'number' && Number.isFinite(num) && num > 0 ? num : undefined;
|
||||
}
|
||||
|
||||
/**
|
||||
* Mirrors the backend resolution chain (packages/api agents/initialize.ts):
|
||||
* explicit conversation setting → agent params → model spec preset →
|
||||
* server-resolved token config lookup.
|
||||
*/
|
||||
export default function useTokenLimits(conversation: TConversation | null): TokenLimits {
|
||||
const { data: startupConfig } = useGetStartupConfig();
|
||||
const { data: tokenConfig } = useTokenConfigQuery();
|
||||
|
||||
const endpoint = conversation?.endpoint ?? '';
|
||||
const agentId = isAgentsEndpoint(endpoint) ? conversation?.agent_id : null;
|
||||
const { data: agent } = useGetAgentByIdQuery(agentId);
|
||||
|
||||
const spec = conversation?.spec;
|
||||
const model = conversation?.model;
|
||||
const maxContextSetting = conversation?.maxContextTokens;
|
||||
|
||||
return useMemo(() => {
|
||||
const specPreset = getModelSpec({ specName: spec, startupConfig })?.preset;
|
||||
|
||||
let lookupEndpoint = endpoint;
|
||||
let lookupModel = model ?? '';
|
||||
if (agent) {
|
||||
lookupEndpoint = agent.provider ?? endpoint;
|
||||
lookupModel = agent.model ?? lookupModel;
|
||||
} else if (specPreset) {
|
||||
lookupEndpoint = specPreset.endpoint ?? lookupEndpoint;
|
||||
lookupModel = lookupModel || (specPreset.model ?? '');
|
||||
}
|
||||
|
||||
const rates = tokenConfig?.[lookupEndpoint]?.[lookupModel];
|
||||
const maxContextTokens =
|
||||
toNumber(maxContextSetting) ??
|
||||
toNumber(agent?.model_parameters?.maxContextTokens) ??
|
||||
toNumber(specPreset?.maxContextTokens) ??
|
||||
rates?.context;
|
||||
|
||||
return { maxContextTokens, rates, endpoint: lookupEndpoint, model: lookupModel };
|
||||
}, [endpoint, model, spec, maxContextSetting, agent, startupConfig, tokenConfig]);
|
||||
}
|
||||
151
client/src/hooks/Chat/useTokenUsage.ts
Normal file
151
client/src/hooks/Chat/useTokenUsage.ts
Normal file
|
|
@ -0,0 +1,151 @@
|
|||
import { useEffect, useMemo, useRef } from 'react';
|
||||
import { useAtomValue, useSetAtom } from 'jotai';
|
||||
import { useQueryClient } from '@tanstack/react-query';
|
||||
import { Constants, QueryKeys } from 'librechat-data-provider';
|
||||
import type { TMessage, TConversation, TModelTokenomics } from 'librechat-data-provider';
|
||||
import type { ContextSnapshot, UsageTotals } from '~/store/usage';
|
||||
import type { BranchTotals } from '~/utils/tokens';
|
||||
import {
|
||||
liveTokensFamily,
|
||||
usageTotalsFamily,
|
||||
branchTotalsFamily,
|
||||
contextSnapshotFamily,
|
||||
} from '~/store/usage';
|
||||
import { useLatestMessageId } from '~/hooks/Messages/useLatestMessage';
|
||||
import { buildIndex, sumBranch } from '~/utils';
|
||||
import useTokenLimits from './useTokenLimits';
|
||||
|
||||
export interface TokenUsageParams {
|
||||
index: number;
|
||||
conversation: TConversation | null;
|
||||
isSubmitting: boolean;
|
||||
}
|
||||
|
||||
export interface TokenUsageView {
|
||||
usedTokens: number;
|
||||
maxTokens?: number;
|
||||
/** 0–100, clamped; 0 when max is unknown */
|
||||
percent: number;
|
||||
/** True when derived from per-message counts instead of a backend snapshot */
|
||||
isEstimate: boolean;
|
||||
snapshot: ContextSnapshot | null;
|
||||
snapshotActive: boolean;
|
||||
branchTotals: BranchTotals;
|
||||
usageTotals: UsageTotals;
|
||||
liveTokens: number;
|
||||
rates?: TModelTokenomics;
|
||||
/** Session cost from provider-reported usage; undefined until usage events arrive */
|
||||
costUSD?: number;
|
||||
}
|
||||
|
||||
/**
|
||||
* View-model for the context usage indicator. Mount only inside the
|
||||
* indicator so its subscriptions never re-render the chat tree.
|
||||
*/
|
||||
export default function useTokenUsage({
|
||||
index,
|
||||
conversation,
|
||||
isSubmitting,
|
||||
}: TokenUsageParams): TokenUsageView {
|
||||
const queryClient = useQueryClient();
|
||||
const conversationKey = conversation?.conversationId ?? Constants.NEW_CONVO;
|
||||
|
||||
const tailId = useLatestMessageId(index);
|
||||
const snapshot = useAtomValue(contextSnapshotFamily(conversationKey));
|
||||
const usageTotals = useAtomValue(usageTotalsFamily(conversationKey));
|
||||
const branchTotals = useAtomValue(branchTotalsFamily(conversationKey));
|
||||
const liveTokens = useAtomValue(liveTokensFamily(conversationKey));
|
||||
const setBranchTotals = useSetAtom(branchTotalsFamily(conversationKey));
|
||||
const limits = useTokenLimits(conversation);
|
||||
|
||||
const isSubmittingRef = useRef(isSubmitting);
|
||||
isSubmittingRef.current = isSubmitting;
|
||||
const tailIdRef = useRef(tailId);
|
||||
tailIdRef.current = tailId;
|
||||
const anchorId = snapshot?.anchorMessageId ?? null;
|
||||
const anchorIdRef = useRef(anchorId);
|
||||
anchorIdRef.current = anchorId;
|
||||
|
||||
useEffect(() => {
|
||||
/** Cache `updated` events fire on every state transition — rebuild the
|
||||
* O(n) index only when the data snapshot reference actually changed */
|
||||
let lastIndexed: TMessage[] | undefined;
|
||||
const rebuild = (messages?: TMessage[]) => {
|
||||
if (messages === lastIndexed && messages !== undefined) {
|
||||
return;
|
||||
}
|
||||
lastIndexed = messages;
|
||||
buildIndex(conversationKey, messages);
|
||||
setBranchTotals(sumBranch(conversationKey, tailIdRef.current, anchorIdRef.current));
|
||||
};
|
||||
|
||||
rebuild(queryClient.getQueryData<TMessage[]>([QueryKeys.messages, conversationKey]));
|
||||
|
||||
const unsubscribe = queryClient.getQueryCache().subscribe((event) => {
|
||||
if (isSubmittingRef.current || event.type !== 'updated') {
|
||||
return;
|
||||
}
|
||||
const queryKey = event.query.queryKey;
|
||||
if (
|
||||
!Array.isArray(queryKey) ||
|
||||
queryKey[0] !== QueryKeys.messages ||
|
||||
queryKey[1] !== conversationKey
|
||||
) {
|
||||
return;
|
||||
}
|
||||
rebuild(event.query.state.data as TMessage[] | undefined);
|
||||
});
|
||||
return unsubscribe;
|
||||
}, [conversationKey, queryClient, setBranchTotals]);
|
||||
|
||||
useEffect(() => {
|
||||
setBranchTotals(sumBranch(conversationKey, tailId, anchorId));
|
||||
}, [conversationKey, tailId, anchorId, setBranchTotals]);
|
||||
|
||||
return useMemo(() => {
|
||||
const snapshotActive =
|
||||
snapshot != null &&
|
||||
(isSubmitting || snapshot.anchorMessageId == null || branchTotals.containsAnchor);
|
||||
|
||||
if (snapshotActive && snapshot) {
|
||||
const breakdown = snapshot.breakdown;
|
||||
const maxTokens = snapshot.contextBudget ?? breakdown.maxContextTokens;
|
||||
const instructionTokens = snapshot.effectiveInstructionTokens ?? breakdown.instructionTokens;
|
||||
const baseUsed =
|
||||
snapshot.remainingContextTokens != null
|
||||
? maxTokens - snapshot.remainingContextTokens
|
||||
: instructionTokens + breakdown.messageTokens;
|
||||
const usedTokens = Math.max(0, baseUsed) + liveTokens;
|
||||
return {
|
||||
usedTokens,
|
||||
maxTokens,
|
||||
percent: maxTokens > 0 ? Math.min((usedTokens / maxTokens) * 100, 100) : 0,
|
||||
isEstimate: false,
|
||||
snapshot,
|
||||
snapshotActive,
|
||||
branchTotals,
|
||||
usageTotals,
|
||||
liveTokens,
|
||||
rates: limits.rates,
|
||||
costUSD: usageTotals.eventCount > 0 ? usageTotals.costUSD : undefined,
|
||||
};
|
||||
}
|
||||
|
||||
const usedTokens = branchTotals.input + branchTotals.output + liveTokens;
|
||||
const maxTokens = limits.maxContextTokens;
|
||||
return {
|
||||
usedTokens,
|
||||
maxTokens,
|
||||
percent:
|
||||
maxTokens != null && maxTokens > 0 ? Math.min((usedTokens / maxTokens) * 100, 100) : 0,
|
||||
isEstimate: true,
|
||||
snapshot,
|
||||
snapshotActive: false,
|
||||
branchTotals,
|
||||
usageTotals,
|
||||
liveTokens,
|
||||
rates: limits.rates,
|
||||
costUSD: usageTotals.eventCount > 0 ? usageTotals.costUSD : undefined,
|
||||
};
|
||||
}, [snapshot, isSubmitting, branchTotals, usageTotals, liveTokens, limits]);
|
||||
}
|
||||
|
|
@ -3,5 +3,6 @@ export { default as useResumableSSE } from './useResumableSSE';
|
|||
export { default as useAdaptiveSSE } from './useAdaptiveSSE';
|
||||
export { default as useResumeOnLoad } from './useResumeOnLoad';
|
||||
export { default as useStepHandler } from './useStepHandler';
|
||||
export { default as useUsageHandler } from './useUsageHandler';
|
||||
export { default as useContentHandler } from './useContentHandler';
|
||||
export { default as useAttachmentHandler } from './useAttachmentHandler';
|
||||
|
|
|
|||
|
|
@ -10,6 +10,7 @@ import {
|
|||
ErrorTypes,
|
||||
StepEvents,
|
||||
apiBaseUrl,
|
||||
UsageEvents,
|
||||
createPayload,
|
||||
ViolationTypes,
|
||||
removeNullishValues,
|
||||
|
|
@ -29,6 +30,7 @@ import {
|
|||
clearAllDrafts,
|
||||
removeConvoFromAllQueries,
|
||||
upsertConvoInAllQueries,
|
||||
countTrailingOutputChars,
|
||||
markStreamStartFailedMetadata,
|
||||
} from '~/utils';
|
||||
import {
|
||||
|
|
@ -39,6 +41,7 @@ import {
|
|||
} from '~/data-provider';
|
||||
import useEventHandlers, { buildCreatedInitialResponse } from './useEventHandlers';
|
||||
import { useAuthContext } from '~/hooks/AuthContext';
|
||||
import useUsageHandler from './useUsageHandler';
|
||||
import store from '~/store';
|
||||
|
||||
type ChatHelpers = Pick<
|
||||
|
|
@ -473,6 +476,15 @@ export default function useResumableSSE(
|
|||
const balanceQuery = useGetUserBalance({
|
||||
enabled: !!isAuthenticated && startupConfig?.balance?.enabled,
|
||||
});
|
||||
const {
|
||||
contextHandler,
|
||||
usageHandler,
|
||||
tapStream,
|
||||
finalizeUsage,
|
||||
backfillUsage,
|
||||
resetLive,
|
||||
seedLive,
|
||||
} = useUsageHandler();
|
||||
|
||||
/**
|
||||
* Subscribe to stream via SSE library (supports custom headers)
|
||||
|
|
@ -530,6 +542,7 @@ export default function useResumableSSE(
|
|||
}
|
||||
try {
|
||||
finalHandler(data, currentSubmission as EventSubmission);
|
||||
finalizeUsage(data, { ...currentSubmission, userMessage });
|
||||
} catch (error) {
|
||||
console.error('[ResumableSSE] Error in finalHandler:', error);
|
||||
setIsSubmitting(false);
|
||||
|
|
@ -589,7 +602,23 @@ export default function useResumableSSE(
|
|||
return;
|
||||
}
|
||||
|
||||
if (data.event === UsageEvents.ON_CONTEXT_USAGE) {
|
||||
contextHandler(data.data, { ...currentSubmission, userMessage });
|
||||
return;
|
||||
}
|
||||
|
||||
if (data.event === UsageEvents.ON_TOKEN_USAGE) {
|
||||
usageHandler(data.data, { ...currentSubmission, userMessage });
|
||||
return;
|
||||
}
|
||||
|
||||
if (data.event != null) {
|
||||
if (
|
||||
data.event === StepEvents.ON_MESSAGE_DELTA ||
|
||||
data.event === StepEvents.ON_REASONING_DELTA
|
||||
) {
|
||||
tapStream(data.data, { ...currentSubmission, userMessage });
|
||||
}
|
||||
if (!isResume && !createdStreamIdsRef.current.has(currentStreamId)) {
|
||||
if (isOAuthStepEvent(data)) {
|
||||
preCreatedStepEvents.push(data);
|
||||
|
|
@ -619,6 +648,24 @@ export default function useResumableSSE(
|
|||
currentSubmission = resumeSubmission;
|
||||
submissionRef.current = resumeSubmission;
|
||||
userMessage = resumeSubmission.userMessage;
|
||||
/**
|
||||
* The run's collected usage is the source of truth at sync time:
|
||||
* pre-snapshot events are never replayed (only represented via
|
||||
* aggregated content), so totals are rebuilt from the backfill
|
||||
* and token-usage replay events are skipped when it's present.
|
||||
*/
|
||||
const backfilledUsage = data.resumeState?.collectedUsage;
|
||||
const hasUsageBackfill = (backfilledUsage?.length ?? 0) > 0;
|
||||
backfillUsage(backfilledUsage ?? [], resumeSubmission);
|
||||
if (data.resumeState?.contextUsage) {
|
||||
contextHandler(data.resumeState.contextUsage, resumeSubmission);
|
||||
}
|
||||
/** Output streamed before this resume is not re-delivered as
|
||||
* deltas — estimate it from the trailing aggregated content */
|
||||
seedLive(
|
||||
countTrailingOutputChars(data.resumeState?.aggregatedContent),
|
||||
resumeSubmission,
|
||||
);
|
||||
|
||||
if (data.resumeState?.runSteps) {
|
||||
for (const runStep of data.resumeState.runSteps) {
|
||||
|
|
@ -700,7 +747,13 @@ export default function useResumableSSE(
|
|||
`[ResumableSSE] Replaying ${data.resumeState.replayEvents.length} resume events`,
|
||||
);
|
||||
for (const replayEvent of data.resumeState.replayEvents) {
|
||||
if (replayEvent.event != null) {
|
||||
if (replayEvent.event === UsageEvents.ON_CONTEXT_USAGE) {
|
||||
contextHandler(replayEvent.data, resumeSubmission);
|
||||
} else if (replayEvent.event === UsageEvents.ON_TOKEN_USAGE) {
|
||||
if (!hasUsageBackfill) {
|
||||
usageHandler(replayEvent.data, resumeSubmission);
|
||||
}
|
||||
} else if (replayEvent.event != null) {
|
||||
stepHandler(replayEvent, resumeSubmission);
|
||||
}
|
||||
}
|
||||
|
|
@ -711,6 +764,12 @@ export default function useResumableSSE(
|
|||
for (const pendingEvent of data.pendingEvents) {
|
||||
if (pendingEvent.event === 'title') {
|
||||
titleHandler(pendingEvent);
|
||||
} else if (pendingEvent.event === UsageEvents.ON_CONTEXT_USAGE) {
|
||||
contextHandler(pendingEvent.data, resumeSubmission);
|
||||
} else if (pendingEvent.event === UsageEvents.ON_TOKEN_USAGE) {
|
||||
if (!hasUsageBackfill) {
|
||||
usageHandler(pendingEvent.data, resumeSubmission);
|
||||
}
|
||||
} else if (pendingEvent.event != null) {
|
||||
stepHandler(pendingEvent, resumeSubmission);
|
||||
} else if (pendingEvent.type != null) {
|
||||
|
|
@ -821,6 +880,7 @@ export default function useResumableSSE(
|
|||
console.log('[ResumableSSE] Server-sent error event received:', e.data);
|
||||
sse.close();
|
||||
removeActiveJob(currentStreamId);
|
||||
resetLive({ ...currentSubmission, userMessage });
|
||||
if (
|
||||
!createdStreamIdsRef.current.has(currentStreamId) &&
|
||||
optimisticStreamIdsRef.current.has(currentStreamId)
|
||||
|
|
@ -994,6 +1054,13 @@ export default function useResumableSSE(
|
|||
balanceQuery,
|
||||
removeActiveJob,
|
||||
queryClient,
|
||||
contextHandler,
|
||||
usageHandler,
|
||||
tapStream,
|
||||
finalizeUsage,
|
||||
backfillUsage,
|
||||
resetLive,
|
||||
seedLive,
|
||||
],
|
||||
);
|
||||
|
||||
|
|
|
|||
|
|
@ -2,13 +2,20 @@ import { useEffect, useState } from 'react';
|
|||
import { v4 } from 'uuid';
|
||||
import { SSE } from 'sse.js';
|
||||
import { useSetRecoilState } from 'recoil';
|
||||
import { request, createPayload, removeNullishValues } from 'librechat-data-provider';
|
||||
import {
|
||||
request,
|
||||
UsageEvents,
|
||||
StepEvents,
|
||||
createPayload,
|
||||
removeNullishValues,
|
||||
} from 'librechat-data-provider';
|
||||
import type { TMessage, TPayload, TSubmission, EventSubmission } from 'librechat-data-provider';
|
||||
import type { EventHandlerParams } from './useEventHandlers';
|
||||
import type { TResData } from '~/common';
|
||||
import { useGetStartupConfig, useGetUserBalance } from '~/data-provider';
|
||||
import { useAuthContext } from '~/hooks/AuthContext';
|
||||
import useEventHandlers from './useEventHandlers';
|
||||
import useUsageHandler from './useUsageHandler';
|
||||
import { clearAllDrafts } from '~/utils';
|
||||
import store from '~/store';
|
||||
|
||||
|
|
@ -60,6 +67,7 @@ export default function useSSE(
|
|||
const balanceQuery = useGetUserBalance({
|
||||
enabled: !!isAuthenticated && startupConfig?.balance?.enabled,
|
||||
});
|
||||
const { contextHandler, usageHandler, tapStream, finalizeUsage, resetLive } = useUsageHandler();
|
||||
|
||||
useEffect(() => {
|
||||
if (submission == null || Object.keys(submission).length === 0) {
|
||||
|
|
@ -96,6 +104,7 @@ export default function useSSE(
|
|||
clearAllDrafts(submission.conversation?.conversationId);
|
||||
try {
|
||||
finalHandler(data, submission as EventSubmission);
|
||||
finalizeUsage(data, { ...submission, userMessage });
|
||||
} catch (error) {
|
||||
console.error('Error in finalHandler:', error);
|
||||
setIsSubmitting(false);
|
||||
|
|
@ -116,7 +125,17 @@ export default function useSSE(
|
|||
createdHandler(data, { ...submission, userMessage } as EventSubmission);
|
||||
} else if (data.event === 'title') {
|
||||
titleHandler(data);
|
||||
} else if (data.event === UsageEvents.ON_CONTEXT_USAGE) {
|
||||
contextHandler(data.data, { ...submission, userMessage });
|
||||
} else if (data.event === UsageEvents.ON_TOKEN_USAGE) {
|
||||
usageHandler(data.data, { ...submission, userMessage });
|
||||
} else if (data.event != null) {
|
||||
if (
|
||||
data.event === StepEvents.ON_MESSAGE_DELTA ||
|
||||
data.event === StepEvents.ON_REASONING_DELTA
|
||||
) {
|
||||
tapStream(data.data, { ...submission, userMessage });
|
||||
}
|
||||
stepHandler(data, { ...submission, userMessage } as EventSubmission);
|
||||
} else if (data.sync != null) {
|
||||
const runId = v4();
|
||||
|
|
@ -162,6 +181,7 @@ export default function useSSE(
|
|||
}
|
||||
|
||||
setCompleted((prev) => new Set(prev.add(streamKey)));
|
||||
resetLive({ ...submission, userMessage });
|
||||
const latestMessages = getMessages();
|
||||
const conversationId = latestMessages?.[latestMessages.length - 1]?.conversationId;
|
||||
try {
|
||||
|
|
@ -206,6 +226,7 @@ export default function useSSE(
|
|||
|
||||
console.log('error in server stream.');
|
||||
(startupConfig?.balance?.enabled ?? false) && balanceQuery.refetch();
|
||||
resetLive({ ...submission, userMessage });
|
||||
|
||||
let data: TResData | undefined = undefined;
|
||||
try {
|
||||
|
|
|
|||
212
client/src/hooks/SSE/useUsageHandler.ts
Normal file
212
client/src/hooks/SSE/useUsageHandler.ts
Normal file
|
|
@ -0,0 +1,212 @@
|
|||
import { useRef, useMemo } from 'react';
|
||||
import { getDefaultStore } from 'jotai';
|
||||
import { useQueryClient } from '@tanstack/react-query';
|
||||
import { Constants, QueryKeys } from 'librechat-data-provider';
|
||||
import type {
|
||||
TMessage,
|
||||
TConversation,
|
||||
TTokenConfigMap,
|
||||
TTokenUsageEvent,
|
||||
TContextUsageEvent,
|
||||
} from 'librechat-data-provider';
|
||||
import {
|
||||
EMPTY_USAGE_TOTALS,
|
||||
liveTokensFamily,
|
||||
calibrationFamily,
|
||||
usageTotalsFamily,
|
||||
branchTotalsFamily,
|
||||
removeUsageAtoms,
|
||||
contextSnapshotFamily,
|
||||
} from '~/store/usage';
|
||||
import { sumBranch, upsertEntries, migrateIndex, calcUsageCost, estimateTokens } from '~/utils';
|
||||
|
||||
const FLUSH_INTERVAL_MS = 250;
|
||||
|
||||
interface UsageSubmissionLike {
|
||||
userMessage?: Pick<TMessage, 'messageId'> | null;
|
||||
conversation?: Partial<Pick<TConversation, 'conversationId' | 'endpoint' | 'model'>> | null;
|
||||
}
|
||||
|
||||
interface FinalDataLike {
|
||||
requestMessage?: Partial<TMessage> | null;
|
||||
responseMessage?: Partial<TMessage> | null;
|
||||
conversation?: Partial<TConversation> | null;
|
||||
}
|
||||
|
||||
export interface UsageHandlers {
|
||||
contextHandler: (data: TContextUsageEvent, submission: UsageSubmissionLike) => void;
|
||||
usageHandler: (data: TTokenUsageEvent, submission: UsageSubmissionLike) => void;
|
||||
tapStream: (data: { delta?: { content?: unknown } }, submission: UsageSubmissionLike) => void;
|
||||
finalizeUsage: (data: FinalDataLike, submission: UsageSubmissionLike) => void;
|
||||
resetLive: (submission: UsageSubmissionLike) => void;
|
||||
/** Replaces accumulated totals with the run's collected usage on resume */
|
||||
backfillUsage: (entries: TTokenUsageEvent[], submission: UsageSubmissionLike) => void;
|
||||
/** Seeds the live estimate from already-streamed output chars on resume */
|
||||
seedLive: (chars: number, submission: UsageSubmissionLike) => void;
|
||||
}
|
||||
|
||||
function getConvoKey(submission: UsageSubmissionLike): string {
|
||||
return submission.conversation?.conversationId ?? Constants.NEW_CONVO;
|
||||
}
|
||||
|
||||
function countDeltaChars(content: unknown): number {
|
||||
const parts = Array.isArray(content) ? content : [content];
|
||||
let chars = 0;
|
||||
for (const part of parts) {
|
||||
if (part == null || typeof part !== 'object') {
|
||||
continue;
|
||||
}
|
||||
const { text, think } = part as { text?: unknown; think?: unknown };
|
||||
if (typeof text === 'string') {
|
||||
chars += text.length;
|
||||
} else if (typeof think === 'string') {
|
||||
chars += think.length;
|
||||
}
|
||||
}
|
||||
return chars;
|
||||
}
|
||||
|
||||
/**
|
||||
* Imperative writers for the per-conversation token usage atoms, driven by
|
||||
* SSE events. All state lives in refs/atoms — the returned handlers are
|
||||
* stable and never cause re-renders themselves.
|
||||
*/
|
||||
export default function useUsageHandler(): UsageHandlers {
|
||||
const queryClient = useQueryClient();
|
||||
/** Streamed chars since the last snapshot or model end (current call only) */
|
||||
const streamCharsRef = useRef(0);
|
||||
/** Provider-confirmed output tokens since the last snapshot (current run) */
|
||||
const confirmedRef = useRef(0);
|
||||
const lastFlushRef = useRef(0);
|
||||
|
||||
return useMemo<UsageHandlers>(() => {
|
||||
const jotai = getDefaultStore();
|
||||
|
||||
const setLive = (convoKey: string, value: number) => {
|
||||
jotai.set(liveTokensFamily(convoKey), value);
|
||||
};
|
||||
|
||||
const contextHandler: UsageHandlers['contextHandler'] = (data, submission) => {
|
||||
const convoKey = getConvoKey(submission);
|
||||
jotai.set(contextSnapshotFamily(convoKey), {
|
||||
...data,
|
||||
anchorMessageId: submission.userMessage?.messageId ?? null,
|
||||
});
|
||||
if (data.calibrationRatio != null && data.calibrationRatio > 0) {
|
||||
jotai.set(calibrationFamily(convoKey), data.calibrationRatio);
|
||||
}
|
||||
streamCharsRef.current = 0;
|
||||
confirmedRef.current = 0;
|
||||
setLive(convoKey, 0);
|
||||
};
|
||||
|
||||
const foldUsage = (data: TTokenUsageEvent, submission: UsageSubmissionLike) => {
|
||||
const convoKey = getConvoKey(submission);
|
||||
const tokenConfig = queryClient.getQueryData<TTokenConfigMap>([QueryKeys.tokenConfig]);
|
||||
const endpoint = submission.conversation?.endpoint ?? '';
|
||||
const model = data.model ?? submission.conversation?.model ?? '';
|
||||
const rates = tokenConfig?.[endpoint]?.[model];
|
||||
|
||||
const totalsAtom = usageTotalsFamily(convoKey);
|
||||
const prev = jotai.get(totalsAtom);
|
||||
jotai.set(totalsAtom, {
|
||||
input: prev.input + (data.input_tokens ?? 0),
|
||||
output: prev.output + (data.output_tokens ?? 0),
|
||||
cacheWrite: prev.cacheWrite + (data.input_token_details?.cache_creation ?? 0),
|
||||
cacheRead: prev.cacheRead + (data.input_token_details?.cache_read ?? 0),
|
||||
costUSD: prev.costUSD + calcUsageCost(data, rates),
|
||||
eventCount: prev.eventCount + 1,
|
||||
});
|
||||
};
|
||||
|
||||
const usageHandler: UsageHandlers['usageHandler'] = (data, submission) => {
|
||||
foldUsage(data, submission);
|
||||
|
||||
if (data.usage_type === 'summarization') {
|
||||
return;
|
||||
}
|
||||
confirmedRef.current += data.output_tokens ?? 0;
|
||||
streamCharsRef.current = 0;
|
||||
setLive(getConvoKey(submission), confirmedRef.current);
|
||||
};
|
||||
|
||||
const tapStream: UsageHandlers['tapStream'] = (data, submission) => {
|
||||
const chars = countDeltaChars(data?.delta?.content);
|
||||
if (chars <= 0) {
|
||||
return;
|
||||
}
|
||||
streamCharsRef.current += chars;
|
||||
const now = Date.now();
|
||||
if (now - lastFlushRef.current < FLUSH_INTERVAL_MS) {
|
||||
return;
|
||||
}
|
||||
lastFlushRef.current = now;
|
||||
const convoKey = getConvoKey(submission);
|
||||
const ratio = jotai.get(calibrationFamily(convoKey));
|
||||
setLive(convoKey, confirmedRef.current + estimateTokens(streamCharsRef.current, ratio));
|
||||
};
|
||||
|
||||
const resetLive: UsageHandlers['resetLive'] = (submission) => {
|
||||
streamCharsRef.current = 0;
|
||||
confirmedRef.current = 0;
|
||||
setLive(getConvoKey(submission), 0);
|
||||
};
|
||||
|
||||
const backfillUsage: UsageHandlers['backfillUsage'] = (entries, submission) => {
|
||||
jotai.set(usageTotalsFamily(getConvoKey(submission)), EMPTY_USAGE_TOTALS);
|
||||
for (const entry of entries) {
|
||||
foldUsage(entry, submission);
|
||||
}
|
||||
};
|
||||
|
||||
const seedLive: UsageHandlers['seedLive'] = (chars, submission) => {
|
||||
if (chars <= 0) {
|
||||
return;
|
||||
}
|
||||
const convoKey = getConvoKey(submission);
|
||||
streamCharsRef.current = chars;
|
||||
confirmedRef.current = 0;
|
||||
setLive(convoKey, estimateTokens(chars, jotai.get(calibrationFamily(convoKey))));
|
||||
};
|
||||
|
||||
const finalizeUsage: UsageHandlers['finalizeUsage'] = (data, submission) => {
|
||||
const fromKey = getConvoKey(submission);
|
||||
const realId = data.conversation?.conversationId ?? fromKey;
|
||||
|
||||
upsertEntries(fromKey, [data.requestMessage, data.responseMessage]);
|
||||
|
||||
if (realId !== fromKey) {
|
||||
migrateIndex(fromKey, realId);
|
||||
jotai.set(contextSnapshotFamily(realId), jotai.get(contextSnapshotFamily(fromKey)));
|
||||
jotai.set(usageTotalsFamily(realId), jotai.get(usageTotalsFamily(fromKey)));
|
||||
jotai.set(calibrationFamily(realId), jotai.get(calibrationFamily(fromKey)));
|
||||
removeUsageAtoms(fromKey);
|
||||
}
|
||||
|
||||
const responseMeta = data.responseMessage?.contextMeta;
|
||||
if (responseMeta?.calibrationRatio != null && responseMeta.calibrationRatio > 0) {
|
||||
jotai.set(calibrationFamily(realId), responseMeta.calibrationRatio);
|
||||
}
|
||||
|
||||
const tailId = data.responseMessage?.messageId ?? data.requestMessage?.messageId ?? null;
|
||||
if (tailId) {
|
||||
const anchorId = submission.userMessage?.messageId ?? null;
|
||||
jotai.set(branchTotalsFamily(realId), sumBranch(realId, tailId, anchorId));
|
||||
}
|
||||
|
||||
streamCharsRef.current = 0;
|
||||
confirmedRef.current = 0;
|
||||
setLive(realId, 0);
|
||||
};
|
||||
|
||||
return {
|
||||
contextHandler,
|
||||
usageHandler,
|
||||
tapStream,
|
||||
finalizeUsage,
|
||||
resetLive,
|
||||
backfillUsage,
|
||||
seedLive,
|
||||
};
|
||||
}, [queryClient]);
|
||||
}
|
||||
|
|
@ -837,6 +837,8 @@
|
|||
"com_ui_branch_error": "Failed to create branch",
|
||||
"com_ui_branch_message": "Create branch from this response",
|
||||
"com_ui_by_author": "by {{0}}",
|
||||
"com_ui_cache_read": "Cache read",
|
||||
"com_ui_cache_write": "Cache write",
|
||||
"com_ui_callback_url": "Callback URL",
|
||||
"com_ui_cancel": "Cancel",
|
||||
"com_ui_cancelled": "Cancelled",
|
||||
|
|
@ -883,6 +885,16 @@
|
|||
"com_ui_contact_admin_if_issue_persists": "Contact the Admin if the issue persists",
|
||||
"com_ui_context": "Context",
|
||||
"com_ui_context_filter_sort": "Filter and Sort by Context",
|
||||
"com_ui_context_free": "Free space",
|
||||
"com_ui_context_messages": "Messages",
|
||||
"com_ui_context_summary": "Summary",
|
||||
"com_ui_context_system": "System prompt",
|
||||
"com_ui_context_tools": "Tool schemas",
|
||||
"com_ui_context_unknown": "Context size unknown",
|
||||
"com_ui_context_usage": "Context usage",
|
||||
"com_ui_context_usage_label": "Context window: {{0}} of {{1}} tokens used ({{2}}%)",
|
||||
"com_ui_context_usage_label_unknown": "Context usage: {{0}} tokens used",
|
||||
"com_ui_context_window": "Context window",
|
||||
"com_ui_continue": "Continue",
|
||||
"com_ui_continue_oauth": "Continue with OAuth",
|
||||
"com_ui_control_bar": "Control bar",
|
||||
|
|
@ -1029,6 +1041,7 @@
|
|||
"com_ui_error_try_following_prefix": "Please try one of the following",
|
||||
"com_ui_error_unexpected": "Oops! Something Unexpected Occurred",
|
||||
"com_ui_error_updating_preferences": "Error updating preferences",
|
||||
"com_ui_estimated": "Estimated from message history",
|
||||
"com_ui_everyone_permission_level": "Everyone's permission level",
|
||||
"com_ui_examples": "Examples",
|
||||
"com_ui_expand": "Expand",
|
||||
|
|
@ -1508,6 +1521,7 @@
|
|||
"com_ui_select_search_model": "Search model by name",
|
||||
"com_ui_select_search_provider": "Search provider by name",
|
||||
"com_ui_select_search_region": "Search region by name",
|
||||
"com_ui_session_cost": "Session cost",
|
||||
"com_ui_set": "Set",
|
||||
"com_ui_share": "Share",
|
||||
"com_ui_share_create_message": "Your name and any messages you add after sharing stay private.",
|
||||
|
|
|
|||
|
|
@ -1,21 +1,22 @@
|
|||
import * as artifacts from './artifacts';
|
||||
import families from './families';
|
||||
import endpoints from './endpoints';
|
||||
import user from './user';
|
||||
import text from './text';
|
||||
import toast from './toast';
|
||||
import submission from './submission';
|
||||
import isTemporary from './temporary';
|
||||
import endpoints from './endpoints';
|
||||
import families from './families';
|
||||
import settings from './settings';
|
||||
import prompts from './prompts';
|
||||
import search from './search';
|
||||
import preset from './preset';
|
||||
import prompts from './prompts';
|
||||
import lang from './language';
|
||||
import settings from './settings';
|
||||
import toast from './toast';
|
||||
import user from './user';
|
||||
import text from './text';
|
||||
import misc from './misc';
|
||||
import isTemporary from './temporary';
|
||||
export * from './agents';
|
||||
export * from './mcp';
|
||||
export * from './favorites';
|
||||
export * from './subagents';
|
||||
export * from './usage';
|
||||
|
||||
export default {
|
||||
...artifacts,
|
||||
|
|
|
|||
60
client/src/store/usage.ts
Normal file
60
client/src/store/usage.ts
Normal file
|
|
@ -0,0 +1,60 @@
|
|||
import { atom } from 'jotai';
|
||||
import { atomFamily } from 'jotai/utils';
|
||||
import type { TContextUsageEvent } from 'librechat-data-provider';
|
||||
import type { BranchTotals } from '~/utils/tokens';
|
||||
import { EMPTY_BRANCH } from '~/utils/tokens';
|
||||
|
||||
/** Latest backend context snapshot, anchored to the run's user message for staleness checks */
|
||||
export interface ContextSnapshot extends TContextUsageEvent {
|
||||
anchorMessageId: string | null;
|
||||
}
|
||||
|
||||
/** Cumulative provider-reported usage for the conversation's current session */
|
||||
export interface UsageTotals {
|
||||
input: number;
|
||||
output: number;
|
||||
cacheWrite: number;
|
||||
cacheRead: number;
|
||||
costUSD: number;
|
||||
eventCount: number;
|
||||
}
|
||||
|
||||
export const EMPTY_USAGE_TOTALS: UsageTotals = {
|
||||
input: 0,
|
||||
output: 0,
|
||||
cacheWrite: 0,
|
||||
cacheRead: 0,
|
||||
costUSD: 0,
|
||||
eventCount: 0,
|
||||
};
|
||||
|
||||
/**
|
||||
* Split into small per-conversation atoms so high-frequency writes (live
|
||||
* stream estimates) re-render only the subscribers that draw them.
|
||||
*/
|
||||
export const branchTotalsFamily = atomFamily((_conversationId: string) =>
|
||||
atom<BranchTotals>(EMPTY_BRANCH),
|
||||
);
|
||||
|
||||
export const contextSnapshotFamily = atomFamily((_conversationId: string) =>
|
||||
atom<ContextSnapshot | null>(null),
|
||||
);
|
||||
|
||||
export const usageTotalsFamily = atomFamily((_conversationId: string) =>
|
||||
atom<UsageTotals>(EMPTY_USAGE_TOTALS),
|
||||
);
|
||||
|
||||
/** Throttled in-flight output token estimate for the current model call */
|
||||
export const liveTokensFamily = atomFamily((_conversationId: string) => atom<number>(0));
|
||||
|
||||
/** Last known provider-vs-estimate calibration ratio for the conversation */
|
||||
export const calibrationFamily = atomFamily((_conversationId: string) => atom<number>(1));
|
||||
|
||||
/** Jotai atomFamily entries are never GC'd — call on conversation switch/cleanup */
|
||||
export function removeUsageAtoms(conversationId: string): void {
|
||||
branchTotalsFamily.remove(conversationId);
|
||||
contextSnapshotFamily.remove(conversationId);
|
||||
usageTotalsFamily.remove(conversationId);
|
||||
liveTokensFamily.remove(conversationId);
|
||||
calibrationFamily.remove(conversationId);
|
||||
}
|
||||
|
|
@ -21,6 +21,7 @@ export * from './presets';
|
|||
export * from './prompts';
|
||||
export * from './textarea';
|
||||
export * from './messages';
|
||||
export * from './tokens';
|
||||
export * from './redirect';
|
||||
export * from './languages';
|
||||
export * from './conversation';
|
||||
|
|
|
|||
|
|
@ -26,12 +26,29 @@ export type BranchSiblingIndex = {
|
|||
siblingIdx: number;
|
||||
};
|
||||
|
||||
/** Keyed by array identity — cache selectors re-derive the tree from the same
|
||||
* immutable snapshot once per subscriber, so repeat builds are pure waste. */
|
||||
const treeMemo = new WeakMap<TMessage[], ReturnType<typeof buildTree>>();
|
||||
|
||||
const buildTreeMemoized = (messages: TMessage[] | null | undefined) => {
|
||||
if (messages == null) {
|
||||
return null;
|
||||
}
|
||||
const cached = treeMemo.get(messages);
|
||||
if (cached !== undefined) {
|
||||
return cached;
|
||||
}
|
||||
const tree = buildTree({ messages });
|
||||
treeMemo.set(messages, tree);
|
||||
return tree;
|
||||
};
|
||||
|
||||
export const selectActiveBranchTail = (
|
||||
messages: TMessage[] | null | undefined,
|
||||
rootSiblingKey: string | null | undefined,
|
||||
getSiblingIndex: SiblingIndexLookup = () => 0,
|
||||
): TMessage | null => {
|
||||
const messagesTree = buildTree({ messages: messages ?? null });
|
||||
const messagesTree = buildTreeMemoized(messages);
|
||||
if (!messagesTree?.length) {
|
||||
return null;
|
||||
}
|
||||
|
|
@ -61,7 +78,7 @@ export const getMessageBranchSiblingParentIds = (
|
|||
messages: TMessage[] | null | undefined,
|
||||
rootSiblingKey: string | null | undefined,
|
||||
): (string | null)[] => {
|
||||
const messagesTree = buildTree({ messages: messages ?? null });
|
||||
const messagesTree = buildTreeMemoized(messages);
|
||||
if (!messagesTree?.length) {
|
||||
return [];
|
||||
}
|
||||
|
|
|
|||
179
client/src/utils/tokens.spec.ts
Normal file
179
client/src/utils/tokens.spec.ts
Normal file
|
|
@ -0,0 +1,179 @@
|
|||
import { Constants } from 'librechat-data-provider';
|
||||
import type { TMessage } from 'librechat-data-provider';
|
||||
import {
|
||||
buildIndex,
|
||||
upsertEntries,
|
||||
migrateIndex,
|
||||
clearIndex,
|
||||
hasIndex,
|
||||
sumBranch,
|
||||
estimateTokens,
|
||||
calcUsageCost,
|
||||
formatCost,
|
||||
countTrailingOutputChars,
|
||||
EMPTY_BRANCH,
|
||||
} from './tokens';
|
||||
|
||||
const CONVO = 'convo-1';
|
||||
|
||||
function msg(
|
||||
messageId: string,
|
||||
parentMessageId: string | null,
|
||||
isCreatedByUser: boolean,
|
||||
tokenCount?: number,
|
||||
): TMessage {
|
||||
return {
|
||||
messageId,
|
||||
parentMessageId,
|
||||
isCreatedByUser,
|
||||
tokenCount,
|
||||
conversationId: CONVO,
|
||||
text: '',
|
||||
} as TMessage;
|
||||
}
|
||||
|
||||
describe('token index', () => {
|
||||
afterEach(() => {
|
||||
clearIndex(CONVO);
|
||||
clearIndex('convo-2');
|
||||
clearIndex(Constants.NEW_CONVO);
|
||||
});
|
||||
|
||||
it('sums only the active branch via the parent chain', () => {
|
||||
buildIndex(CONVO, [
|
||||
msg('u1', Constants.NO_PARENT, true, 10),
|
||||
msg('a1', 'u1', false, 20),
|
||||
msg('u2', 'a1', true, 30),
|
||||
msg('a2', 'u2', false, 40),
|
||||
/** Sibling branch that must not be counted */
|
||||
msg('a2-alt', 'u2', false, 999),
|
||||
]);
|
||||
|
||||
const totals = sumBranch(CONVO, 'a2');
|
||||
expect(totals.input).toBe(40);
|
||||
expect(totals.output).toBe(60);
|
||||
expect(totals.total).toBe(4);
|
||||
expect(totals.counted).toBe(4);
|
||||
|
||||
const altTotals = sumBranch(CONVO, 'a2-alt');
|
||||
expect(altTotals.output).toBe(1019);
|
||||
});
|
||||
|
||||
it('flags whether the anchor message is on the branch', () => {
|
||||
buildIndex(CONVO, [
|
||||
msg('u1', Constants.NO_PARENT, true, 10),
|
||||
msg('a1', 'u1', false, 20),
|
||||
msg('a1-alt', 'u1', false, 25),
|
||||
]);
|
||||
|
||||
expect(sumBranch(CONVO, 'a1', 'a1').containsAnchor).toBe(true);
|
||||
expect(sumBranch(CONVO, 'a1-alt', 'a1').containsAnchor).toBe(false);
|
||||
});
|
||||
|
||||
it('tracks uncounted messages and tolerates missing parents', () => {
|
||||
buildIndex(CONVO, [msg('u2', 'missing-parent', true, undefined), msg('a2', 'u2', false, 15)]);
|
||||
|
||||
const totals = sumBranch(CONVO, 'a2');
|
||||
expect(totals.total).toBe(2);
|
||||
expect(totals.counted).toBe(1);
|
||||
expect(totals.output).toBe(15);
|
||||
});
|
||||
|
||||
it('returns EMPTY_BRANCH without an index or tail', () => {
|
||||
expect(sumBranch('unknown', 'x')).toBe(EMPTY_BRANCH);
|
||||
buildIndex(CONVO, []);
|
||||
expect(sumBranch(CONVO, null)).toBe(EMPTY_BRANCH);
|
||||
});
|
||||
|
||||
it('upserts incrementally and overwrites by messageId', () => {
|
||||
buildIndex(CONVO, [msg('u1', Constants.NO_PARENT, true, 10)]);
|
||||
upsertEntries(CONVO, [msg('a1', 'u1', false, 20), null, undefined]);
|
||||
expect(sumBranch(CONVO, 'a1').output).toBe(20);
|
||||
|
||||
upsertEntries(CONVO, [msg('a1', 'u1', false, 35)]);
|
||||
expect(sumBranch(CONVO, 'a1').output).toBe(35);
|
||||
});
|
||||
|
||||
it('migrates the index to a new conversation id', () => {
|
||||
buildIndex(Constants.NEW_CONVO, [msg('u1', Constants.NO_PARENT, true, 10)]);
|
||||
migrateIndex(Constants.NEW_CONVO, 'convo-2');
|
||||
|
||||
expect(hasIndex(Constants.NEW_CONVO)).toBe(false);
|
||||
expect(sumBranch('convo-2', 'u1').input).toBe(10);
|
||||
});
|
||||
});
|
||||
|
||||
describe('estimateTokens', () => {
|
||||
it('estimates chars/4 scaled by calibration ratio', () => {
|
||||
expect(estimateTokens(400)).toBe(100);
|
||||
expect(estimateTokens(400, 1.1)).toBe(110);
|
||||
expect(estimateTokens(0)).toBe(0);
|
||||
expect(estimateTokens(100, 0)).toBe(25);
|
||||
});
|
||||
});
|
||||
|
||||
describe('calcUsageCost', () => {
|
||||
const rates = { prompt: 3, completion: 15, cacheWrite: 3.75, cacheRead: 0.3 };
|
||||
|
||||
it('prices additive cache usage (Anthropic pattern)', () => {
|
||||
const cost = calcUsageCost(
|
||||
{
|
||||
input_tokens: 1000,
|
||||
output_tokens: 500,
|
||||
input_token_details: { cache_creation: 2000, cache_read: 10000 },
|
||||
},
|
||||
rates,
|
||||
);
|
||||
expect(cost).toBeCloseTo((1000 * 3 + 2000 * 3.75 + 10000 * 0.3 + 500 * 15) / 1e6);
|
||||
});
|
||||
|
||||
it('prices inclusive cache usage (OpenAI pattern)', () => {
|
||||
const cost = calcUsageCost(
|
||||
{
|
||||
input_tokens: 10000,
|
||||
output_tokens: 500,
|
||||
input_token_details: { cache_read: 4000 },
|
||||
},
|
||||
rates,
|
||||
);
|
||||
expect(cost).toBeCloseTo((6000 * 3 + 4000 * 0.3 + 500 * 15) / 1e6);
|
||||
});
|
||||
|
||||
it('returns 0 without rates', () => {
|
||||
expect(calcUsageCost({ input_tokens: 100, output_tokens: 100 })).toBe(0);
|
||||
expect(calcUsageCost({ input_tokens: 100 }, { context: 1000 })).toBe(0);
|
||||
});
|
||||
});
|
||||
|
||||
describe('formatCost', () => {
|
||||
it('formats across magnitude bands', () => {
|
||||
expect(formatCost(0)).toBe('$0.00');
|
||||
expect(formatCost(0.004)).toBe('<$0.01');
|
||||
expect(formatCost(0.0523)).toBe('$0.0523');
|
||||
expect(formatCost(1.234)).toBe('$1.23');
|
||||
});
|
||||
});
|
||||
|
||||
describe('countTrailingOutputChars', () => {
|
||||
const text = (value: string) => ({ type: 'text', text: value });
|
||||
const think = (value: string) => ({ type: 'think', think: value });
|
||||
const tool = () => ({ type: 'tool_call', tool_call: { id: 'tc' } });
|
||||
|
||||
it('counts the trailing run of text and think parts', () => {
|
||||
expect(countTrailingOutputChars([text('aaaa'), tool(), think('bb'), text('ccc')])).toBe(5);
|
||||
});
|
||||
|
||||
it('skips trailing in-progress tool parts before collecting', () => {
|
||||
expect(countTrailingOutputChars([text('aaaa'), tool(), text('ccc'), tool()])).toBe(3);
|
||||
});
|
||||
|
||||
it('stops at the previous tool boundary', () => {
|
||||
expect(countTrailingOutputChars([text('aaaa'), tool(), text('ccc')])).toBe(3);
|
||||
});
|
||||
|
||||
it('handles empty and non-output content', () => {
|
||||
expect(countTrailingOutputChars(undefined)).toBe(0);
|
||||
expect(countTrailingOutputChars([])).toBe(0);
|
||||
expect(countTrailingOutputChars([tool()])).toBe(0);
|
||||
});
|
||||
});
|
||||
228
client/src/utils/tokens.ts
Normal file
228
client/src/utils/tokens.ts
Normal file
|
|
@ -0,0 +1,228 @@
|
|||
import { Constants } from 'librechat-data-provider';
|
||||
import type { TMessage, TTokenUsageEvent, TModelTokenomics } from 'librechat-data-provider';
|
||||
|
||||
export interface TokenEntry {
|
||||
tokenCount: number;
|
||||
isCreatedByUser: boolean;
|
||||
parentMessageId: string | null;
|
||||
}
|
||||
|
||||
export interface BranchTotals {
|
||||
/** Sum of user-message token counts on the active branch */
|
||||
input: number;
|
||||
/** Sum of assistant-message token counts on the active branch */
|
||||
output: number;
|
||||
/** Messages on the branch with a known tokenCount */
|
||||
counted: number;
|
||||
/** Total messages on the branch */
|
||||
total: number;
|
||||
tailId: string | null;
|
||||
/** Whether the latest run's anchor message is on this branch */
|
||||
containsAnchor: boolean;
|
||||
}
|
||||
|
||||
export const EMPTY_BRANCH: BranchTotals = {
|
||||
input: 0,
|
||||
output: 0,
|
||||
counted: 0,
|
||||
total: 0,
|
||||
tailId: null,
|
||||
containsAnchor: false,
|
||||
};
|
||||
|
||||
/** Module-level token index: conversationId → messageId → entry. Not render state. */
|
||||
const registry = new Map<string, Map<string, TokenEntry>>();
|
||||
|
||||
function toEntry(message: Partial<TMessage>): TokenEntry {
|
||||
return {
|
||||
tokenCount: typeof message.tokenCount === 'number' ? message.tokenCount : 0,
|
||||
isCreatedByUser: message.isCreatedByUser === true,
|
||||
parentMessageId: message.parentMessageId ?? null,
|
||||
};
|
||||
}
|
||||
|
||||
/** Full O(n) rebuild — only on discrete cache replacements (load, refetch, edits) */
|
||||
export function buildIndex(conversationId: string, messages?: TMessage[] | null): void {
|
||||
const index = new Map<string, TokenEntry>();
|
||||
if (messages != null) {
|
||||
for (const message of messages) {
|
||||
if (message?.messageId) {
|
||||
index.set(message.messageId, toEntry(message));
|
||||
}
|
||||
}
|
||||
}
|
||||
registry.set(conversationId, index);
|
||||
}
|
||||
|
||||
/** Incremental upsert — final SSE events touch at most two entries */
|
||||
export function upsertEntries(
|
||||
conversationId: string,
|
||||
messages: (Partial<TMessage> | null | undefined)[],
|
||||
): void {
|
||||
let index = registry.get(conversationId);
|
||||
if (!index) {
|
||||
index = new Map<string, TokenEntry>();
|
||||
registry.set(conversationId, index);
|
||||
}
|
||||
for (const message of messages) {
|
||||
if (message?.messageId) {
|
||||
index.set(message.messageId, toEntry(message));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/** Re-keys the index when `new` resolves to a real conversation id */
|
||||
export function migrateIndex(fromId: string, toId: string): void {
|
||||
if (fromId === toId) {
|
||||
return;
|
||||
}
|
||||
const index = registry.get(fromId);
|
||||
if (!index) {
|
||||
return;
|
||||
}
|
||||
registry.delete(fromId);
|
||||
registry.set(toId, index);
|
||||
}
|
||||
|
||||
export function clearIndex(conversationId: string): void {
|
||||
registry.delete(conversationId);
|
||||
}
|
||||
|
||||
export function hasIndex(conversationId: string): boolean {
|
||||
return registry.has(conversationId);
|
||||
}
|
||||
|
||||
/** O(depth) walk from the branch tail up the parent chain */
|
||||
export function sumBranch(
|
||||
conversationId: string,
|
||||
tailId: string | null | undefined,
|
||||
anchorId?: string | null,
|
||||
): BranchTotals {
|
||||
const index = registry.get(conversationId);
|
||||
if (!index || !tailId) {
|
||||
return EMPTY_BRANCH;
|
||||
}
|
||||
|
||||
const totals = { input: 0, output: 0, counted: 0, total: 0, containsAnchor: false };
|
||||
let currentId: string | null = tailId;
|
||||
let guard = index.size;
|
||||
|
||||
while (currentId && currentId !== Constants.NO_PARENT && guard-- > 0) {
|
||||
const entry: TokenEntry | undefined = index.get(currentId);
|
||||
if (!entry) {
|
||||
break;
|
||||
}
|
||||
totals.total += 1;
|
||||
if (anchorId != null && currentId === anchorId) {
|
||||
totals.containsAnchor = true;
|
||||
}
|
||||
if (entry.tokenCount > 0) {
|
||||
totals.counted += 1;
|
||||
if (entry.isCreatedByUser) {
|
||||
totals.input += entry.tokenCount;
|
||||
} else {
|
||||
totals.output += entry.tokenCount;
|
||||
}
|
||||
}
|
||||
currentId = entry.parentMessageId;
|
||||
}
|
||||
|
||||
return { ...totals, tailId };
|
||||
}
|
||||
|
||||
function getOutputChars(part: unknown): number | null {
|
||||
if (part == null || typeof part !== 'object') {
|
||||
return null;
|
||||
}
|
||||
const { text, think } = part as { text?: unknown; think?: unknown };
|
||||
if (typeof text === 'string') {
|
||||
return text.length;
|
||||
}
|
||||
if (typeof think === 'string') {
|
||||
return think.length;
|
||||
}
|
||||
return null;
|
||||
}
|
||||
|
||||
/**
|
||||
* Chars of the trailing contiguous text/think parts in aggregated content —
|
||||
* the model output produced since the last tool boundary, which the latest
|
||||
* context snapshot does not yet account for. Trailing non-text parts
|
||||
* (in-progress tool calls) are skipped before collecting.
|
||||
*/
|
||||
export function countTrailingOutputChars(content?: unknown[] | null): number {
|
||||
if (!Array.isArray(content) || content.length === 0) {
|
||||
return 0;
|
||||
}
|
||||
let i = content.length - 1;
|
||||
while (i >= 0 && getOutputChars(content[i]) == null) {
|
||||
i--;
|
||||
}
|
||||
let chars = 0;
|
||||
for (; i >= 0; i--) {
|
||||
const partChars = getOutputChars(content[i]);
|
||||
if (partChars == null) {
|
||||
break;
|
||||
}
|
||||
chars += partChars;
|
||||
}
|
||||
return chars;
|
||||
}
|
||||
|
||||
/** Rough live estimate for streaming text, calibrated by the last known provider ratio */
|
||||
export function estimateTokens(charCount: number, calibrationRatio = 1): number {
|
||||
if (charCount <= 0) {
|
||||
return 0;
|
||||
}
|
||||
const ratio = calibrationRatio > 0 ? calibrationRatio : 1;
|
||||
return Math.round((charCount / 4) * ratio);
|
||||
}
|
||||
|
||||
/**
|
||||
* USD cost of one model call. Mirrors the cache-detection heuristic used in
|
||||
* token accounting: cache counts are additive when they exceed base input
|
||||
* (Anthropic), otherwise cache reads are included in input tokens (OpenAI).
|
||||
*/
|
||||
export function calcUsageCost(usage: TTokenUsageEvent, rates?: TModelTokenomics): number {
|
||||
if (!rates || rates.prompt == null || rates.completion == null) {
|
||||
return 0;
|
||||
}
|
||||
const input = usage.input_tokens ?? 0;
|
||||
const output = usage.output_tokens ?? 0;
|
||||
const cacheWrite = usage.input_token_details?.cache_creation ?? 0;
|
||||
const cacheRead = usage.input_token_details?.cache_read ?? 0;
|
||||
const writeRate = rates.cacheWrite ?? rates.prompt;
|
||||
const readRate = rates.cacheRead ?? rates.prompt;
|
||||
|
||||
const cacheIsAdditive = cacheWrite + cacheRead > input;
|
||||
const baseInput = cacheIsAdditive ? input : Math.max(0, input - cacheRead - cacheWrite);
|
||||
|
||||
return (
|
||||
(baseInput * rates.prompt +
|
||||
cacheWrite * writeRate +
|
||||
cacheRead * readRate +
|
||||
output * rates.completion) /
|
||||
1e6
|
||||
);
|
||||
}
|
||||
|
||||
export function formatTokens(count: number): string {
|
||||
const formatted = new Intl.NumberFormat(undefined, {
|
||||
notation: 'compact',
|
||||
maximumFractionDigits: 1,
|
||||
}).format(count);
|
||||
return formatted.replace(/\.0(?=[A-Za-z]|$)/, '');
|
||||
}
|
||||
|
||||
export function formatCost(usd: number): string {
|
||||
if (usd <= 0) {
|
||||
return '$0.00';
|
||||
}
|
||||
if (usd < 0.01) {
|
||||
return '<$0.01';
|
||||
}
|
||||
if (usd < 1) {
|
||||
return `$${usd.toFixed(4)}`;
|
||||
}
|
||||
return `$${usd.toFixed(2)}`;
|
||||
}
|
||||
|
|
@ -10,6 +10,8 @@
|
|||
* from the conversation and the agents' advertised tools.
|
||||
*/
|
||||
const { FakeChatModel } = require('@librechat/agents');
|
||||
const { ChatGenerationChunk } = require('@langchain/core/outputs');
|
||||
const { AIMessageChunk } = require('@langchain/core/messages');
|
||||
|
||||
const MOCK_REPLY = process.env.MOCK_LLM_REPLY || 'E2E mock reply: pong';
|
||||
const CHUNK_DELAY_MS = Number(process.env.MOCK_LLM_CHUNK_DELAY_MS) || 10;
|
||||
|
|
@ -298,9 +300,41 @@ function replyResponses(text) {
|
|||
return null;
|
||||
}
|
||||
|
||||
/**
|
||||
* Attaches synthetic usage_metadata on a final empty chunk (the OpenAI
|
||||
* streaming pattern) so token-usage SSE events flow end to end in mock runs.
|
||||
*/
|
||||
class UsageEmittingFakeChatModel extends FakeChatModel {
|
||||
async *_streamResponseChunks(messages, options, runManager) {
|
||||
let outputChars = 0;
|
||||
for await (const chunk of super._streamResponseChunks(messages, options, runManager)) {
|
||||
outputChars += typeof chunk.text === 'string' ? chunk.text.length : 0;
|
||||
yield chunk;
|
||||
}
|
||||
const inputChars = (messages ?? []).reduce(
|
||||
(sum, message) => sum + getContentText(message?.content).length,
|
||||
0,
|
||||
);
|
||||
const input_tokens = Math.max(1, Math.ceil(inputChars / 4));
|
||||
const output_tokens = Math.max(1, Math.ceil(outputChars / 4));
|
||||
yield new ChatGenerationChunk({
|
||||
text: '',
|
||||
message: new AIMessageChunk({
|
||||
content: '',
|
||||
usage_metadata: { input_tokens, output_tokens, total_tokens: input_tokens + output_tokens },
|
||||
}),
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
function overrideModel({ graph, responses, sleep, toolCalls, thrownError }) {
|
||||
if (!thrownError) {
|
||||
graph.overrideTestModel(responses, sleep ?? CHUNK_DELAY_MS, toolCalls);
|
||||
graph.overrideModel = new UsageEmittingFakeChatModel({
|
||||
responses,
|
||||
sleep: sleep ?? CHUNK_DELAY_MS,
|
||||
emitCustomEvent: true,
|
||||
toolCalls,
|
||||
});
|
||||
return;
|
||||
}
|
||||
|
||||
|
|
|
|||
50
e2e/specs/mock/usage.spec.ts
Normal file
50
e2e/specs/mock/usage.spec.ts
Normal file
|
|
@ -0,0 +1,50 @@
|
|||
import { expect, test } from '@playwright/test';
|
||||
import type { Page } from '@playwright/test';
|
||||
import {
|
||||
MOCK_ENDPOINTS,
|
||||
NEW_CHAT_PATH,
|
||||
mockReply,
|
||||
sendMessage,
|
||||
selectMockEndpoint,
|
||||
} from './helpers';
|
||||
|
||||
const gauge = (page: Page) => page.getByTestId('token-usage');
|
||||
const gaugeMeter = (page: Page) => gauge(page).getByRole('meter');
|
||||
|
||||
async function expectGaugeAboveZero(page: Page) {
|
||||
await expect(gauge(page)).toBeVisible({ timeout: 20000 });
|
||||
await expect(gaugeMeter(page)).toHaveAttribute('aria-valuenow', /[1-9]/, { timeout: 20000 });
|
||||
}
|
||||
|
||||
test.describe('context usage gauge', () => {
|
||||
test('tracks usage from live SSE events and survives reload', async ({ page }) => {
|
||||
test.setTimeout(120000);
|
||||
await page.goto(NEW_CHAT_PATH, { timeout: 10000 });
|
||||
|
||||
// REQUIRED so the message streams without a real key.
|
||||
await selectMockEndpoint(page, MOCK_ENDPOINTS[0]);
|
||||
|
||||
const response = await sendMessage(page, 'hello');
|
||||
expect(response.ok()).toBeTruthy();
|
||||
await expect(mockReply(page)).toBeVisible({ timeout: 20000 });
|
||||
await expect(page).toHaveURL(/\/c\/(?!new)/, { timeout: 15000 });
|
||||
|
||||
/** Live path: the agents pipeline's context snapshot + usage events fill the gauge */
|
||||
await expectGaugeAboveZero(page);
|
||||
|
||||
/** Breakdown popover: context section always, usage section from on_token_usage events */
|
||||
await gauge(page).hover();
|
||||
const popover = page.getByRole('region', { name: 'Context usage' });
|
||||
await expect(popover).toBeVisible({ timeout: 10000 });
|
||||
await expect(popover.getByText('Context window')).toBeVisible();
|
||||
await expect(popover.getByText('Input', { exact: true })).toBeVisible();
|
||||
await expect(popover.getByText('Output', { exact: true })).toBeVisible();
|
||||
await page.keyboard.press('Escape');
|
||||
|
||||
/** Fallback path: after reload there is no snapshot — the gauge rebuilds
|
||||
* from per-message tokenCount history returned by the messages query */
|
||||
await page.reload({ timeout: 15000 });
|
||||
await expect(mockReply(page)).toBeVisible({ timeout: 20000 });
|
||||
await expectGaugeAboveZero(page);
|
||||
});
|
||||
});
|
||||
|
|
@ -16,6 +16,26 @@ import { tokenConfigCache } from '~/cache';
|
|||
|
||||
const { PROXY } = process.env;
|
||||
|
||||
/**
|
||||
* Cache key for an endpoint's fetched token config. User-scoped when the
|
||||
* model fetch can resolve per-user: user-provided key/URL, or header
|
||||
* templates forwarded against an admin-trusted base URL — making the
|
||||
* response, and therefore the derived token config, user-specific.
|
||||
*/
|
||||
export function getTokenConfigKey(
|
||||
endpointConfig: Partial<TEndpoint>,
|
||||
endpoint: string,
|
||||
userId: string,
|
||||
): string {
|
||||
const hasTokenConfig = (endpointConfig as Record<string, unknown>).tokenConfig != null;
|
||||
const userProvidesKey = isUserProvided(extractEnvVariable(endpointConfig.apiKey ?? ''));
|
||||
const userProvidesURL = isUserProvided(extractEnvVariable(endpointConfig.baseURL ?? ''));
|
||||
const willForwardUserScopedHeaders = !!endpointConfig?.headers && !userProvidesURL;
|
||||
return !hasTokenConfig && (userProvidesKey || userProvidesURL || willForwardUserScopedHeaders)
|
||||
? `${endpoint}:${userId}`
|
||||
: endpoint;
|
||||
}
|
||||
|
||||
/**
|
||||
* Builds custom options from endpoint configuration
|
||||
*/
|
||||
|
|
@ -138,17 +158,7 @@ export async function initializeCustom({
|
|||
const cache = tokenConfigCache();
|
||||
/** tokenConfig is an optional extended property on custom endpoints */
|
||||
const hasTokenConfig = (endpointConfig as Record<string, unknown>).tokenConfig != null;
|
||||
// When `endpointConfig.headers` will be forwarded to the model fetch (i.e.
|
||||
// base URL is admin-trusted, so the security guard below leaves them in
|
||||
// place), header templates may resolve against the current user — making
|
||||
// the response, and therefore the derived token config, user-specific.
|
||||
// User-scope the token-config cache key in that case so a cached entry
|
||||
// for one user can't be served to another.
|
||||
const willForwardUserScopedHeaders = !!endpointConfig?.headers && !userProvidesURL;
|
||||
const tokenKey =
|
||||
!hasTokenConfig && (userProvidesKey || userProvidesURL || willForwardUserScopedHeaders)
|
||||
? `${endpoint}:${userId}`
|
||||
: endpoint;
|
||||
const tokenKey = getTokenConfigKey(endpointConfig, endpoint, userId);
|
||||
|
||||
const cachedConfig =
|
||||
!hasTokenConfig &&
|
||||
|
|
|
|||
|
|
@ -5,3 +5,4 @@ export * from './custom';
|
|||
export * from './google';
|
||||
export * from './models';
|
||||
export * from './openai';
|
||||
export * from './pricing';
|
||||
|
|
|
|||
89
packages/api/src/endpoints/pricing.spec.ts
Normal file
89
packages/api/src/endpoints/pricing.spec.ts
Normal file
|
|
@ -0,0 +1,89 @@
|
|||
import mongoose from 'mongoose';
|
||||
import { EModelEndpoint } from 'librechat-data-provider';
|
||||
import { createTxMethods } from '@librechat/data-schemas';
|
||||
import type { EndpointTokenConfig } from '~/types';
|
||||
import { matchModelName, findMatchingPattern } from '~/utils';
|
||||
import { buildTokenConfigMap } from './pricing';
|
||||
|
||||
const { getValueKey, getMultiplier, getCacheMultiplier } = createTxMethods(mongoose, {
|
||||
matchModelName,
|
||||
findMatchingPattern,
|
||||
});
|
||||
|
||||
const deps = { getValueKey, getMultiplier, getCacheMultiplier };
|
||||
|
||||
describe('buildTokenConfigMap', () => {
|
||||
it('resolves context windows without pricing by default', () => {
|
||||
const map = buildTokenConfigMap(
|
||||
{
|
||||
modelsConfig: {
|
||||
[EModelEndpoint.openAI]: ['gpt-4o'],
|
||||
[EModelEndpoint.anthropic]: ['claude-3-5-sonnet-20241022'],
|
||||
},
|
||||
},
|
||||
deps,
|
||||
);
|
||||
|
||||
expect(map[EModelEndpoint.openAI]['gpt-4o'].context).toBeGreaterThan(100000);
|
||||
expect(map[EModelEndpoint.anthropic]['claude-3-5-sonnet-20241022'].context).toBe(200000);
|
||||
expect(map[EModelEndpoint.openAI]['gpt-4o'].prompt).toBeUndefined();
|
||||
expect(map[EModelEndpoint.openAI]['gpt-4o'].completion).toBeUndefined();
|
||||
});
|
||||
|
||||
it('includes pattern-matched pricing when enabled', () => {
|
||||
const map = buildTokenConfigMap(
|
||||
{
|
||||
modelsConfig: {
|
||||
[EModelEndpoint.anthropic]: ['claude-3-5-sonnet-20241022'],
|
||||
},
|
||||
includePricing: true,
|
||||
},
|
||||
deps,
|
||||
);
|
||||
|
||||
const tokenomics = map[EModelEndpoint.anthropic]['claude-3-5-sonnet-20241022'];
|
||||
expect(tokenomics.prompt).toBe(3);
|
||||
expect(tokenomics.completion).toBe(15);
|
||||
expect(tokenomics.cacheWrite).toBe(3.75);
|
||||
expect(tokenomics.cacheRead).toBe(0.3);
|
||||
});
|
||||
|
||||
it('treats unknown endpoints as custom for context lookups', () => {
|
||||
const map = buildTokenConfigMap(
|
||||
{
|
||||
modelsConfig: { MyProxy: ['gpt-4o-mini'] },
|
||||
},
|
||||
deps,
|
||||
);
|
||||
|
||||
expect(map.MyProxy['gpt-4o-mini'].context).toBeGreaterThan(0);
|
||||
});
|
||||
|
||||
it('prefers endpoint token config overrides for context and rates', () => {
|
||||
const override: EndpointTokenConfig = {
|
||||
'custom-model': { prompt: 1.5, completion: 4.5, context: 32000 },
|
||||
};
|
||||
const map = buildTokenConfigMap(
|
||||
{
|
||||
modelsConfig: { MyProxy: ['custom-model', 'gpt-4o-mini'] },
|
||||
endpointTokenConfigs: { MyProxy: override },
|
||||
includePricing: true,
|
||||
},
|
||||
deps,
|
||||
);
|
||||
|
||||
expect(map.MyProxy['custom-model']).toEqual({
|
||||
context: 32000,
|
||||
prompt: 1.5,
|
||||
completion: 4.5,
|
||||
});
|
||||
/** Models absent from the override fall back to static tables */
|
||||
expect(map.MyProxy['gpt-4o-mini'].context).toBeGreaterThan(0);
|
||||
expect(map.MyProxy['gpt-4o-mini'].prompt).toBeGreaterThan(0);
|
||||
});
|
||||
|
||||
it('skips endpoints with empty model lists', () => {
|
||||
const map = buildTokenConfigMap({ modelsConfig: { [EModelEndpoint.agents]: [] } }, deps);
|
||||
expect(map[EModelEndpoint.agents]).toBeUndefined();
|
||||
});
|
||||
});
|
||||
98
packages/api/src/endpoints/pricing.ts
Normal file
98
packages/api/src/endpoints/pricing.ts
Normal file
|
|
@ -0,0 +1,98 @@
|
|||
import { EModelEndpoint } from 'librechat-data-provider';
|
||||
import type { TModelsConfig, TTokenConfigMap, TModelTokenomics } from 'librechat-data-provider';
|
||||
import type { TxMethods } from '@librechat/data-schemas';
|
||||
import type { EndpointTokenConfig } from '~/types';
|
||||
import { getModelMaxTokens, maxTokensMap } from '~/utils';
|
||||
|
||||
export interface TokenomicsDeps {
|
||||
getValueKey: TxMethods['getValueKey'];
|
||||
getMultiplier: TxMethods['getMultiplier'];
|
||||
getCacheMultiplier: TxMethods['getCacheMultiplier'];
|
||||
}
|
||||
|
||||
export interface TokenConfigParams {
|
||||
/** endpoint → model list, from the resolved models config */
|
||||
modelsConfig: TModelsConfig;
|
||||
/** Per-endpoint overrides: fetched (e.g. OpenRouter) or yaml `tokenConfig` */
|
||||
endpointTokenConfigs?: Record<string, EndpointTokenConfig | undefined>;
|
||||
/** Include USD-per-1M rates; gated by `interface.contextCost` */
|
||||
includePricing?: boolean;
|
||||
}
|
||||
|
||||
/**
|
||||
* Resolves context windows (and optionally pricing) for every configured
|
||||
* model, server-side, so the client never reimplements pattern matching.
|
||||
*/
|
||||
export function buildTokenConfigMap(
|
||||
params: TokenConfigParams,
|
||||
deps: TokenomicsDeps,
|
||||
): TTokenConfigMap {
|
||||
const { modelsConfig, endpointTokenConfigs, includePricing = false } = params;
|
||||
const map: TTokenConfigMap = {};
|
||||
|
||||
for (const [endpoint, models] of Object.entries(modelsConfig)) {
|
||||
if (!Array.isArray(models) || models.length === 0) {
|
||||
continue;
|
||||
}
|
||||
const override = endpointTokenConfigs?.[endpoint];
|
||||
const endpointKey = (
|
||||
maxTokensMap[endpoint] != null ? endpoint : EModelEndpoint.custom
|
||||
) as EModelEndpoint;
|
||||
|
||||
const entry: Record<string, TModelTokenomics> = {};
|
||||
for (const model of models) {
|
||||
const tokenomics: TModelTokenomics = {};
|
||||
const context = override
|
||||
? (getModelMaxTokens(model, endpointKey, override) ?? getModelMaxTokens(model, endpointKey))
|
||||
: getModelMaxTokens(model, endpointKey);
|
||||
if (context != null) {
|
||||
tokenomics.context = context;
|
||||
}
|
||||
|
||||
if (includePricing) {
|
||||
const overrideRates = override?.[model];
|
||||
if (overrideRates?.prompt != null || overrideRates?.completion != null) {
|
||||
tokenomics.prompt = overrideRates.prompt;
|
||||
tokenomics.completion = overrideRates.completion;
|
||||
} else {
|
||||
const valueKey = deps.getValueKey(model, endpoint);
|
||||
tokenomics.prompt = deps.getMultiplier({
|
||||
valueKey,
|
||||
model,
|
||||
endpoint,
|
||||
tokenType: 'prompt',
|
||||
});
|
||||
tokenomics.completion = deps.getMultiplier({
|
||||
valueKey,
|
||||
model,
|
||||
endpoint,
|
||||
tokenType: 'completion',
|
||||
});
|
||||
const cacheWrite = deps.getCacheMultiplier({
|
||||
valueKey,
|
||||
model,
|
||||
endpoint,
|
||||
cacheType: 'write',
|
||||
});
|
||||
const cacheRead = deps.getCacheMultiplier({
|
||||
valueKey,
|
||||
model,
|
||||
endpoint,
|
||||
cacheType: 'read',
|
||||
});
|
||||
if (cacheWrite != null) {
|
||||
tokenomics.cacheWrite = cacheWrite;
|
||||
}
|
||||
if (cacheRead != null) {
|
||||
tokenomics.cacheRead = cacheRead;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
entry[model] = tokenomics;
|
||||
}
|
||||
map[endpoint] = entry;
|
||||
}
|
||||
|
||||
return map;
|
||||
}
|
||||
|
|
@ -1,5 +1,5 @@
|
|||
import { Constants, parseTextParts } from 'librechat-data-provider';
|
||||
import { logger, getTenantId, SYSTEM_TENANT_ID } from '@librechat/data-schemas';
|
||||
import { Constants, UsageEvents, parseTextParts } from 'librechat-data-provider';
|
||||
import type { Agents, TMessageContentParts } from 'librechat-data-provider';
|
||||
import type { StandardGraph } from '@librechat/agents';
|
||||
import type {
|
||||
|
|
@ -177,6 +177,9 @@ class GenerationJobManagerClass {
|
|||
/** Serializes replay-event read/modify/write updates per stream. */
|
||||
private replayEventWriteQueues = new Map<string, Promise<void>>();
|
||||
|
||||
/** Serializes token-usage read/modify/write updates per stream. */
|
||||
private tokenUsageWriteQueues = new Map<string, Promise<void>>();
|
||||
|
||||
private cleanupInterval: NodeJS.Timeout | null = null;
|
||||
|
||||
/** Whether we're using Redis stores */
|
||||
|
|
@ -649,6 +652,7 @@ class GenerationJobManagerClass {
|
|||
this.jobStore.clearContentState(streamId);
|
||||
this.runStepBuffers?.delete(streamId);
|
||||
this.replayEventWriteQueues.delete(streamId);
|
||||
this.tokenUsageWriteQueues.delete(streamId);
|
||||
|
||||
// For error jobs, DON'T delete immediately - keep around so late-connecting
|
||||
// clients can receive the error. This handles the race condition where error
|
||||
|
|
@ -794,6 +798,7 @@ class GenerationJobManagerClass {
|
|||
this.jobStore.clearContentState(streamId);
|
||||
this.runStepBuffers?.delete(streamId);
|
||||
this.replayEventWriteQueues.delete(streamId);
|
||||
this.tokenUsageWriteQueues.delete(streamId);
|
||||
|
||||
// Immediate cleanup if configured (default: true)
|
||||
if (this._cleanupOnComplete) {
|
||||
|
|
@ -1053,6 +1058,8 @@ class GenerationJobManagerClass {
|
|||
await this.trackUserMessage(streamId, event);
|
||||
await this.trackTitleEvent(streamId, event);
|
||||
await this.trackReplayEvent(streamId, event);
|
||||
await this.trackContextUsage(streamId, event);
|
||||
await this.trackTokenUsage(streamId, event);
|
||||
|
||||
// For Redis mode, persist chunk for later reconstruction (fire-and-forget for resumability)
|
||||
if (this._isRedis) {
|
||||
|
|
@ -1150,6 +1157,48 @@ class GenerationJobManagerClass {
|
|||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Persist the latest context usage snapshot (one per model call) so a
|
||||
* resuming client can restore the context gauge without waiting for the
|
||||
* next model call.
|
||||
*/
|
||||
private async trackContextUsage(streamId: string, event: t.ServerSentEvent): Promise<void> {
|
||||
if (!('event' in event) || event.event !== UsageEvents.ON_CONTEXT_USAGE) {
|
||||
return;
|
||||
}
|
||||
|
||||
await this.jobStore.updateJob(streamId, {
|
||||
contextUsage: JSON.stringify((event as { data?: unknown }).data ?? null),
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Chains a read/modify/write job update onto the stream's queue so
|
||||
* concurrent writers can't clobber each other's merged state.
|
||||
*/
|
||||
private async queueJobWrite(
|
||||
queues: Map<string, Promise<void>>,
|
||||
streamId: string,
|
||||
write: () => Promise<void>,
|
||||
): Promise<void> {
|
||||
const previousWrite = queues.get(streamId) ?? Promise.resolve();
|
||||
const nextWrite = previousWrite
|
||||
.catch(() => {
|
||||
// Keep the queue moving even if a prior metadata write failed.
|
||||
})
|
||||
.then(write);
|
||||
|
||||
queues.set(streamId, nextWrite);
|
||||
|
||||
try {
|
||||
await nextWrite;
|
||||
} finally {
|
||||
if (queues.get(streamId) === nextWrite) {
|
||||
queues.delete(streamId);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Persist replay-only stream events that are needed to reconstruct active
|
||||
* UI state on resume but are not represented by aggregated message content.
|
||||
|
|
@ -1159,22 +1208,45 @@ class GenerationJobManagerClass {
|
|||
return;
|
||||
}
|
||||
|
||||
const previousWrite = this.replayEventWriteQueues.get(streamId) ?? Promise.resolve();
|
||||
const nextWrite = previousWrite
|
||||
.catch(() => {
|
||||
// Keep the queue moving even if a prior replay metadata write failed.
|
||||
})
|
||||
.then(() => this.persistReplayEvent(streamId, event));
|
||||
await this.queueJobWrite(this.replayEventWriteQueues, streamId, () =>
|
||||
this.persistReplayEvent(streamId, event),
|
||||
);
|
||||
}
|
||||
|
||||
this.replayEventWriteQueues.set(streamId, nextWrite);
|
||||
/**
|
||||
* Persist per-model-call token usage so resuming clients can rebuild
|
||||
* usage totals on any replica (the live collectedUsage array only exists
|
||||
* on the generating instance).
|
||||
*/
|
||||
private async trackTokenUsage(streamId: string, event: t.ServerSentEvent): Promise<void> {
|
||||
if (!('event' in event) || event.event !== UsageEvents.ON_TOKEN_USAGE) {
|
||||
return;
|
||||
}
|
||||
|
||||
try {
|
||||
await nextWrite;
|
||||
} finally {
|
||||
if (this.replayEventWriteQueues.get(streamId) === nextWrite) {
|
||||
this.replayEventWriteQueues.delete(streamId);
|
||||
await this.queueJobWrite(this.tokenUsageWriteQueues, streamId, () =>
|
||||
this.persistTokenUsage(streamId, event as { data?: unknown }),
|
||||
);
|
||||
}
|
||||
|
||||
private async persistTokenUsage(streamId: string, event: { data?: unknown }): Promise<void> {
|
||||
const jobData = await this.jobStore.getJob(streamId);
|
||||
if (!jobData || event.data == null) {
|
||||
return;
|
||||
}
|
||||
|
||||
let tokenUsage: unknown[] = [];
|
||||
if (jobData.tokenUsage) {
|
||||
try {
|
||||
tokenUsage = JSON.parse(jobData.tokenUsage) as unknown[];
|
||||
} catch {
|
||||
tokenUsage = [];
|
||||
}
|
||||
}
|
||||
tokenUsage.push(event.data);
|
||||
|
||||
await this.jobStore.updateJob(streamId, {
|
||||
tokenUsage: JSON.stringify(tokenUsage),
|
||||
});
|
||||
}
|
||||
|
||||
private async persistReplayEvent(streamId: string, event: t.ServerSentEvent): Promise<void> {
|
||||
|
|
@ -1342,10 +1414,32 @@ class GenerationJobManagerClass {
|
|||
}
|
||||
}
|
||||
|
||||
let contextUsage: t.ResumeState['contextUsage'];
|
||||
if (jobData.contextUsage) {
|
||||
try {
|
||||
contextUsage = JSON.parse(jobData.contextUsage) as t.ResumeState['contextUsage'];
|
||||
} catch {
|
||||
// Ignore malformed persisted context usage.
|
||||
}
|
||||
}
|
||||
|
||||
/** Persisted per model call by trackTokenUsage — unlike the live
|
||||
* collectedUsage reference, this survives cross-replica resumes. */
|
||||
let collectedUsage: t.ResumeState['collectedUsage'];
|
||||
if (jobData.tokenUsage) {
|
||||
try {
|
||||
const parsed = JSON.parse(jobData.tokenUsage) as t.ResumeState['collectedUsage'];
|
||||
collectedUsage = parsed && parsed.length > 0 ? parsed : undefined;
|
||||
} catch {
|
||||
// Ignore malformed persisted token usage.
|
||||
}
|
||||
}
|
||||
|
||||
logger.debug(`[GenerationJobManager] getResumeState:`, {
|
||||
streamId,
|
||||
runStepsLength: runSteps.length,
|
||||
aggregatedContentLength: aggregatedContent.length,
|
||||
collectedUsageLength: collectedUsage?.length ?? 0,
|
||||
});
|
||||
|
||||
return {
|
||||
|
|
@ -1359,6 +1453,8 @@ class GenerationJobManagerClass {
|
|||
model: jobData.model,
|
||||
titleEvent,
|
||||
replayEvents,
|
||||
collectedUsage,
|
||||
contextUsage,
|
||||
};
|
||||
}
|
||||
|
||||
|
|
@ -1576,6 +1672,7 @@ class GenerationJobManagerClass {
|
|||
this.syncRunningJobMetrics();
|
||||
this.runStepBuffers?.clear();
|
||||
this.replayEventWriteQueues.clear();
|
||||
this.tokenUsageWriteQueues.clear();
|
||||
|
||||
logger.debug('[GenerationJobManager] Destroyed');
|
||||
}
|
||||
|
|
|
|||
96
packages/api/src/stream/__tests__/usageResume.spec.ts
Normal file
96
packages/api/src/stream/__tests__/usageResume.spec.ts
Normal file
|
|
@ -0,0 +1,96 @@
|
|||
import { InMemoryEventTransport } from '~/stream/implementations/InMemoryEventTransport';
|
||||
import { InMemoryJobStore } from '~/stream/implementations/InMemoryJobStore';
|
||||
import { GenerationJobManagerClass } from '~/stream/GenerationJobManager';
|
||||
|
||||
jest.spyOn(console, 'log').mockImplementation();
|
||||
|
||||
function createInMemoryManager(): GenerationJobManagerClass {
|
||||
const manager = new GenerationJobManagerClass();
|
||||
manager.configure({
|
||||
jobStore: new InMemoryJobStore({ ttlAfterComplete: 60000 }),
|
||||
eventTransport: new InMemoryEventTransport(),
|
||||
isRedis: false,
|
||||
});
|
||||
manager.initialize();
|
||||
return manager;
|
||||
}
|
||||
|
||||
describe('GenerationJobManager usage resume state', () => {
|
||||
let manager: GenerationJobManagerClass | undefined;
|
||||
|
||||
afterEach(async () => {
|
||||
await manager?.destroy();
|
||||
manager = undefined;
|
||||
});
|
||||
|
||||
test('accumulates persisted token usage events in resume state', async () => {
|
||||
manager = createInMemoryManager();
|
||||
const streamId = `usage-resume-${Date.now()}`;
|
||||
await manager.createJob(streamId, 'user-1', streamId);
|
||||
|
||||
const firstUsage = {
|
||||
input_tokens: 1200,
|
||||
output_tokens: 300,
|
||||
input_token_details: { cache_creation: 100, cache_read: 800 },
|
||||
model: 'claude-3-5-sonnet',
|
||||
provider: 'anthropic',
|
||||
};
|
||||
const secondUsage = {
|
||||
input_tokens: 1600,
|
||||
output_tokens: 120,
|
||||
model: 'claude-3-5-sonnet',
|
||||
provider: 'anthropic',
|
||||
usage_type: 'summarization',
|
||||
};
|
||||
|
||||
await manager.emitChunk(streamId, { event: 'on_token_usage', data: firstUsage });
|
||||
await manager.emitChunk(streamId, { event: 'on_token_usage', data: secondUsage });
|
||||
|
||||
const resumeState = await manager.getResumeState(streamId);
|
||||
expect(resumeState?.collectedUsage).toEqual([firstUsage, secondUsage]);
|
||||
});
|
||||
|
||||
test('omits collected usage when none was recorded', async () => {
|
||||
manager = createInMemoryManager();
|
||||
const streamId = `usage-resume-empty-${Date.now()}`;
|
||||
await manager.createJob(streamId, 'user-1', streamId);
|
||||
|
||||
const resumeState = await manager.getResumeState(streamId);
|
||||
expect(resumeState?.collectedUsage).toBeUndefined();
|
||||
expect(resumeState?.contextUsage).toBeUndefined();
|
||||
});
|
||||
|
||||
test('persists the latest context usage snapshot for resume', async () => {
|
||||
manager = createInMemoryManager();
|
||||
const streamId = `context-resume-${Date.now()}`;
|
||||
await manager.createJob(streamId, 'user-1', streamId);
|
||||
|
||||
const makeSnapshot = (messageTokens: number) => ({
|
||||
runId: 'run-1',
|
||||
agentId: 'agent-1',
|
||||
breakdown: {
|
||||
maxContextTokens: 200000,
|
||||
instructionTokens: 1500,
|
||||
systemMessageTokens: 1000,
|
||||
dynamicInstructionTokens: 0,
|
||||
toolSchemaTokens: 500,
|
||||
summaryTokens: 0,
|
||||
toolCount: 3,
|
||||
messageCount: 4,
|
||||
messageTokens,
|
||||
availableForMessages: 188500,
|
||||
},
|
||||
contextBudget: 190000,
|
||||
effectiveInstructionTokens: 1500,
|
||||
prePruneContextTokens: messageTokens,
|
||||
remainingContextTokens: 190000 - 1500 - messageTokens,
|
||||
calibrationRatio: 1.05,
|
||||
});
|
||||
|
||||
await manager.emitChunk(streamId, { event: 'on_context_usage', data: makeSnapshot(4000) });
|
||||
await manager.emitChunk(streamId, { event: 'on_context_usage', data: makeSnapshot(9000) });
|
||||
|
||||
const resumeState = await manager.getResumeState(streamId);
|
||||
expect(resumeState?.contextUsage).toEqual(makeSnapshot(9000));
|
||||
});
|
||||
});
|
||||
|
|
@ -923,6 +923,8 @@ export class RedisJobStore implements IJobStore {
|
|||
promptTokens: data.promptTokens ? parseInt(data.promptTokens, 10) : undefined,
|
||||
titleEvent: data.titleEvent || undefined,
|
||||
replayEvents: data.replayEvents || undefined,
|
||||
contextUsage: data.contextUsage || undefined,
|
||||
tokenUsage: data.tokenUsage || undefined,
|
||||
};
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -48,6 +48,12 @@ export interface SerializableJobData {
|
|||
/** Serialized replay-only stream events for active-stream resume */
|
||||
replayEvents?: string;
|
||||
|
||||
/** Serialized latest context usage snapshot for active-stream resume */
|
||||
contextUsage?: string;
|
||||
|
||||
/** Serialized token usage events for active-stream resume (cross-replica safe) */
|
||||
tokenUsage?: string;
|
||||
|
||||
/** Endpoint metadata for abort handling - avoids storing functions */
|
||||
endpoint?: string;
|
||||
iconURL?: string;
|
||||
|
|
|
|||
|
|
@ -148,6 +148,8 @@ export const deletePreset = () => `${BASE_URL}/api/presets/delete`;
|
|||
|
||||
export const aiEndpoints = () => `${BASE_URL}/api/endpoints`;
|
||||
|
||||
export const tokenConfig = () => `${BASE_URL}/api/endpoints/token-config`;
|
||||
|
||||
export const models = () => `${BASE_URL}/api/models`;
|
||||
|
||||
export const tokenizer = () => `${BASE_URL}/api/tokenizer`;
|
||||
|
|
|
|||
|
|
@ -1167,6 +1167,8 @@ export const interfaceSchema = z
|
|||
retainAgentFiles: z.boolean().optional(),
|
||||
runCode: z.boolean().optional(),
|
||||
webSearch: z.boolean().optional(),
|
||||
contextUsage: z.boolean().optional(),
|
||||
contextCost: z.boolean().optional(),
|
||||
peoplePicker: z
|
||||
.object({
|
||||
users: z.boolean().optional(),
|
||||
|
|
@ -1236,6 +1238,8 @@ export const interfaceSchema = z
|
|||
autoSubmitFromUrl: true,
|
||||
runCode: true,
|
||||
webSearch: true,
|
||||
contextUsage: true,
|
||||
contextCost: false,
|
||||
peoplePicker: {
|
||||
users: true,
|
||||
groups: true,
|
||||
|
|
|
|||
|
|
@ -245,6 +245,10 @@ export const getAIEndpoints = (): Promise<t.TEndpointsConfig> => {
|
|||
return request.get(endpoints.aiEndpoints());
|
||||
};
|
||||
|
||||
export const getTokenConfig = (): Promise<t.TTokenConfigMap> => {
|
||||
return request.get(endpoints.tokenConfig());
|
||||
};
|
||||
|
||||
export const getModels = async (): Promise<t.TModelsConfig> => {
|
||||
return request.get(endpoints.models());
|
||||
};
|
||||
|
|
|
|||
|
|
@ -12,6 +12,7 @@ export enum QueryKeys {
|
|||
models = 'models',
|
||||
balance = 'balance',
|
||||
endpoints = 'endpoints',
|
||||
tokenConfig = 'tokenConfig',
|
||||
presets = 'presets',
|
||||
searchResults = 'searchResults',
|
||||
tokenCount = 'tokenCount',
|
||||
|
|
|
|||
|
|
@ -748,6 +748,8 @@ export const tMessageSchema = z.object({
|
|||
feedback: feedbackSchema.optional(),
|
||||
/** metadata */
|
||||
metadata: z.record(z.unknown()).optional(),
|
||||
/** Output tokens for assistant messages, calibrated prompt-side estimate for user messages */
|
||||
tokenCount: z.number().optional(),
|
||||
contextMeta: z
|
||||
.object({
|
||||
calibrationRatio: z
|
||||
|
|
|
|||
|
|
@ -454,6 +454,18 @@ export type TEndpointsConfig =
|
|||
|
||||
export type TModelsConfig = Record<string, string[]>;
|
||||
|
||||
/** Server-resolved context window and pricing for one model. Rates are USD per 1M tokens. */
|
||||
export type TModelTokenomics = {
|
||||
context?: number;
|
||||
prompt?: number;
|
||||
completion?: number;
|
||||
cacheWrite?: number;
|
||||
cacheRead?: number;
|
||||
};
|
||||
|
||||
/** endpoint → model → resolved tokenomics, from GET /api/endpoints/token-config */
|
||||
export type TTokenConfigMap = Record<string, Record<string, TModelTokenomics>>;
|
||||
|
||||
export type TUpdateTokenCountResponse = {
|
||||
count: number;
|
||||
};
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
/* eslint-disable @typescript-eslint/no-namespace */
|
||||
import type { FunctionToolCall, SummaryContentPart } from './assistants';
|
||||
import type { TTokenUsageEvent, TContextUsageEvent } from './runs';
|
||||
import type { TAttachment, TPlugin } from 'src/schemas';
|
||||
import { StepTypes, ContentTypes, ToolCallTypes } from './runs';
|
||||
|
||||
|
|
@ -231,6 +232,10 @@ export namespace Agents {
|
|||
data?: unknown;
|
||||
[key: string]: unknown;
|
||||
}>;
|
||||
/** Cumulative provider-reported usage for the run; backfills usage totals on resume */
|
||||
collectedUsage?: TTokenUsageEvent[];
|
||||
/** Latest context window snapshot; restores the usage gauge on resume */
|
||||
contextUsage?: TContextUsageEvent;
|
||||
}
|
||||
/**
|
||||
* Represents a run step delta i.e. any changed fields on a run step during
|
||||
|
|
|
|||
|
|
@ -40,6 +40,57 @@ export enum StepEvents {
|
|||
ON_SUBAGENT_UPDATE = 'on_subagent_update',
|
||||
}
|
||||
|
||||
/** Token-tracking event names streamed to the client (separate from StepEvents dispatch). */
|
||||
export enum UsageEvents {
|
||||
ON_CONTEXT_USAGE = 'on_context_usage',
|
||||
ON_TOKEN_USAGE = 'on_token_usage',
|
||||
}
|
||||
|
||||
/** Mirrors TokenBudgetBreakdown from @librechat/agents (data-provider cannot import it). */
|
||||
export type TTokenBudgetBreakdown = {
|
||||
maxContextTokens: number;
|
||||
instructionTokens: number;
|
||||
systemMessageTokens: number;
|
||||
dynamicInstructionTokens: number;
|
||||
toolSchemaTokens: number;
|
||||
summaryTokens: number;
|
||||
toolCount: number;
|
||||
messageCount: number;
|
||||
messageTokens: number;
|
||||
availableForMessages: number;
|
||||
};
|
||||
|
||||
/** Per-model-call context snapshot, dispatched after pruning and before the LLM call. */
|
||||
export type TContextUsageEvent = {
|
||||
runId?: string;
|
||||
agentId?: string;
|
||||
breakdown: TTokenBudgetBreakdown;
|
||||
/** Usable budget this call: maxContextTokens minus output reserve */
|
||||
contextBudget?: number;
|
||||
/** Calibrated instruction overhead actually applied this call */
|
||||
effectiveInstructionTokens?: number;
|
||||
/** Calibrated message tokens before pruning (excluding instructions) */
|
||||
prePruneContextTokens?: number;
|
||||
/** Tokens still free after instructions + pruned messages */
|
||||
remainingContextTokens?: number;
|
||||
calibrationRatio?: number;
|
||||
};
|
||||
|
||||
/** Provider-reported usage for a single completed model call. */
|
||||
export type TTokenUsageEvent = {
|
||||
input_tokens?: number;
|
||||
output_tokens?: number;
|
||||
total_tokens?: number;
|
||||
input_token_details?: {
|
||||
cache_creation?: number;
|
||||
cache_read?: number;
|
||||
};
|
||||
model?: string;
|
||||
provider?: string;
|
||||
usage_type?: 'summarization';
|
||||
runId?: string;
|
||||
};
|
||||
|
||||
/** Lifecycle phase carried on subagent-progress envelopes (mirrors SDK SubagentUpdatePhase). */
|
||||
export type SubagentUpdatePhase =
|
||||
| 'start'
|
||||
|
|
|
|||
|
|
@ -14,6 +14,7 @@ export {
|
|||
cacheTokenValues,
|
||||
premiumTokenValues,
|
||||
defaultRate,
|
||||
createTxMethods,
|
||||
permissionBitSupersets,
|
||||
partitionIssues,
|
||||
validateSkillName,
|
||||
|
|
|
|||
|
|
@ -89,7 +89,7 @@ import { createAgentMethods, type AgentMethods, type AgentDeps } from './agent';
|
|||
import { createConfigMethods, type ConfigMethods } from './config';
|
||||
|
||||
export { RoleConflictError, DEFAULT_REFRESH_TOKEN_EXPIRY, DEFAULT_SESSION_EXPIRY };
|
||||
export { tokenValues, cacheTokenValues, premiumTokenValues, defaultRate };
|
||||
export { tokenValues, cacheTokenValues, premiumTokenValues, defaultRate, createTxMethods };
|
||||
export { permissionBitSupersets };
|
||||
export {
|
||||
partitionIssues,
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue