mirror of
https://github.com/sipeed/picoclaw.git
synced 2026-07-28 01:27:58 +00:00
feat(pico): add SetTurnUsage and usage payload helper
This commit is contained in:
@@ -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 == "" {
|
||||
|
||||
@@ -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"
|
||||
PayloadKeyToolCalls = "tool_calls"
|
||||
PayloadKeyModelName = "model_name"
|
||||
PayloadKeyUsage = "usage"
|
||||
|
||||
MessageKindThought = "thought"
|
||||
MessageKindToolCalls = "tool_calls"
|
||||
|
||||
Reference in New Issue
Block a user