mirror of
https://github.com/danny-avila/LibreChat.git
synced 2026-07-10 16:23:44 +00:00
🪟 feat: Faithful Over-Window Context Estimate via Prune Mirror and Overhead Reserve (#13959)
Some checks are pending
Docker Dev Branch Images Build / build (Dockerfile, lc-dev, node) (push) Waiting to run
Docker Dev Branch Images Build / build (Dockerfile.multi, lc-dev-api, api-build) (push) Waiting to run
GitNexus Index / index (push) Waiting to run
GitNexus Index / post-index (push) Blocked by required conditions
Some checks are pending
Docker Dev Branch Images Build / build (Dockerfile, lc-dev, node) (push) Waiting to run
Docker Dev Branch Images Build / build (Dockerfile.multi, lc-dev-api, api-build) (push) Waiting to run
GitNexus Index / index (push) Waiting to run
GitNexus Index / post-index (push) Blocked by required conditions
* ✨ feat: Mirror send-path pruning in the over-window context estimate For a snapshot-less branch whose tokens exceed the window, the send path prunes oldest-first (getMessagesWithinTokenLimit), so the next call can sit well under the window. The gauge previously clamped the full sum to 100%, hiding that headroom. Add prunedBranchTokens — a newest->oldest walk that keeps messages until the next would overflow the message budget (max minus the summary baseline), mirroring the pruner — and use it on the estimate path in place of the clamp. Approximation: omits the instruction/tool overhead and tool-call pairing the real pruner accounts for (unknowable for a snapshot-less branch); superseded by an exact snapshot once the branch is generated. * ✨ feat: Reserve cached instruction/tool overhead in the snapshot-less estimate The over-window prune mirror and the gauge couldn't account for the fixed instruction + tool-schema overhead the next call always sends, because a snapshot-less branch has no breakdown. The backend already emits that overhead in the ON_CONTEXT_USAGE breakdown, so cache it per agent/model (keyed endpoint::model::agentId, already inclusive of tool schemas) from the live usage events, then reserve it from the prune budget and add it to used so the estimate is consistent with snapshots. Falls back to message-only until the agent has run once this session. Surfaced as a System row in the estimate breakdown. * 🩹 fix: Address Codex review on the over-window estimate - Key the overhead cache by agentId when present. useTokenLimits resolves an agent to its real provider/model, so the reader keyed `provider::model::agent` while the writer stored `agents::::agent` — a cache miss for the main agents case. Both sides now resolve to `agent:<id>` (non-agent configs: endpoint:model). - Skip the overhead reserve when a summary baseline exists: computeSummaryUsedTokens already folds instruction/tool overhead into that marker, so adding it again double-counted on summarized branches. - Collapse the breakdown's input/output/estimated rows into one pruned Messages row when over-window pruning ran, so the popover matches the gauge instead of summing to the discarded pre-prune history.
This commit is contained in:
parent
0d1103f62f
commit
0789a04d11
7 changed files with 253 additions and 18 deletions
|
|
@ -145,13 +145,28 @@ export default function Breakdown({ view, showCost, currency }: BreakdownProps)
|
|||
max={maxTokens}
|
||||
/>
|
||||
)}
|
||||
<Row label={localize('com_ui_input')} value={view.branchTotals.input} />
|
||||
<Row
|
||||
label={localize('com_ui_output')}
|
||||
value={view.branchTotals.output + view.liveTokens}
|
||||
/>
|
||||
{view.estimatedTokens > 0 && (
|
||||
<Row label={localize('com_ui_context_estimated')} value={view.estimatedTokens} />
|
||||
{view.messagesPruned ? (
|
||||
/** Over-window: the per-category split no longer describes what's
|
||||
* sent, so show the pruned message total (incl. in-flight). */
|
||||
<Row
|
||||
label={localize('com_ui_context_messages')}
|
||||
value={view.messageTokens + view.liveTokens}
|
||||
max={maxTokens}
|
||||
/>
|
||||
) : (
|
||||
<>
|
||||
<Row label={localize('com_ui_input')} value={view.branchTotals.input} />
|
||||
<Row
|
||||
label={localize('com_ui_output')}
|
||||
value={view.branchTotals.output + view.liveTokens}
|
||||
/>
|
||||
{view.estimatedTokens > 0 && (
|
||||
<Row label={localize('com_ui_context_estimated')} value={view.estimatedTokens} />
|
||||
)}
|
||||
</>
|
||||
)}
|
||||
{view.overheadTokens > 0 && (
|
||||
<Row label={localize('com_ui_context_system')} value={view.overheadTokens} />
|
||||
)}
|
||||
{maxTokens == null && (
|
||||
<p className="text-xs text-text-secondary">{localize('com_ui_context_unknown')}</p>
|
||||
|
|
|
|||
|
|
@ -6,6 +6,8 @@ import type { TMessage, TConversation, TModelTokenomics } from 'librechat-data-p
|
|||
import type { BranchTotals, BranchUsage } from '~/utils/tokens';
|
||||
import type { ContextSnapshot } from '~/store/usage';
|
||||
import {
|
||||
overheadKey,
|
||||
getModelOverhead,
|
||||
liveTokensFamily,
|
||||
totalUsageFamily,
|
||||
removeUsageAtoms,
|
||||
|
|
@ -21,6 +23,7 @@ import {
|
|||
clearIndex,
|
||||
mergeUsage,
|
||||
sumTotalUsage,
|
||||
prunedBranchTokens,
|
||||
findBranchSnapshotAnchor,
|
||||
} from '~/utils';
|
||||
import { useLatestMessageId } from '~/hooks/Messages/useLatestMessage';
|
||||
|
|
@ -56,6 +59,15 @@ export interface TokenUsageView {
|
|||
/** Estimated tokens for count-less messages (in-flight tail excluded while
|
||||
* streaming); 0 on snapshots. Rendered as its own breakdown row. */
|
||||
estimatedTokens: number;
|
||||
/** Cached instruction + tool overhead applied to a snapshot-less estimate; 0 on
|
||||
* snapshots (which carry their own breakdown) and until the agent has run. */
|
||||
overheadTokens: number;
|
||||
/** Final message-token portion of a snapshot-less estimate (pruned when over
|
||||
* window, excludes live); 0 on snapshots. */
|
||||
messageTokens: number;
|
||||
/** True when over-window pruning replaced the raw message sum, so the breakdown
|
||||
* shows a single pruned Messages row instead of input/output/estimated. */
|
||||
messagesPruned: boolean;
|
||||
rates?: TModelTokenomics;
|
||||
}
|
||||
|
||||
|
|
@ -238,6 +250,9 @@ export default function useTokenUsage({
|
|||
totalCost: totalUsage.cost,
|
||||
liveTokens,
|
||||
estimatedTokens: 0,
|
||||
overheadTokens: 0,
|
||||
messageTokens: 0,
|
||||
messagesPruned: false,
|
||||
rates: limits.rates,
|
||||
};
|
||||
}
|
||||
|
|
@ -252,23 +267,56 @@ export default function useTokenUsage({
|
|||
* from re-summing the discarded pre-summary history (which otherwise pins the
|
||||
* gauge at 100% after a compaction). */
|
||||
const maxTokens = limits.maxContextTokens;
|
||||
const liveOnTail = liveTokens > 0;
|
||||
/** Fixed instruction + tool-schema overhead for this agent/model (the latter is
|
||||
* already folded into `instructionTokens`), cached from live usage events. The
|
||||
* client can't otherwise know it for a snapshot-less branch, so reserve it from
|
||||
* the prune budget and add it to used — making over-window pruning faithful and
|
||||
* the gauge consistent with snapshots. Skipped when a summary baseline exists:
|
||||
* `computeSummaryUsedTokens` already folds the overhead into that marker, so
|
||||
* adding it again would double-count. 0 until the agent has run once this
|
||||
* session (then falls back to message-only, as before). */
|
||||
const overheadTokens =
|
||||
branchTotals.summaryBaseline > 0
|
||||
? 0
|
||||
: getModelOverhead(
|
||||
overheadKey(
|
||||
limits.endpoint ?? conversation?.endpoint,
|
||||
limits.model ?? conversation?.model,
|
||||
conversation?.agent_id,
|
||||
),
|
||||
);
|
||||
/** When a stream is live the tail is the in-flight response, already counted
|
||||
* by `liveTokens`; drop its static estimate so a resumed/partial response
|
||||
* isn't double-counted on the estimate path. */
|
||||
const estimatedTokens = Math.max(
|
||||
0,
|
||||
branchTotals.estTokens - (liveTokens > 0 ? branchTotals.tailEstTokens : 0),
|
||||
branchTotals.estTokens - (liveOnTail ? branchTotals.tailEstTokens : 0),
|
||||
);
|
||||
const rawUsed =
|
||||
branchTotals.input +
|
||||
branchTotals.output +
|
||||
estimatedTokens +
|
||||
branchTotals.summaryBaseline +
|
||||
liveTokens;
|
||||
/** The send path prunes an over-window branch before calling the model, so the
|
||||
* live gauge never actually exceeds the window; clamp the display to the
|
||||
* window rather than show impossible values (e.g. 50k / 8k). */
|
||||
const usedTokens = maxTokens != null && maxTokens > 0 ? Math.min(rawUsed, maxTokens) : rawUsed;
|
||||
const rawMessageTokens = branchTotals.input + branchTotals.output + estimatedTokens;
|
||||
let messageTokens = rawMessageTokens;
|
||||
/** The send path prunes an over-window branch oldest-first before calling the
|
||||
* model, so the next call can sit well under the window even when the full
|
||||
* branch exceeds it. Mirror that: when the raw sum overflows the message window
|
||||
* (max minus the always-sent summary baseline and instruction overhead), report
|
||||
* the newest messages that actually fit instead of clamping the whole branch to
|
||||
* 100%. */
|
||||
if (maxTokens != null && maxTokens > 0) {
|
||||
const messageBudget = Math.max(0, maxTokens - branchTotals.summaryBaseline - overheadTokens);
|
||||
if (messageTokens > messageBudget) {
|
||||
messageTokens = prunedBranchTokens(
|
||||
conversationKey,
|
||||
branchTotals.tailId,
|
||||
messageBudget,
|
||||
liveOnTail,
|
||||
);
|
||||
}
|
||||
}
|
||||
/** When pruning replaced the raw sum, the per-category input/output/estimated
|
||||
* rows no longer describe what's sent, so the breakdown collapses them into a
|
||||
* single pruned Messages row to stay consistent with the gauge. */
|
||||
const messagesPruned = messageTokens < rawMessageTokens;
|
||||
const usedTokens = overheadTokens + branchTotals.summaryBaseline + messageTokens + liveTokens;
|
||||
return {
|
||||
usedTokens,
|
||||
maxTokens,
|
||||
|
|
@ -285,6 +333,9 @@ export default function useTokenUsage({
|
|||
totalCost: totalUsage.cost,
|
||||
liveTokens,
|
||||
estimatedTokens,
|
||||
overheadTokens,
|
||||
messageTokens,
|
||||
messagesPruned,
|
||||
rates: limits.rates,
|
||||
};
|
||||
}, [
|
||||
|
|
@ -297,5 +348,7 @@ export default function useTokenUsage({
|
|||
liveTokens,
|
||||
limits,
|
||||
branchSnapshot,
|
||||
conversationKey,
|
||||
conversation,
|
||||
]);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -9,10 +9,12 @@ import type {
|
|||
} from 'librechat-data-provider';
|
||||
import type { ContextSnapshot } from '~/store/usage';
|
||||
import {
|
||||
overheadKey,
|
||||
markUsageFolded,
|
||||
liveTokensFamily,
|
||||
totalUsageFamily,
|
||||
removeUsageAtoms,
|
||||
setModelOverhead,
|
||||
clearUsageFolded,
|
||||
calibrationFamily,
|
||||
pendingUsageFamily,
|
||||
|
|
@ -150,6 +152,18 @@ export default function useUsageHandler(): UsageHandlers {
|
|||
if (data.calibrationRatio != null && data.calibrationRatio > 0) {
|
||||
jotai.set(calibrationFamily(convoKey), data.calibrationRatio);
|
||||
}
|
||||
/** Cache the fixed instruction+tool overhead (already inclusive of tool
|
||||
* schemas) per agent/model so a snapshot-less branch of the same config can
|
||||
* reserve it in its prune budget and gauge. */
|
||||
const overhead = data.effectiveInstructionTokens ?? data.breakdown?.instructionTokens ?? 0;
|
||||
setModelOverhead(
|
||||
overheadKey(
|
||||
submission.conversation?.endpoint,
|
||||
submission.conversation?.model,
|
||||
data.agentId,
|
||||
),
|
||||
overhead,
|
||||
);
|
||||
streamCharsRef.current = 0;
|
||||
confirmedRef.current = 0;
|
||||
setLive(convoKey, 0);
|
||||
|
|
|
|||
34
client/src/store/usage.spec.ts
Normal file
34
client/src/store/usage.spec.ts
Normal file
|
|
@ -0,0 +1,34 @@
|
|||
import { overheadKey, setModelOverhead, getModelOverhead } from './usage';
|
||||
|
||||
describe('model overhead cache', () => {
|
||||
it('keys agents by agentId so the resolved reader and raw writer agree', () => {
|
||||
/** The writer stores under the raw `agents` submission (no resolved model);
|
||||
* the reader resolves to the agent's real provider/model. Both must produce
|
||||
* the same key, or the cache misses for the main agents case. */
|
||||
const writerKey = overheadKey('agents', '', 'agent_123');
|
||||
const readerKey = overheadKey('openAI', 'gpt-4', 'agent_123');
|
||||
expect(writerKey).toBe('agent:agent_123');
|
||||
expect(readerKey).toBe('agent:agent_123');
|
||||
|
||||
setModelOverhead(writerKey, 1500);
|
||||
expect(getModelOverhead(readerKey)).toBe(1500);
|
||||
});
|
||||
|
||||
it('keys non-agent configs by endpoint:model and round-trips overhead', () => {
|
||||
const key = overheadKey('openAI', 'gpt-4', null);
|
||||
expect(key).toBe('openAI::gpt-4');
|
||||
|
||||
/** Unknown config defaults to 0 (estimate falls back to message-only). */
|
||||
expect(getModelOverhead(key)).toBe(0);
|
||||
|
||||
setModelOverhead(key, 800);
|
||||
expect(getModelOverhead(key)).toBe(800);
|
||||
|
||||
/** Non-positive overhead is ignored, keeping the last good value. */
|
||||
setModelOverhead(key, 0);
|
||||
expect(getModelOverhead(key)).toBe(800);
|
||||
|
||||
/** A different config is isolated. */
|
||||
expect(getModelOverhead(overheadKey('google', 'gemini', null))).toBe(0);
|
||||
});
|
||||
});
|
||||
|
|
@ -127,6 +127,43 @@ export function clearUsageFolded(conversationId: string): void {
|
|||
foldedUsageKeys.delete(conversationId);
|
||||
}
|
||||
|
||||
/**
|
||||
* Per-agent/model instruction+tool overhead — the system prompt + tool-schema
|
||||
* tokens the next call always sends — cached from the breakdown the live
|
||||
* `ON_CONTEXT_USAGE` event emits. Keyed by agent/model (NOT per conversation),
|
||||
* so a snapshot-less branch (import / never-generated) can reuse the overhead
|
||||
* once the same agent has run anywhere this session, making its prune budget and
|
||||
* gauge account for the fixed overhead the client can't otherwise know. Never
|
||||
* cleared on convo switch; bounded by the number of distinct configs used.
|
||||
*/
|
||||
const modelOverhead = new Map<string, number>();
|
||||
|
||||
/** Stable cache key the writer (usage handler) and reader (estimate) build
|
||||
* identically. An agent resolves to its real provider/model only after its data
|
||||
* loads, so the reader (resolved) and writer (raw `agents` submission) would key
|
||||
* differently — key by `agentId` when present so both agree; non-agent configs
|
||||
* key by endpoint:model. */
|
||||
export function overheadKey(
|
||||
endpoint?: string | null,
|
||||
model?: string | null,
|
||||
agentId?: string | null,
|
||||
): string {
|
||||
if (agentId != null && agentId !== '') {
|
||||
return `agent:${agentId}`;
|
||||
}
|
||||
return `${endpoint ?? ''}::${model ?? ''}`;
|
||||
}
|
||||
|
||||
export function setModelOverhead(key: string, tokens: number): void {
|
||||
if (tokens > 0) {
|
||||
modelOverhead.set(key, tokens);
|
||||
}
|
||||
}
|
||||
|
||||
export function getModelOverhead(key: string): number {
|
||||
return modelOverhead.get(key) ?? 0;
|
||||
}
|
||||
|
||||
/** Jotai atomFamily entries are never GC'd — call on conversation switch/cleanup */
|
||||
export function removeUsageAtoms(conversationId: string): void {
|
||||
branchTotalsFamily.remove(conversationId);
|
||||
|
|
|
|||
|
|
@ -10,6 +10,7 @@ import {
|
|||
mergeUsage,
|
||||
setEntryUsage,
|
||||
sumTotalUsage,
|
||||
prunedBranchTokens,
|
||||
findBranchSnapshotAnchor,
|
||||
estimateTokens,
|
||||
normalizeUsageUnits,
|
||||
|
|
@ -218,6 +219,34 @@ describe('token index', () => {
|
|||
expect(totals.tailEstTokens).toBe(5);
|
||||
});
|
||||
|
||||
describe('prunedBranchTokens (over-window mirror of getMessagesWithinTokenLimit)', () => {
|
||||
/** u1 ← a1(huge, old) ← u2 ← a2(tail). */
|
||||
const buildChain = () =>
|
||||
buildIndex(CONVO, [
|
||||
msg('u1', Constants.NO_PARENT, true, 2),
|
||||
msg('a1', 'u1', false, 10),
|
||||
msg('u2', 'a1', true, 2),
|
||||
msg('a2', 'u2', false, 2),
|
||||
]);
|
||||
|
||||
it('keeps the newest messages that fit and stops at the first overflow', () => {
|
||||
buildChain();
|
||||
/** Budget 8: a2(2)+u2(2)=4 fit; a1(10) would overflow → pruned. */
|
||||
expect(prunedBranchTokens(CONVO, 'a2', 8, false)).toBe(4);
|
||||
});
|
||||
|
||||
it('returns the full branch sum when it fits the budget', () => {
|
||||
buildChain();
|
||||
expect(prunedBranchTokens(CONVO, 'a2', 100, false)).toBe(16);
|
||||
});
|
||||
|
||||
it('skips the in-flight tail when excludeTail is set', () => {
|
||||
buildChain();
|
||||
/** Skip a2; a1(10)+u2(2)+u1(2)=14 all fit under 100. */
|
||||
expect(prunedBranchTokens(CONVO, 'a2', 100, true)).toBe(14);
|
||||
});
|
||||
});
|
||||
|
||||
it('caps the branch at a summary marker instead of re-summing compacted history', () => {
|
||||
const summarized = {
|
||||
messageId: 'a2',
|
||||
|
|
|
|||
|
|
@ -387,6 +387,59 @@ export function sumBranch(
|
|||
return { ...totals, tailEstTokens, tailId, usage, summaryBaseline };
|
||||
}
|
||||
|
||||
/**
|
||||
* Message tokens that would actually be sent for an over-window branch. The send
|
||||
* path prunes oldest-first to fit (`getMessagesWithinTokenLimit`), so walk the
|
||||
* branch newest→oldest and stop once the next message would exceed `budget`,
|
||||
* mirroring its "newest-that-fits" behavior for the gauge. Approximation: it omits
|
||||
* the instruction/tool-schema overhead and tool-call pairing the real pruner also
|
||||
* accounts for, which the client can't know for a snapshot-less branch — close
|
||||
* enough for an estimate, and superseded by an exact snapshot once generated.
|
||||
* `budget` is the message window (max minus the always-sent summary baseline);
|
||||
* when `excludeTail`, the in-flight tail response is skipped (it rides on
|
||||
* `liveTokens`). Per-message contribution matches `sumBranch`: stored `tokenCount`
|
||||
* when counted, else the char-based `estTokens`.
|
||||
*/
|
||||
export function prunedBranchTokens(
|
||||
conversationId: string,
|
||||
tailId: string | null | undefined,
|
||||
budget: number,
|
||||
excludeTail: boolean,
|
||||
): number {
|
||||
const index = registry.get(conversationId);
|
||||
if (!index || !tailId || budget <= 0) {
|
||||
return 0;
|
||||
}
|
||||
|
||||
let total = 0;
|
||||
let currentId: string | null = tailId;
|
||||
let guard = index.size;
|
||||
let isTail = true;
|
||||
|
||||
while (currentId && currentId !== Constants.NO_PARENT && guard-- > 0) {
|
||||
const entry: TokenEntry | undefined = index.get(currentId);
|
||||
if (!entry) {
|
||||
break;
|
||||
}
|
||||
const skip = isTail && excludeTail;
|
||||
isTail = false;
|
||||
if (!skip) {
|
||||
const contribution = entry.tokenCount > 0 ? entry.tokenCount : entry.estTokens;
|
||||
if (total + contribution > budget) {
|
||||
break;
|
||||
}
|
||||
total += contribution;
|
||||
}
|
||||
/** Pre-summary turns are subsumed by the baseline the caller already reserved,
|
||||
* so stop after counting the summarizing turn — mirrors `sumBranch`. */
|
||||
if (entry.summaryUsedTokens != null && entry.summaryUsedTokens > 0) {
|
||||
break;
|
||||
}
|
||||
currentId = entry.parentMessageId;
|
||||
}
|
||||
return total;
|
||||
}
|
||||
|
||||
/**
|
||||
* Sums provider usage/cost across EVERY message in the conversation (all
|
||||
* branches, including regenerated/abandoned responses) — the conversation
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue