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

2026-08-14 19:10:48 +08:00
parent bb5bc619c9
commit ba2ee4b0a0
6 changed files with 124 additions and 0 deletions
@@ -14,11 +14,13 @@ require (

 require (
 	filippo.io/edwards25519 v1.2.0 // indirect
+	github.com/dlclark/regexp2 v1.10.0 // indirect
 	github.com/dustin/go-humanize v1.0.1 // indirect
 	github.com/google/uuid v1.6.0 // indirect
 	github.com/mattn/go-isatty v0.0.24 // indirect
 	github.com/mattn/go-runewidth v0.0.3 // indirect
 	github.com/ncruces/go-strftime v1.0.0 // indirect
+	github.com/pkoukk/tiktoken-go v0.1.8 // indirect
 	github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // indirect
 	github.com/tidwall/match v1.1.1 // indirect
 	github.com/tidwall/pretty v1.2.1 // indirect
@@ -1,5 +1,7 @@
 filippo.io/edwards25519 v1.2.0 h1:crnVqOiS4jqYleHd9vaKZ+HKtHfllngJIiOpNpoJsjo=
 filippo.io/edwards25519 v1.2.0/go.mod h1:xzAOLCNug/yB62zG1bQ8uziwrIqIuxhctzJT18Q77mc=
+github.com/dlclark/regexp2 v1.10.0 h1:+/GIL799phkJqYW+3YbOd8LCcbHzT0Pbo8zl70MHsq0=
+github.com/dlclark/regexp2 v1.10.0/go.mod h1:DHkYz0B9wPfa6wondMfaivmHpzrQ3v9q8cnmRbL6yW8=
 github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY=
 github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto=
 github.com/expr-lang/expr v1.17.8 h1:W1loDTT+0PQf5YteHSTpju2qfUfNoBt4yw9+wOEU9VM=
@@ -22,6 +24,8 @@ github.com/openai/openai-go v1.12.0 h1:NBQCnXzqOTv5wsgNC36PrFEiskGfO5wccfCWDo9S1
 github.com/openai/openai-go v1.12.0/go.mod h1:g461MYGXEXBVdV5SaR/5tNzNbSfwTBBefwc+LlDCK0Y=
 github.com/peterh/liner v1.2.2 h1:aJ4AOodmL+JxOZZEL2u9iJf8omNRpqHc/EbrK+3mAXw=
 github.com/peterh/liner v1.2.2/go.mod h1:xFwJyiKIXJZUKItq5dGHZSTBRAuG/CpeNpWLyiNRNwI=
+github.com/pkoukk/tiktoken-go v0.1.8 h1:85ENo+3FpWgAACBaEUVp+lctuTcYUO7BtmfhlN/QTRo=
+github.com/pkoukk/tiktoken-go v0.1.8/go.mod h1:9NiV+i9mJKGj1rYOT+njbv+ZwA/zJxYdewGl6qVatpg=
 github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec h1:W09IVJc94icq4NjY3clb7Lk8O1qJ8BdBEF8z0ibU0rE=
 github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec/go.mod h1:qqbHyh8v60DhA7CoWK5oRCqLrMHRGoxYCSS9EjAz6Eo=
 github.com/tidwall/gjson v1.14.2/go.mod h1:/wbyibRr2FHMks5tjHJ5F8dMZh3AcwJEMf5vlfC0lxk=
@@ -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})
@@ -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)
+	}
+}
@@ -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)
@@ -8,6 +8,7 @@ import (
 	"io"
 	"log"
 	"os"
+	"path/filepath"
 	"strings"

 	"github.com/peterh/liner"
@@ -19,6 +20,7 @@ import (
 )

 func main() {
+	os.Setenv("TIKTOKEN_CACHE_DIR", filepath.Join("data", "tiktoken"))
 	cfg, err := config.Load()
 	if err != nil {
 		log.Fatalf("加载配置失败: %v", err)