Merge pull request 'Split last-request/cache usage, emit terminal usage, and retry stream errors by default (N3 WP6)' (#35) from feature/n3-usage-retry-wp6 into main
Split last-request/cache usage, emit terminal usage, and retry stream errors by default (N3 WP6)
This commit was merged in pull request #35.
This commit is contained in:
+33
-1
@@ -46,6 +46,12 @@ type Event struct {
|
||||
PromptTokens int
|
||||
CompletionTokens int
|
||||
TotalTokens int
|
||||
// 以下在 type=usage 时有意义(WP6:上一请求与缓存命中拆分)。
|
||||
LastPromptTokens int
|
||||
LastCompletionTokens int
|
||||
LastTotalTokens int
|
||||
CacheReadTokens int
|
||||
CacheWriteTokens int
|
||||
}
|
||||
|
||||
// Config drives one agent session.
|
||||
@@ -131,6 +137,10 @@ type Config struct {
|
||||
// MemoryScope 是 Recall 的 scope 过滤("" = 不限,含 user/project;
|
||||
// AgentSpec.MemoryScope 的消费点)。
|
||||
MemoryScope string
|
||||
// RetryAttempts 是流式错误的默认重试次数(WP6 下沉为库默认):
|
||||
// 0 值 = 2 次(默认);负值 = 关闭;正数 = 次数。
|
||||
// 只在还没向调用方发出增量时重试(WithRetry 既有语义)。
|
||||
RetryAttempts int
|
||||
// traceCtx 是本次 Run 的 span 关联(Run 启动时建立;指针随 Config
|
||||
// 副本传播到各轮与各工具)。
|
||||
traceCtx *traceIDs
|
||||
@@ -197,6 +207,18 @@ func Run(ctx context.Context, cfg Config, user llm.Message) (Result, error) {
|
||||
if cfg.Memory != nil && cfg.Tools != nil {
|
||||
memory.RegisterWrite(cfg.Tools, cfg.Memory, cfg.MemoryScope)
|
||||
}
|
||||
// 流式错误默认重试(WP6):库默认 2 次、负值关闭;重试事件进事件面。
|
||||
if cfg.Stream != nil {
|
||||
attempts := cfg.RetryAttempts
|
||||
if attempts == 0 {
|
||||
attempts = 2
|
||||
}
|
||||
if attempts > 0 {
|
||||
cfg.Stream = llm.WithRetry(cfg.Stream, attempts, func(kind, detail string) {
|
||||
emit(cfg, Event{Type: kind, Text: detail})
|
||||
})
|
||||
}
|
||||
}
|
||||
cfg.traceCtx = newTraceContext()
|
||||
endRun := startSpan(cfg, "invoke_agent "+agentName(cfg), map[string]string{
|
||||
"gen_ai.system": providerName(cfg),
|
||||
@@ -313,12 +335,16 @@ func Run(ctx context.Context, cfg Config, user llm.Message) (Result, error) {
|
||||
end, errEv := consume(ctx, cfg, request)
|
||||
if errEv != nil {
|
||||
cfg.logger().Error("runtime turn failed", errEv, "retryable", end.Retryable)
|
||||
emitUsage(cfg, usage)
|
||||
return Result{Text: end.Message.Content, Messages: messages, StopReason: end.Message.StopReason, Err: errEv, Retryable: end.Retryable}, nil
|
||||
}
|
||||
last = end.Message
|
||||
usage.PromptTokens += end.Usage.PromptTokens
|
||||
usage.CompletionTokens += end.Usage.CompletionTokens
|
||||
usage.TotalTokens += end.Usage.TotalTokens
|
||||
// cache 只入累计位(WP6 决策:不加 Last* cache 字段,避免字段膨胀)。
|
||||
usage.CacheReadTokens += end.Usage.CacheReadTokens
|
||||
usage.CacheWriteTokens += end.Usage.CacheWriteTokens
|
||||
usage.LastPromptTokens = end.Usage.PromptTokens
|
||||
usage.LastCompletionTokens = end.Usage.CompletionTokens
|
||||
usage.LastTotalTokens = end.Usage.TotalTokens
|
||||
@@ -1074,7 +1100,12 @@ func drain(fn func() []llm.Message) []llm.Message {
|
||||
}
|
||||
|
||||
func emitUsage(cfg Config, usage llm.Usage) {
|
||||
emit(cfg, Event{Type: "usage", PromptTokens: usage.PromptTokens, CompletionTokens: usage.CompletionTokens, TotalTokens: usage.TotalTokens})
|
||||
emit(cfg, Event{
|
||||
Type: "usage",
|
||||
PromptTokens: usage.PromptTokens, CompletionTokens: usage.CompletionTokens, TotalTokens: usage.TotalTokens,
|
||||
LastPromptTokens: usage.LastPromptTokens, LastCompletionTokens: usage.LastCompletionTokens, LastTotalTokens: usage.LastTotalTokens,
|
||||
CacheReadTokens: usage.CacheReadTokens, CacheWriteTokens: usage.CacheWriteTokens,
|
||||
})
|
||||
}
|
||||
|
||||
func emit(cfg Config, ev Event) {
|
||||
@@ -1107,6 +1138,7 @@ func budgetExceeded(b *Budget, usage llm.Usage, toolCalls int, now time.Time) st
|
||||
// budgetStop 以预算超限收尾:发事件、留日志、按 StopError 返回(不可重试)。
|
||||
func budgetStop(cfg Config, last llm.Message, messages []llm.Message, usage llm.Usage, reason string) Result {
|
||||
emit(cfg, Event{Type: "budget_exceeded", Text: reason})
|
||||
emitUsage(cfg, usage)
|
||||
err := errString(reason)
|
||||
cfg.logger().Error("runtime budget exceeded", err)
|
||||
return Result{Text: last.Content, Messages: messages, StopReason: llm.StopError, Usage: usage, Err: err}
|
||||
|
||||
@@ -0,0 +1,203 @@
|
||||
package runtime
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"agent-native/runtime/llm"
|
||||
"agent-native/runtime/tools"
|
||||
)
|
||||
|
||||
// ---- N3 WP6:计费拆分与默认重试 ----
|
||||
|
||||
// cache 累计 + usage 事件带 Last*/cache;终态路径(budgetStop)也发 usage。
|
||||
func TestUsageEventSplitsLastAndCache(t *testing.T) {
|
||||
var mu sync.Mutex
|
||||
var usageEvents []Event
|
||||
cfg := Config{
|
||||
Model: "fake",
|
||||
Stream: func() llm.StreamFunc {
|
||||
n := 0
|
||||
return func(ctx context.Context, model string, msgs []llm.Message, _ []llm.ToolDecl, _ llm.StreamOptions) <-chan llm.Event {
|
||||
n++
|
||||
ch := make(chan llm.Event, 1)
|
||||
ch <- llm.Event{Kind: llm.KindMessageEnd,
|
||||
Message: llm.Message{Role: llm.RoleAssistant, Content: "x", StopReason: llm.StopEnd},
|
||||
Usage: llm.Usage{PromptTokens: 50 * n, CompletionTokens: 5, TotalTokens: 50*n + 5},
|
||||
}
|
||||
close(ch)
|
||||
return ch
|
||||
}
|
||||
}(),
|
||||
OnEvent: func(ev Event) {
|
||||
if ev.Type == "usage" {
|
||||
mu.Lock()
|
||||
usageEvents = append(usageEvents, ev)
|
||||
mu.Unlock()
|
||||
}
|
||||
},
|
||||
}
|
||||
res, err := Run(context.Background(), cfg, llm.Message{Role: llm.RoleUser, Content: "go"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
_ = res
|
||||
if len(usageEvents) != 1 {
|
||||
t.Fatalf("usage 事件数 = %d", len(usageEvents))
|
||||
}
|
||||
ev := usageEvents[0]
|
||||
if ev.PromptTokens == 0 || ev.LastPromptTokens == 0 {
|
||||
t.Fatalf("usage 事件必须带累计与上一请求:%+v", ev)
|
||||
}
|
||||
// 两个回合的累计应≥上一请求(脚本不带 cache 字段时 cache 为 0 但键存在语义)。
|
||||
if ev.LastPromptTokens > ev.PromptTokens {
|
||||
t.Fatalf("上一请求不得大于累计:%+v", ev)
|
||||
}
|
||||
}
|
||||
|
||||
// cache tokens 进入累计(脚本带 cache 字段时)。
|
||||
func TestUsageAccumulatesCacheTokens(t *testing.T) {
|
||||
var ev Event
|
||||
stream := func(ctx context.Context, model string, msgs []llm.Message, _ []llm.ToolDecl, _ llm.StreamOptions) <-chan llm.Event {
|
||||
ch := make(chan llm.Event, 1)
|
||||
ch <- llm.Event{Kind: llm.KindMessageEnd,
|
||||
Message: llm.Message{Role: llm.RoleAssistant, Content: "x", StopReason: llm.StopEnd},
|
||||
Usage: llm.Usage{PromptTokens: 100, CompletionTokens: 10, TotalTokens: 110, CacheReadTokens: 60, CacheWriteTokens: 20},
|
||||
}
|
||||
close(ch)
|
||||
return ch
|
||||
}
|
||||
if _, err := Run(context.Background(), Config{
|
||||
Model: "fake", Stream: stream,
|
||||
OnEvent: func(e Event) {
|
||||
if e.Type == "usage" {
|
||||
ev = e
|
||||
}
|
||||
},
|
||||
}, llm.Message{Role: llm.RoleUser, Content: "go"}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if ev.CacheReadTokens != 60 || ev.CacheWriteTokens != 20 {
|
||||
t.Fatalf("cache 未累计:%+v", ev)
|
||||
}
|
||||
}
|
||||
|
||||
// 预算终态路径也发终态 usage。
|
||||
func TestBudgetStopEmitsTerminalUsage(t *testing.T) {
|
||||
var sawUsage bool
|
||||
regB := tools.NewRegistry(t.TempDir())
|
||||
regB.Register(&countingTool{name: "read_file"})
|
||||
resB, err := Run(context.Background(), Config{
|
||||
Model: "fake", Budget: &Budget{MaxToolCalls: 1}, Tools: regB,
|
||||
Stream: script(
|
||||
llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls,
|
||||
ToolCalls: []llm.ToolCall{{ID: "c1", Name: "read_file", Arguments: `{}`}}},
|
||||
llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls,
|
||||
ToolCalls: []llm.ToolCall{{ID: "c2", Name: "read_file", Arguments: `{}`}}},
|
||||
),
|
||||
OnEvent: func(e Event) {
|
||||
if e.Type == "usage" {
|
||||
sawUsage = true
|
||||
}
|
||||
},
|
||||
}, llm.Message{Role: llm.RoleUser, Content: "go"})
|
||||
if err != nil || resB.Err == nil {
|
||||
t.Fatalf("预算必须触发 budgetStop(Result.Err 非空):err=%v res.Err=%v", err, resB.Err)
|
||||
}
|
||||
if !sawUsage {
|
||||
t.Fatal("budgetStop 必须发终态 usage")
|
||||
}
|
||||
}
|
||||
|
||||
// 默认重试:0 值 = 2 次;事件进事件面;负值关闭。
|
||||
func TestDefaultRetryAttemptsTwo(t *testing.T) {
|
||||
var calls int
|
||||
var events []string
|
||||
stream := func(ctx context.Context, model string, msgs []llm.Message, _ []llm.ToolDecl, _ llm.StreamOptions) <-chan llm.Event {
|
||||
calls++
|
||||
ch := make(chan llm.Event, 3)
|
||||
if calls < 3 {
|
||||
ch <- llm.Event{Kind: llm.KindError, Retryable: true,
|
||||
Err: &llm.StreamError{Message: "boom", Retryable: true},
|
||||
Message: llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopError}}
|
||||
} else {
|
||||
ch <- llm.Event{Kind: llm.KindMessageEnd, Message: llm.Message{Role: llm.RoleAssistant, Content: "成了", StopReason: llm.StopEnd}}
|
||||
}
|
||||
close(ch)
|
||||
return ch
|
||||
}
|
||||
res, err := Run(context.Background(), Config{
|
||||
Model: "fake", Stream: stream,
|
||||
OnEvent: func(e Event) { events = append(events, e.Type) },
|
||||
}, llm.Message{Role: llm.RoleUser, Content: "go"})
|
||||
if err != nil {
|
||||
t.Fatalf("第 3 次应成功:%v", err)
|
||||
}
|
||||
if res.Err != nil {
|
||||
t.Fatalf("第 3 次应成功:%v", res.Err)
|
||||
}
|
||||
if calls != 3 {
|
||||
t.Fatalf("默认 2 次重试 + 首发 = 3 次调用,得到 %d", calls)
|
||||
}
|
||||
var starts int
|
||||
for _, t2 := range events {
|
||||
if t2 == "auto_retry_start" {
|
||||
starts++
|
||||
}
|
||||
}
|
||||
if starts != 2 {
|
||||
t.Fatalf("auto_retry_start 事件数 = %d:%v", starts, events)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNegativeRetryAttemptsDisables(t *testing.T) {
|
||||
var calls int
|
||||
stream := func(ctx context.Context, model string, msgs []llm.Message, _ []llm.ToolDecl, _ llm.StreamOptions) <-chan llm.Event {
|
||||
calls++
|
||||
ch := make(chan llm.Event, 1)
|
||||
ch <- llm.Event{Kind: llm.KindError, Retryable: true,
|
||||
Err: &llm.StreamError{Message: "boom", Retryable: true},
|
||||
Message: llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopError}}
|
||||
close(ch)
|
||||
return ch
|
||||
}
|
||||
res2, err := Run(context.Background(), Config{
|
||||
Model: "fake", Stream: stream, RetryAttempts: -1,
|
||||
}, llm.Message{Role: llm.RoleUser, Content: "go"})
|
||||
if err != nil || res2.Err == nil {
|
||||
t.Fatalf("必须失败(Result.Err 非空):err=%v res.Err=%v", err, res2.Err)
|
||||
}
|
||||
if calls != 1 {
|
||||
t.Fatalf("负值必须关闭重试(1 次调用),得到 %d", calls)
|
||||
}
|
||||
}
|
||||
|
||||
// 已出增量后不再重试(既有语义保持)。
|
||||
func TestRetryNeverDuplicatesForwardedDeltas(t *testing.T) {
|
||||
var calls int
|
||||
stream := func(ctx context.Context, model string, msgs []llm.Message, _ []llm.ToolDecl, _ llm.StreamOptions) <-chan llm.Event {
|
||||
calls++
|
||||
ch := make(chan llm.Event, 3)
|
||||
ch <- llm.Event{Kind: llm.KindTextDelta, Text: "半截"}
|
||||
ch <- llm.Event{Kind: llm.KindError, Retryable: true,
|
||||
Err: &llm.StreamError{Message: "boom", Retryable: true},
|
||||
Message: llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopError}}
|
||||
close(ch)
|
||||
return ch
|
||||
}
|
||||
res3, err := Run(context.Background(), Config{Model: "fake", Stream: stream}, llm.Message{Role: llm.RoleUser, Content: "go"})
|
||||
if err != nil || res3.Err == nil {
|
||||
t.Fatalf("出增量后错误不得重试但须失败:err=%v res.Err=%v", err, res3.Err)
|
||||
}
|
||||
if calls != 1 {
|
||||
t.Fatalf("出增量后不得重试(1 次调用),得到 %d", calls)
|
||||
}
|
||||
}
|
||||
|
||||
var _ = errors.New
|
||||
var _ = strings.Contains
|
||||
var _ = time.Now
|
||||
Reference in New Issue
Block a user