account for stripCold in cost
This commit is contained in:
parent
e4c24e8a5f
commit
329edc3218
|
|
@ -61,7 +61,7 @@ func TestStripColdCancellation(t *testing.T) {
|
|||
done := make(chan struct{})
|
||||
|
||||
go func() {
|
||||
h.stripCold(ctx, b)
|
||||
_, _ = h.stripCold(ctx, b)
|
||||
close(done)
|
||||
}()
|
||||
|
||||
|
|
@ -99,7 +99,7 @@ func TestPreRunCompactionCancellation(t *testing.T) {
|
|||
Backend: b,
|
||||
},
|
||||
history: newHistory("test"),
|
||||
log: noOpLogger,
|
||||
log: noOpLogger,
|
||||
}
|
||||
|
||||
// Add enough messages to exceed the compact threshold.
|
||||
|
|
|
|||
|
|
@ -381,10 +381,10 @@ func (s *History) estimateTokens() int {
|
|||
// stripCold summarizes large tool-result messages outside the hot tail
|
||||
// using a single batched LLM call. Messages in the last hotTailSize slots
|
||||
// are left verbatim.
|
||||
func (s *History) stripCold(ctx context.Context, b backend.Backend) {
|
||||
func (s *History) stripCold(ctx context.Context, b backend.Backend) (backend.Usage, bool) {
|
||||
hot := len(s.messages) - hotTailSize
|
||||
if hot < 0 {
|
||||
return
|
||||
return backend.Usage{}, false
|
||||
}
|
||||
|
||||
// Collect cold-zone tool results exceeding 200 chars.
|
||||
|
|
@ -404,7 +404,7 @@ func (s *History) stripCold(ctx context.Context, b backend.Backend) {
|
|||
}
|
||||
}
|
||||
if len(items) == 0 {
|
||||
return
|
||||
return backend.Usage{}, false
|
||||
}
|
||||
|
||||
// Build a single prompt listing each item.
|
||||
|
|
@ -416,17 +416,27 @@ func (s *History) stripCold(ctx context.Context, b backend.Backend) {
|
|||
|
||||
ch, err := b.ChatStream(ctx, []backend.Message{{Role: "user", Content: sb.String()}}, nil, backend.GenerationParams{})
|
||||
if err != nil {
|
||||
return
|
||||
return backend.Usage{}, false
|
||||
}
|
||||
var resp strings.Builder
|
||||
var usage backend.Usage
|
||||
var reported bool
|
||||
for ev := range ch {
|
||||
if ev.Content != "" {
|
||||
resp.WriteString(ev.Content)
|
||||
}
|
||||
if ev.Done {
|
||||
usage = ev.Usage
|
||||
reported = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !reported {
|
||||
usage.InputTokens = len(sb.String()) / 4
|
||||
usage.OutputTokens = resp.Len() / 4
|
||||
}
|
||||
estimated := !reported
|
||||
|
||||
summaries := make(map[string]string)
|
||||
json.Unmarshal([]byte(resp.String()), &summaries)
|
||||
|
||||
|
|
@ -437,6 +447,7 @@ func (s *History) stripCold(ctx context.Context, b backend.Backend) {
|
|||
s.messages[it.idx].Content = summary
|
||||
}
|
||||
}
|
||||
return usage, estimated
|
||||
}
|
||||
|
||||
// resolveCompactionModel returns the model to use for compaction.
|
||||
|
|
|
|||
|
|
@ -204,7 +204,11 @@ func (ag *Agent) executeTurn(ctx context.Context, input string) string {
|
|||
|
||||
// Strip cold-zone tool results (one batched LLM call per turn).
|
||||
if ag.history != nil {
|
||||
ag.history.stripCold(actCtx, ag.runtime.Backend)
|
||||
usage, estimated := ag.history.stripCold(actCtx, ag.runtime.Backend)
|
||||
if usage.InputTokens > 0 || usage.OutputTokens > 0 {
|
||||
ag.history.addUsage(usage, estimated)
|
||||
ag.notifyChange()
|
||||
}
|
||||
}
|
||||
|
||||
// Inject the preamble system message into history before running.
|
||||
|
|
@ -257,6 +261,9 @@ func (ag *Agent) executeTurn(ctx context.Context, input string) string {
|
|||
ag.flush()
|
||||
|
||||
if err != nil {
|
||||
if ag.history != nil {
|
||||
ag.history.recordTurnCost(ag.runtime.Backend.Model())
|
||||
}
|
||||
// Keep completed work — only remove cancelled tool results.
|
||||
if ag.history != nil {
|
||||
ag.history.removeCancelledToolResults()
|
||||
|
|
|
|||
Loading…
Reference in New Issue