139 lines
3.1 KiB
Go
139 lines
3.1 KiB
Go
package stream
|
|
|
|
import (
|
|
"context"
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
"unicode"
|
|
|
|
"aichat/message"
|
|
)
|
|
|
|
type Stats struct {
|
|
PromptTokens int `json:"prompt_tokens"`
|
|
CompletionTokens int `json:"completion_tokens"`
|
|
ToolPromptTokens int `json:"tool_prompt_tokens"`
|
|
ToolCompletionTokens int `json:"tool_completion_tokens"`
|
|
TotalTokens int `json:"total_tokens"`
|
|
CompletionTokensPerSec float64 `json:"completion_tokens_per_sec"`
|
|
PeakCompletionTokensPerSec float64 `json:"peak_completion_tokens_per_sec"`
|
|
Estimated bool `json:"estimated"`
|
|
}
|
|
|
|
type Tracker struct {
|
|
mu sync.Mutex
|
|
promptTokens int
|
|
completionTokens int
|
|
toolPromptTokens int
|
|
toolCompletionTokens int
|
|
}
|
|
|
|
type tokenUsageContextKey struct{}
|
|
|
|
func NewTracker() *Tracker {
|
|
return &Tracker{}
|
|
}
|
|
|
|
func ContextWithTracker(ctx context.Context, tracker *Tracker) context.Context {
|
|
if tracker == nil {
|
|
return ctx
|
|
}
|
|
return context.WithValue(ctx, tokenUsageContextKey{}, tracker)
|
|
}
|
|
|
|
func TrackerFromContext(ctx context.Context) *Tracker {
|
|
tracker, _ := ctx.Value(tokenUsageContextKey{}).(*Tracker)
|
|
return tracker
|
|
}
|
|
|
|
func (t *Tracker) AddTool(promptTokens, completionTokens int) {
|
|
if t == nil {
|
|
return
|
|
}
|
|
t.mu.Lock()
|
|
defer t.mu.Unlock()
|
|
t.toolPromptTokens += promptTokens
|
|
t.toolCompletionTokens += completionTokens
|
|
}
|
|
|
|
func (t *Tracker) SetModel(promptTokens, completionTokens int) {
|
|
if t == nil {
|
|
return
|
|
}
|
|
t.mu.Lock()
|
|
defer t.mu.Unlock()
|
|
t.promptTokens = promptTokens
|
|
t.completionTokens = completionTokens
|
|
}
|
|
|
|
func (t *Tracker) Snapshot(tokensPerSecond, peakTokensPerSecond float64) Stats {
|
|
if t == nil {
|
|
return Stats{Estimated: true}
|
|
}
|
|
t.mu.Lock()
|
|
defer t.mu.Unlock()
|
|
total := t.promptTokens + t.completionTokens + t.toolPromptTokens + t.toolCompletionTokens
|
|
return Stats{
|
|
PromptTokens: t.promptTokens,
|
|
CompletionTokens: t.completionTokens,
|
|
ToolPromptTokens: t.toolPromptTokens,
|
|
ToolCompletionTokens: t.toolCompletionTokens,
|
|
TotalTokens: total,
|
|
CompletionTokensPerSec: tokensPerSecond,
|
|
PeakCompletionTokensPerSec: peakTokensPerSecond,
|
|
Estimated: true,
|
|
}
|
|
}
|
|
|
|
func EstimateChatMessagesTokens(messages []message.ChatMessage) int {
|
|
total := 0
|
|
for _, msg := range messages {
|
|
total += EstimateTokenCount(msg.Role) + EstimateTokenCount(msg.Content) + 4
|
|
if msg.ImageURL != "" || msg.ImageURLAlias != "" {
|
|
total += 85
|
|
}
|
|
}
|
|
return total
|
|
}
|
|
|
|
func EstimateTokenCount(text string) int {
|
|
text = strings.TrimSpace(text)
|
|
if text == "" {
|
|
return 0
|
|
}
|
|
tokens := 0
|
|
asciiRunes := 0
|
|
flushASCII := func() {
|
|
if asciiRunes > 0 {
|
|
tokens += (asciiRunes + 3) / 4
|
|
asciiRunes = 0
|
|
}
|
|
}
|
|
for _, r := range text {
|
|
if unicode.IsSpace(r) {
|
|
flushASCII()
|
|
continue
|
|
}
|
|
if r <= unicode.MaxASCII {
|
|
asciiRunes++
|
|
continue
|
|
}
|
|
flushASCII()
|
|
tokens++
|
|
}
|
|
flushASCII()
|
|
if tokens == 0 {
|
|
return 1
|
|
}
|
|
return tokens
|
|
}
|
|
|
|
func TokensPerSecond(tokens int, start time.Time) float64 {
|
|
elapsed := time.Since(start).Seconds()
|
|
if tokens <= 0 || elapsed <= 0 {
|
|
return 0
|
|
}
|
|
return float64(tokens) / elapsed
|
|
}
|