Blizzard
faa1871760
token 用量估算计量（CJK≈1/字、ASCII≈1/4字，无需分词器，护栏够用）。 单任务硬上限（dispatcher）：Budget 挂 ctx 沿图透传，各 LLM 节点（对话/ReAct/compose/ 报告）入口计输入 token、出口计输出，触顶即中止整图——防失控成本（死循环/超大报告）。 报告路径优雅降级：触顶跳过剩余章节出部分稿，不整体失败。预算来源 Meta.token_budget 或 env TASK_TOKEN_BUDGET（默认 20 万）。 单用户日预算（gateway）：dispatcher 收尾经 NATS 回写 UsageEvent → 网关按用户按天累计 Redis（48h 过期自滚动）→ 提交前门控 USER_DAILY_TOKEN_BUDGET（0=不限，超额 402）。 /billing 升级为真实用量：当日已用 / 日预算 / 余额。 契约 UsageEvent + MetaTokenBudget + SubjectUsage；bus Publish/SubscribeUsage； orchestrator SetUsageSink + 预算触顶 failed（不计熔断）。harness budget 5 单测，三模块全绿。 live：单任务 budget=30 → failed(已用约689)；用户日 budget=200 → /billing remaining=0 → 402。 至此 harness 由「测温计」完成向「恒温器」的演进（评测闭环/纠偏/忠实度/脱敏/输入护栏/预算六项）。 Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
83 lines
2.3 KiB
Go
83 lines
2.3 KiB
Go
package harness
|
|
|
|
import (
|
|
"context"
|
|
"sync"
|
|
"unicode"
|
|
)
|
|
|
|
// EstimateTokens 粗估一段文本的 token 数（无需依赖模型/分词器，作预算护栏足够）：
|
|
// CJK 字符约 1 token/字；其余（拉丁/数字/空白/标点）约 1 token/4 字符。偏保守（宁高勿低）。
|
|
func EstimateTokens(s string) int {
|
|
cjk, other := 0, 0
|
|
for _, r := range s {
|
|
if unicode.Is(unicode.Han, r) || unicode.Is(unicode.Hiragana, r) ||
|
|
unicode.Is(unicode.Katakana, r) || unicode.Is(unicode.Hangul, r) {
|
|
cjk++
|
|
} else {
|
|
other++
|
|
}
|
|
}
|
|
return cjk + (other+3)/4 // 向上取整
|
|
}
|
|
|
|
// Budget 是单任务的 token 预算计量与封顶（并发安全：流式回调与节点循环可能并发计量）。
|
|
// max ≤ 0 表示不限额（仅计量、不中止）。
|
|
type Budget struct {
|
|
mu sync.Mutex
|
|
max int
|
|
prompt int
|
|
comp int
|
|
}
|
|
|
|
// NewBudget 建一个上限为 max 的任务预算（max≤0 → 不限额，纯计量）。
|
|
func NewBudget(max int) *Budget { return &Budget{max: max} }
|
|
|
|
// AddPrompt / AddComplete 累计输入/输出 token（按文本估算）。
|
|
func (b *Budget) AddPrompt(text string) { b.add(EstimateTokens(text), 0) }
|
|
func (b *Budget) AddComplete(text string) { b.add(0, EstimateTokens(text)) }
|
|
|
|
func (b *Budget) add(p, c int) {
|
|
if b == nil {
|
|
return
|
|
}
|
|
b.mu.Lock()
|
|
b.prompt += p
|
|
b.comp += c
|
|
b.mu.Unlock()
|
|
}
|
|
|
|
// Exceeded 报告是否已触顶（max≤0 恒为 false）。
|
|
func (b *Budget) Exceeded() bool {
|
|
if b == nil {
|
|
return false
|
|
}
|
|
b.mu.Lock()
|
|
defer b.mu.Unlock()
|
|
return b.max > 0 && b.prompt+b.comp > b.max
|
|
}
|
|
|
|
// Snapshot 返回当前用量快照（prompt / completion / total）。
|
|
func (b *Budget) Snapshot() (prompt, comp, total int) {
|
|
if b == nil {
|
|
return 0, 0, 0
|
|
}
|
|
b.mu.Lock()
|
|
defer b.mu.Unlock()
|
|
return b.prompt, b.comp, b.prompt + b.comp
|
|
}
|
|
|
|
// budgetKey 是 context 里承载任务预算的私有键（避免跨包碰撞）。
|
|
type budgetKey struct{}
|
|
|
|
// WithBudget 把任务预算挂到 context，沿图执行透传（免改各节点函数签名）。
|
|
func WithBudget(ctx context.Context, b *Budget) context.Context {
|
|
return context.WithValue(ctx, budgetKey{}, b)
|
|
}
|
|
|
|
// BudgetFrom 取出 context 里的任务预算（无则返回 nil，调用方按"不限额"处理）。
|
|
func BudgetFrom(ctx context.Context) *Budget {
|
|
b, _ := ctx.Value(budgetKey{}).(*Budget)
|
|
return b
|
|
}
|