《互联网产品的 AI 架构·模型接入与能力边界》6.2 上下文预算与裁剪

// file: src/prompt/context-budget.ts
// tier: L1
// verify: node --test test/prompt/context-budget.test.ts

import { AiError } from "../core/errors.ts";

export type ContextKind = "system" | "history" | "retrieval" | "tool";
export type TrimmableKind = Exclude<ContextKind, "system">;

export const CONTEXT_KINDS: ContextKind[] = ["system", "history", "retrieval", "tool"];
export const TRIMMABLE_KINDS: TrimmableKind[] = ["history", "retrieval", "tool"];

export interface ContextSegment {
  id: string;
  kind: ContextKind;
  text: string;
  /** 同类内越大越先保留：历史用新旧序，检索用相关性，工具结果用新旧序。 */
  priority: number;
}

export interface BudgetPlan {
  maxTokens: number;
  /** 每类硬上限；system 超限直接失败，其余按优先级裁到上限以内。 */
  caps: Record<ContextKind, number>;
  /** 全局仍超预算时的裁剪顺序，必须恰好包含三个可裁剪类别。 */
  trimOrder: TrimmableKind[];
}

export interface DroppedSegment {
  id: string;
  kind: ContextKind;
  tokens: number;
  reason: "class-cap" | "global-budget";
}

export interface BudgetResult {
  kept: ContextSegment[];
  dropped: DroppedSegment[];
  usedTokens: number;
  byKind: Record<ContextKind, number>;
}

const CJK = /[　-〿㐀-䶿一-鿿豈-﫿＀-￯]/u;

/**
 * 可复现的近似口径：CJK 与全角字符按 1 字符 1 token，其余每 4 字符 1 token，
 * 逐段向上取整。它不是任何提供方的真实分词结果，只用于预算比较与回归测试。
 */
export function estimateTokens(text: string): number {
  if (text === "") return 0;
  let tokens = 0;
  let latin = 0;
  for (const char of text) {
    if (CJK.test(char)) {
      tokens += 1;
      if (latin > 0) {
        tokens += Math.ceil(latin / 4);
        latin = 0;
      }
    } else {
      latin += 1;
    }
  }
  if (latin > 0) tokens += Math.ceil(latin / 4);
  return tokens;
}

function fail(message: string): never {
  throw new AiError("invalid-request", `context-budget: ${message}`);
}

function validatePlan(plan: BudgetPlan): void {
  if (!Number.isInteger(plan.maxTokens) || plan.maxTokens <= 0) {
    fail("maxTokens must be a positive integer");
  }
  for (const kind of CONTEXT_KINDS) {
    const cap = plan.caps[kind];
    if (cap !== Number.POSITIVE_INFINITY && (!Number.isInteger(cap) || cap < 0)) {
      fail(`caps.${kind} must be a non-negative integer or Infinity`);
    }
  }
  const order = plan.trimOrder;
  if (
    order.length !== TRIMMABLE_KINDS.length ||
    new Set(order).size !== order.length ||
    TRIMMABLE_KINDS.some((kind) => !order.includes(kind)) ||
    order.some((kind) => !TRIMMABLE_KINDS.includes(kind))
  ) {
    fail("trimOrder must contain history, retrieval and tool exactly once");
  }
}

/**
 * 两阶段裁剪：先按类别上限保留高优先级片段，再按 trimOrder 从低优先级开始
 * 释放全局预算。system 只检查、不裁剪；它放不下时整个请求失败。
 */
export function planContext(plan: BudgetPlan, segments: ContextSegment[]): BudgetResult {
  validatePlan(plan);

  const seen = new Set<string>();
  for (const segment of segments) {
    if (segment.id.trim() === "") fail("segment id must not be empty");
    if (seen.has(segment.id)) fail(`duplicate segment id "${segment.id}"`);
    seen.add(segment.id);
    if (!CONTEXT_KINDS.includes(segment.kind)) fail(`unknown kind "${segment.kind}"`);
    if (!Number.isFinite(segment.priority)) fail(`priority for "${segment.id}" must be finite`);
  }

  const items = segments.map((segment, index) => ({
    segment,
    index,
    tokens: estimateTokens(segment.text),
  }));

  const systemTokens = items
    .filter((item) => item.segment.kind === "system")
    .reduce((sum, item) => sum + item.tokens, 0);
  if (systemTokens > plan.caps.system || systemTokens > plan.maxTokens) {
    throw new AiError(
      "budget-exceeded",
      `context-budget: system prompt needs ${systemTokens} tokens and is never trimmed`,
    );
  }

  const kept = new Set<string>();
  for (const item of items) {
    if (item.segment.kind === "system") kept.add(item.segment.id);
  }
  const dropped: DroppedSegment[] = [];

  for (const kind of TRIMMABLE_KINDS) {
    const candidates = items
      .filter((item) => item.segment.kind === kind)
      .sort((left, right) => right.segment.priority - left.segment.priority || left.index - right.index);
    let used = 0;
    for (const item of candidates) {
      if (used + item.tokens <= plan.caps[kind]) {
        kept.add(item.segment.id);
        used += item.tokens;
      } else {
        dropped.push({ id: item.segment.id, kind, tokens: item.tokens, reason: "class-cap" });
      }
    }
  }

  let total = items
    .filter((item) => kept.has(item.segment.id))
    .reduce((sum, item) => sum + item.tokens, 0);

  for (const kind of plan.trimOrder) {
    if (total <= plan.maxTokens) break;
    const candidates = items
      .filter((item) => item.segment.kind === kind && kept.has(item.segment.id))
      .sort((left, right) => left.segment.priority - right.segment.priority || right.index - left.index);
    for (const item of candidates) {
      if (total <= plan.maxTokens) break;
      kept.delete(item.segment.id);
      dropped.push({ id: item.segment.id, kind, tokens: item.tokens, reason: "global-budget" });
      total -= item.tokens;
    }
  }

  if (total > plan.maxTokens) {
    throw new AiError("budget-exceeded", `context-budget: context needs ${total} tokens after trimming`);
  }

  const position = new Map<string, number>();
  items.forEach((item) => position.set(item.segment.id, item.index));
  dropped.sort((left, right) => position.get(left.id)! - position.get(right.id)!);

  const byKind: Record<ContextKind, number> = { system: 0, history: 0, retrieval: 0, tool: 0 };
  const keptSegments: ContextSegment[] = [];
  for (const item of items) {
    if (!kept.has(item.segment.id)) continue;
    keptSegments.push(item.segment);
    byKind[item.segment.kind] += item.tokens;
  }

  return { kept: keptSegments, dropped, usedTokens: total, byKind };
}

node --test test/prompt/context-budget.test.ts

✔ token 估算口径可复现 (1.054167ms)
✔ system 永远保留，先牺牲工具结果 (0.708375ms)
✔ 类别上限先裁掉低优先级片段 (0.108667ms)
✔ 全局超预算按 trimOrder 依次释放 (0.124625ms)
✔ 保留片段维持原始输入顺序 (0.110834ms)
✔ system 超限直接失败，不裁 system (0.257084ms)
✔ trimOrder 必须覆盖三个可裁剪类别 (0.081917ms)
✔ 同一输入重复运行结果完全一致 (0.103834ms)
ℹ tests 8
ℹ suites 0
ℹ pass 8
ℹ fail 0
ℹ cancelled 0
ℹ skipped 0
ℹ todo 0
ℹ duration_ms 141.792