拆分 main 功能模块

Co-Authored-By: Claude <noreply@anthropic.com>

stream/tokens.go

@@ -0,0 +1,138 @@
+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
+}