/context 增加 token 统计：tiktoken 精确分词与窗口使用百分比

internal/bot/bot.go

@@ -11,6 +11,7 @@ import (
 	"github.com/openai/openai-go/packages/param"
 	"github.com/openai/openai-go/shared"
 	"github.com/openai/openai-go/shared/constant"
+	"github.com/pkoukk/tiktoken-go"
 	"github.com/tidwall/gjson"
 	"myaibot/internal/config"
 	"myaibot/internal/store"
@@ -134,6 +135,43 @@ func (b *Bot) ContextWindow() int64 {
 	return 0
 }
 
+// ContextStats 统计当前上下文的 token 使用量与窗口总大小（0 表示未配置）。
+func (b *Bot) ContextStats() (used, total int64) {
+	total = b.ContextWindow()
+	used += estimateTokens(b.systemPrompt)
+	for _, msg := range b.history {
+		var content string
+		switch {
+		case msg.OfUser != nil:
+			content = msg.OfUser.Content.OfString.Value
+		case msg.OfAssistant != nil:
+			content = msg.OfAssistant.Content.OfString.Value
+		case msg.OfSystem != nil:
+			content = msg.OfSystem.Content.OfString.Value
+		}
+		used += estimateTokens(content)
+	}
+	return used, total
+}
+
+var tke *tiktoken.Tiktoken
+
+// estimateTokens 用 o200k_base 词表精确统计 token；
+// 初始化失败（如无法下载词表）时回退为字符数/2 估算。
+func estimateTokens(s string) int64 {
+	if s == "" {
+		return 0
+	}
+	if tke == nil {
+		t, err := tiktoken.GetEncoding("o200k_base")
+		if err != nil {
+			return int64(len([]rune(s)) / 2)
+		}
+		tke = t
+	}
+	return int64(len(tke.Encode(s, nil, nil)))
+}
+
 func (b *Bot) SessionMessages() []store.Message {
 	out := make([]store.Message, 0, len(b.history)+1)
 	out = append(out, store.Message{Role: "system", Content: b.systemPrompt})

internal/bot/stats_test.go

@@ -0,0 +1,56 @@
+package bot
+
+import (
+	"testing"
+
+	"github.com/openai/openai-go"
+
+	"myaibot/internal/config"
+)
+
+func TestEstimateTokensEmpty(t *testing.T) {
+	if n := estimateTokens(""); n != 0 {
+		t.Errorf("空串应为 0, got %d", n)
+	}
+}
+
+func TestEstimateTokensKnown(t *testing.T) {
+	cases := []struct {
+		text string
+		want int64
+	}{
+		{"hello", 1},
+		{"hello world", 2},
+		{"你好", 1},
+		{"你是一个乐于助人的 AI 助手。", 12},
+	}
+	for _, c := range cases {
+		if n := estimateTokens(c.text); n != c.want {
+			t.Errorf("estimateTokens(%q) = %d, want %d", c.text, n, c.want)
+		}
+	}
+}
+
+func TestContextStats(t *testing.T) {
+	b := &Bot{
+		systemPrompt: "你是一个乐于助人的 AI 助手。",
+		provider: &config.Provider{
+			Name: "p",
+			Models: []config.ModelConfig{
+				{Name: "m", ContextWindow: 1000},
+			},
+		},
+		model: "m",
+		history: []openai.ChatCompletionMessageParamUnion{
+			openai.UserMessage("hello"),
+			openai.AssistantMessage("world"),
+		},
+	}
+	used, total := b.ContextStats()
+	if total != 1000 {
+		t.Errorf("total = %d, want 1000", total)
+	}
+	if used != 14 { // 系统提示 12 + hello 1 + world 1
+		t.Errorf("used = %d, want 14", used)
+	}
+}