上下文窗口预算分配:AI编程助手在有限token下的价值最大化策略

0 阅读

上下文窗口的稀缺与浪费

AI编程助手的能力,越来越看上下文质量。同样的模型,喂得好产出就稳,喂得差就胡说。但上下文窗口是稀缺资源。塞满无关代码,关键信息被挤出,回答跑偏。

常见的浪费场景很多。把整个文件塞进去,但实际只有两行相关。依赖列表全量塞,但当前任务只用一两个。历史对话全保留,旧指令反复干扰。

示例放太多,留给真正代码的预算所剩无几。上下文管理不是"塞越多越好"。而是"在有限预算内塞最高价值信息"。这需要明确的预算分配策略。

让每一类信息按权重竞争token,而非无序堆积。把上下文当预算管理,是AI编程工程化的关键一步。

预算分配与裁剪机制

上下文预算要分桶管理。每桶有上限,互不侵占。

系统提示桶。定义助手行为规范,几乎不变,固定分配。这部分省不得,否则输出风格漂移。

任务代码桶。当前编辑的代码块,是最核心的上下文。优先级最高,预算占比最大。

依赖代码桶。被引用的函数定义、类型签名。按引用频次召回,频次低的不进。

历史对话桶。之前的指令与回复。按时间衰减,旧消息优先裁剪。

示例桶。Few-shot示例锚定输出格式。数量不固定,剩余预算才分配。预算分配后,还要做信息密度排序。

密度高的内容(核心代码)放前面。密度低的(历史闲聊)放后面。超预算时从尾部裁剪,保留高密度部分。下面是预算分配的数据流:

flowchart TD
    A[总预算: N tokens] --> B[分桶: 系统/代码/依赖/历史/示例]
    B --> C[各自召回候选]
    C --> D[按信息密度排序]
    D --> E[按桶上限填充]
    E --> F{超预算?}
    F -->|是| G[尾部裁剪: 历史>示例>依赖]
    F -->|否| H[拼接最终上下文]
    G --> H
    style H fill:#e8f5e9

关键在"裁剪顺序"。不是按桶上限直接砍。而是按信息价值密度排序后,从价值最低的尾部裁。保证留在窗口内的都是高价值信息。

生产级实现

下面用代码描述一个上下文预算分配器。支持分桶、密度排序、超预算裁剪。带异常处理与默认回退。

from dataclasses import dataclass, field
from typing import Callable

@dataclass
class Chunk:
    """上下文片段:内容 + token 估算 + 密度分。
    密度分衡量该片段对当前任务的价值。
    分数低的优先被裁掉。"""
    bucket: str          # 所属桶
    content: str
    tokens: int
    density: float       # 信息密度评分
    source: str = ""     # 来源标记,便于追溯

# 默认桶预算占比:合计 1.0
DEFAULT_BUDGET_RATIO = {
    "system": 0.10,
    "code": 0.40,
    "deps": 0.20,
    "history": 0.15,
    "examples": 0.15,
}

def estimate_tokens(text: str) -> int:
    """粗略估算 token 数。
    中文按 1.5 字符/token,英文按 4 字符/token。
    真实系统应接 tokenizer,这里只做近似。
    估算偏差会直接影响预算分配,必要时校准。"""
    cn_chars = sum(1 for c in text if "\u4e00" <= c <= "\u9fff")
    other_chars = len(text) - cn_chars
    return int(cn_chars / 1.5 + other_chars / 4)

@dataclass
class BudgetAllocator:
    """预算分配器:按桶比例分配 token 上限。
    设计上桶之间互不侵占,避免高优桶被低优吃掉。
    比例可配置,但一旦设定要稳定,便于回归对比。"""
    total_budget: int
    ratios: dict[str, float] = field(
        default_factory=lambda: dict(DEFAULT_BUDGET_RATIO)
    )

    def bucket_limits(self) -> dict[str, int]:
        # 计算每个桶的 token 上限,舍入到整数
        return {
            b: int(self.total_budget * r)
            for b, r in self.ratios.items()
        }

    def allocate(self, candidates: list[Chunk]) -> list[Chunk]:
        """主分配函数:按桶分组,按密度降序填充。
        单桶超限时从尾部裁剪,保留高密度片段。
        异常时回退为空列表,不抛断流。"""
        if not candidates:
            return []
        try:
            limits = self.bucket_limits()
            result: list[Chunk] = []
            # 按桶分组
            buckets: dict[str, list[Chunk]] = {}
            for c in candidates:
                buckets.setdefault(c.bucket, []).append(c)

            for bucket, chunks in buckets.items():
                limit = limits.get(bucket, 0)
                # 按密度降序,密度高的优先保留
                chunks.sort(key=lambda x: x.density, reverse=True)
                used = 0
                for c in chunks:
                    if used + c.tokens <= limit:
                        result.append(c)
                        used += c.tokens
                    else:
                        # 超预算:当前片段及之后全部裁掉
                        break
            # 全局二次校验:防止单桶估算偏差导致超总预算
            return self._trim_to_total(result)
        except Exception as e:
            # 分配异常时回退为空,避免污染下游
            print(f"[allocator] 分配异常: {e}")
            return []

    def _trim_to_total(self, selected: list[Chunk]) -> list[Chunk]:
        """全局裁剪:若总 token 超预算,按密度升序裁掉尾部。
        裁剪顺序与桶内填充相反,先裁价值最低的。
        保证最终结果严格不超总预算。"""
        total = sum(c.tokens for c in selected)
        if total <= self.total_budget:
            return selected
        # 按密度升序,密度低的先被裁
        selected.sort(key=lambda x: x.density)
        while selected and total > self.total_budget:
            removed = selected.pop(0)
            total -= removed.tokens
        # 裁完按桶再重新分组,保持输出顺序稳定
        selected.sort(key=lambda x: (x.bucket, -x.density))
        return selected

def build_candidates(
    system_prompt: str,
    code_blocks: list[str],
    deps: list[str],
    history: list[str],
    examples: list[str],
    relevance: Callable[[str], float],
) -> list[Chunk]:
    """构建候选片段集合。
    relevance 函数为每个片段算密度分。
    业务自定义密度算法,工具不绑死标准。"""
    candidates: list[Chunk] = []
    if system_prompt:
        candidates.append(Chunk(
            bucket="system",
            content=system_prompt,
            tokens=estimate_tokens(system_prompt),
            density=1.0,  # 系统提示密度最高,永不裁
        ))
    for code in code_blocks:
        candidates.append(Chunk(
            bucket="code",
            content=code,
            tokens=estimate_tokens(code),
            density=relevance(code),
            source="editor",
        ))
    for dep in deps:
        candidates.append(Chunk(
            bucket="deps",
            content=dep,
            tokens=estimate_tokens(dep),
            density=relevance(dep),
            source="deps_index",
        ))
    for i, msg in enumerate(history):
        # 历史按位置衰减:越靠后密度越低
        density = 0.5 * (1 - i / max(len(history), 1))
        candidates.append(Chunk(
            bucket="history",
            content=msg,
            tokens=estimate_tokens(msg),
            density=density,
        ))
    for ex in examples:
        candidates.append(Chunk(
            bucket="examples",
            content=ex,
            tokens=estimate_tokens(ex),
            density=0.3,  # 示例密度固定低,预算紧时先裁
        ))
    return candidates

if __name__ == "__main__":
    allocator = BudgetAllocator(total_budget=4000)
    candidates = build_candidates(
        system_prompt="你是代码助手",
        code_blocks=["def add(a, b): return a + b"],
        deps=["def mul(a, b): return a * b"],
        history=["之前问过减法", "之前问过除法"],
        examples=["示例:输入1+1,输出2"],
        relevance=lambda x: 0.7,
    )
    selected = allocator.allocate(candidates)
    total = sum(c.tokens for c in selected)
    print(f"选中 {len(selected)} 片段,共 {total} tokens")

真实系统会接向量化召回。按当前任务嵌入,检索最相关的代码与依赖。并把分配结果做缓存,相似任务复用上次分配。

上下文窗口的预算分配的代价与边界

预算分配能提效,但有几个坑。

召回偏差。检索召回的"相关代码"可能不准。召回召回错了,再好的分配也救不回。召回质量是分配质量的前置条件。

token估算偏差。粗估与真实token数有差。差几个百分点,分配就漂移。关键任务应接真实tokenizer校准。

压缩损失。为塞更多内容,做摘要压缩。压缩丢细节,可能正好丢了关键。压缩是权衡,不是免费午餐。

桶比例僵化。固定比例适配不了所有任务。改个bug不需要示例,写新功能要更多依赖。比例应按任务类型动态调整。

缓存失效。缓存的分配结果,环境变了就废。代码改了,缓存的依赖召回就过时。缓存要带版本,源码变更即失效。

上下文预算管理的几个被忽视的实践要点:第一,密度评分必须结合任务意图,同样是函数定义,与任务相关的密度高,无关的密度低,不能只看代码本身。第二,裁剪后要输出"被裁掉了什么"的清单,让调用方能感知到上下文是否被截断,否则调用方会误以为模型看到了全部信息。第三,预算分配要预留约10%的"机动空间",给系统提示或紧急召回留余量,避免每次都顶满。最后,不同任务的桶比例要分开维护,写新功能、改bug、做重构三类的最佳配比差异很大,用一套比例硬套会劣化效果。

总结

上下文窗口是AI编程助手的稀缺资源。机制上以分桶预算管理,按信息密度排序填充,超预算从尾部裁剪。工程上靠分配器与召回缓存形成可控的上下文组装。落地路线:先定义桶与比例;实现token估算与密度评分;做分桶填充与全局裁剪;接向量化召回提升相关性;按任务类型动态调比例。在有限token内塞入高价值信息,是AI编程工程化的基本功。