/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)
+	}
+}

internal/cli/cli.go

@@ -32,6 +32,21 @@ func formatWindow(n int64) string {
 	}
 }
 
+func thousands(n int64) string {
+	s := strconv.FormatInt(n, 10)
+	if len(s) <= 3 {
+		return s
+	}
+	var b strings.Builder
+	for i, c := range s {
+		if i > 0 && (len(s)-i)%3 == 0 {
+			b.WriteByte(',')
+		}
+		b.WriteRune(c)
+	}
+	return b.String()
+}
+
 func (h *Handler) Handle(input string) bool {
 	fields := strings.Fields(input)
 	cmd, args := fields[0], fields[1:]
@@ -89,6 +104,13 @@ func (h *Handler) Handle(input string) bool {
 		fmt.Printf("思考强度已设置为 %s\n", args[0])
 	case "/context":
 		fmt.Print(h.bot.ContextDump())
+		used, total := h.bot.ContextStats()
+		if total <= 0 {
+			fmt.Printf("上下文: 约 %s tokens（窗口大小未配置）\n", thousands(used))
+			return true
+		}
+		pct := float64(used) / float64(total) * 100
+		fmt.Printf("上下文窗口使用: %s / %s tokens (%.2f%%)\n", thousands(used), thousands(total), pct)
 	case "/tools":
 		for _, t := range h.bot.Tools() {
 			fmt.Println("  " + t)