feat(pico): add SetTurnUsage and usage payload helper

This commit is contained in:
Andy Lo-A-Foe
2026-06-22 13:07:48 +02:00
parent 287853ab58
commit 052c742fe7
3 changed files with 64 additions and 0 deletions
+27
View File
@@ -531,6 +531,8 @@ type picoStreamer struct {
channel *PicoChannel
chatID string
modelName string
turnInputTokens int
turnOutputTokens int
messageID string
reasoningID string
throttleInterval time.Duration
@@ -553,6 +555,17 @@ func (s *picoStreamer) SetModelName(modelName string) {
s.modelName = strings.TrimSpace(modelName)
}
// SetTurnUsage records the real per-turn LLM token usage to emit on finalize.
func (s *picoStreamer) SetTurnUsage(inputTokens, outputTokens int) {
if s == nil {
return
}
s.mu.Lock()
defer s.mu.Unlock()
s.turnInputTokens = inputTokens
s.turnOutputTokens = outputTokens
}
func (s *picoStreamer) Update(ctx context.Context, content string) error {
s.mu.Lock()
defer s.mu.Unlock()
@@ -1403,6 +1416,20 @@ func setContextUsagePayload(payload map[string]any, u *bus.ContextUsage) {
}
}
// setTurnUsagePayload attaches real per-turn LLM token usage to the payload.
// Input and output are kept separate (billed at different rates); total is a
// convenience sum. Omitted entirely when both counts are zero.
func setTurnUsagePayload(payload map[string]any, inputTokens, outputTokens int) {
if inputTokens <= 0 && outputTokens <= 0 {
return
}
payload[PayloadKeyUsage] = map[string]any{
"input_tokens": inputTokens,
"output_tokens": outputTokens,
"total_tokens": inputTokens + outputTokens,
}
}
func picoToolCallsPayload(msg bus.OutboundMessage) ([]utils.VisibleToolCall, bool) {
raw := strings.TrimSpace(msg.Context.Raw[PayloadKeyToolCalls])
if raw == "" {
+36
View File
@@ -0,0 +1,36 @@
package pico
import "testing"
func TestSetTurnUsagePayload(t *testing.T) {
t.Run("populates usage block when counts present", func(t *testing.T) {
payload := map[string]any{PayloadKeyContent: "hi"}
setTurnUsagePayload(payload, 1234, 567)
raw, ok := payload[PayloadKeyUsage]
if !ok {
t.Fatalf("expected %q key in payload", PayloadKeyUsage)
}
usage, ok := raw.(map[string]any)
if !ok {
t.Fatalf("usage block is not a map: %T", raw)
}
if usage["input_tokens"] != 1234 {
t.Errorf("input_tokens = %v, want 1234", usage["input_tokens"])
}
if usage["output_tokens"] != 567 {
t.Errorf("output_tokens = %v, want 567", usage["output_tokens"])
}
if usage["total_tokens"] != 1801 {
t.Errorf("total_tokens = %v, want 1801", usage["total_tokens"])
}
})
t.Run("omits usage block when both counts zero", func(t *testing.T) {
payload := map[string]any{PayloadKeyContent: "hi"}
setTurnUsagePayload(payload, 0, 0)
if _, ok := payload[PayloadKeyUsage]; ok {
t.Errorf("expected no %q key when counts are zero", PayloadKeyUsage)
}
})
}
+1
View File
@@ -28,6 +28,7 @@ const (
PayloadKeyPlaceholder = "placeholder"
PayloadKeyToolCalls = "tool_calls"
PayloadKeyModelName = "model_name"
PayloadKeyUsage = "usage"
MessageKindThought = "thought"
MessageKindToolCalls = "tool_calls"