From ba2ee4b0a00deccba77b01320b60a0c7cb55d4ba Mon Sep 17 00:00:00 2001
From: kevin <kevin@lmve.net>
Date: Fri, 14 Aug 2026 19:10:48 +0800
Subject: [PATCH] =?UTF-8?q?/context=20=E5=A2=9E=E5=8A=A0=20token=20?=
 =?UTF-8?q?=E7=BB=9F=E8=AE=A1=EF=BC=9Atiktoken=20=E7=B2=BE=E7=A1=AE?=
 =?UTF-8?q?=E5=88=86=E8=AF=8D=E4=B8=8E=E7=AA=97=E5=8F=A3=E4=BD=BF=E7=94=A8?=
 =?UTF-8?q?=E7=99=BE=E5=88=86=E6=AF=94?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 go.mod                     |  2 ++
 go.sum                     |  4 +++
 internal/bot/bot.go        | 38 ++++++++++++++++++++++++++
 internal/bot/stats_test.go | 56 ++++++++++++++++++++++++++++++++++++++
 internal/cli/cli.go        | 22 +++++++++++++++
 main.go                    |  2 ++
 6 files changed, 124 insertions(+)
 create mode 100644 internal/bot/stats_test.go

diff --git a/go.mod b/go.mod
index abfb33b..57a9355 100644
--- a/go.mod
+++ b/go.mod
@@ -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
diff --git a/go.sum b/go.sum
index 18a7ab3..b5e9fbc 100644
--- a/go.sum
+++ b/go.sum
@@ -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=
diff --git a/internal/bot/bot.go b/internal/bot/bot.go
index 643e0aa..59424d9 100644
--- a/internal/bot/bot.go
+++ b/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})
diff --git a/internal/bot/stats_test.go b/internal/bot/stats_test.go
new file mode 100644
index 0000000..cb52e38
--- /dev/null
+++ b/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)
+	}
+}
diff --git a/internal/cli/cli.go b/internal/cli/cli.go
index cc01146..9cd2ba4 100644
--- a/internal/cli/cli.go
+++ b/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)
diff --git a/main.go b/main.go
index 5711298..0552d5f 100644
--- a/main.go
+++ b/main.go
@@ -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)