feat(pico): add SetTurnUsage and usage payload helper
This commit is contained in:
parent
287853ab58
commit
052c742fe7
3 changed files with 64 additions and 0 deletions
|
|
@ -531,6 +531,8 @@ type picoStreamer struct {
|
||||||
channel *PicoChannel
|
channel *PicoChannel
|
||||||
chatID string
|
chatID string
|
||||||
modelName string
|
modelName string
|
||||||
|
turnInputTokens int
|
||||||
|
turnOutputTokens int
|
||||||
messageID string
|
messageID string
|
||||||
reasoningID string
|
reasoningID string
|
||||||
throttleInterval time.Duration
|
throttleInterval time.Duration
|
||||||
|
|
@ -553,6 +555,17 @@ func (s *picoStreamer) SetModelName(modelName string) {
|
||||||
s.modelName = strings.TrimSpace(modelName)
|
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 {
|
func (s *picoStreamer) Update(ctx context.Context, content string) error {
|
||||||
s.mu.Lock()
|
s.mu.Lock()
|
||||||
defer s.mu.Unlock()
|
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) {
|
func picoToolCallsPayload(msg bus.OutboundMessage) ([]utils.VisibleToolCall, bool) {
|
||||||
raw := strings.TrimSpace(msg.Context.Raw[PayloadKeyToolCalls])
|
raw := strings.TrimSpace(msg.Context.Raw[PayloadKeyToolCalls])
|
||||||
if raw == "" {
|
if raw == "" {
|
||||||
|
|
|
||||||
36
pkg/channels/pico/pico_usage_test.go
Normal file
36
pkg/channels/pico/pico_usage_test.go
Normal 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)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
@ -28,6 +28,7 @@ const (
|
||||||
PayloadKeyPlaceholder = "placeholder"
|
PayloadKeyPlaceholder = "placeholder"
|
||||||
PayloadKeyToolCalls = "tool_calls"
|
PayloadKeyToolCalls = "tool_calls"
|
||||||
PayloadKeyModelName = "model_name"
|
PayloadKeyModelName = "model_name"
|
||||||
|
PayloadKeyUsage = "usage"
|
||||||
|
|
||||||
MessageKindThought = "thought"
|
MessageKindThought = "thought"
|
||||||
MessageKindToolCalls = "tool_calls"
|
MessageKindToolCalls = "tool_calls"
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue