From 755868a33145ffd8a9b6a76bdbcebd1163a3d375 Mon Sep 17 00:00:00 2001 From: chengqian Date: Tue, 8 Sep 2026 15:30:41 +0800 Subject: [PATCH] =?UTF-8?q?feat(relay):=E6=96=B0=E5=A2=9Etoken=E4=BC=B0?= =?UTF-8?q?=E7=AE=97=E6=8E=A5=E5=8F=A3?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- internal/cmd/cmd.go | 1 + internal/handler/relay/relay.go | 32 +++++ relay/handler/count_tokens.go | 221 +++++++++++++++++++++++++++++ relay/handler/count_tokens_test.go | 91 ++++++++++++ 4 files changed, 345 insertions(+) create mode 100644 relay/handler/count_tokens.go create mode 100644 relay/handler/count_tokens_test.go diff --git a/internal/cmd/cmd.go b/internal/cmd/cmd.go index 0eb76be..9a48fcd 100644 --- a/internal/cmd/cmd.go +++ b/internal/cmd/cmd.go @@ -531,6 +531,7 @@ func registerRelayRoutes(server *ghttp.Server) { group.POST("/responses/{id}/cancel", relay.HandleResponsesCancel) group.DELETE("/responses/{id}", relay.HandleResponsesDelete) group.POST("/messages", relay.HandleMessages) + group.POST("/messages/count_tokens", relay.HandleCountTokens) group.POST("/audio/speech", relay.HandleAudioSpeech) group.POST("/audio/transcriptions", relay.HandleAudioTranscription) group.POST("/audio/translations", relay.HandleAudioTranslation) diff --git a/internal/handler/relay/relay.go b/internal/handler/relay/relay.go index e71e18f..de04d37 100644 --- a/internal/handler/relay/relay.go +++ b/internal/handler/relay/relay.go @@ -236,6 +236,38 @@ func HandleMessages(r *ghttp.Request) { recordAudit(r, rc, capture, body, "/messages", latencyMs, firstTokenMs(billingResult)) } +// HandleCountTokens 处理 /v1/messages/count_tokens（Claude 原生格式） +// 本地估算输入 token 数：不请求上游、不计费，仅返回 {"input_tokens": N} 供客户端做上下文窗口管理参考。 +func HandleCountTokens(r *ghttp.Request) { + body := r.GetBody() + if len(body) == 0 { + writeClaudeCountTokensError(r, http.StatusBadRequest, "request body is empty") + return + } + + monitor.RecordBytesIn(len(body)) + + tokens, err := handler.CountClaudeTokens(body) + if err != nil { + writeClaudeCountTokensError(r, http.StatusBadRequest, err.Error()) + return + } + + r.Response.WriteJsonExit(g.Map{"input_tokens": tokens}) +} + +// writeClaudeCountTokensError 以 Claude 原生错误格式回写 count_tokens 的校验错误。 +func writeClaudeCountTokensError(r *ghttp.Request, status int, message string) { + r.Response.WriteHeader(status) + r.Response.WriteJsonExit(g.Map{ + "type": "error", + "error": g.Map{ + "type": "invalid_request_error", + "message": message, + }, + }) +} + // HandleGeminiGenerateContent 处理 /v1beta/models/{model}:generateContent（Gemini 原生格式） func HandleGeminiGenerateContent(r *ghttp.Request) { g.Log().Debugf(r.Context(), "[HandleGeminiGenerateContent] Entry: method=%s, path=%s, uri=%s", diff --git a/relay/handler/count_tokens.go b/relay/handler/count_tokens.go new file mode 100644 index 0000000..5c40387 --- /dev/null +++ b/relay/handler/count_tokens.go @@ -0,0 +1,221 @@ +package handler + +import ( + "encoding/json" + "errors" + "math" + "strings" + + "github.com/qianfree/team-api/relay/dto" +) + +// count_tokens 端点的本地估算参数。 +// +// 说明：本平台不内置精确分词器（tiktoken 等），POST /v1/messages/count_tokens +// 采用「本地粗略估算」策略——不请求上游、不计费，仅供客户端做上下文窗口管理的近似参考， +// 返回值不保证与上游精确计数一致。估算按字符类别加权：CJK 表意文字 token 密度远高于拉丁字母。 +const ( + latinTokenWeight = 0.25 // 拉丁字母/数字/空白/标点：约 4 字符/token + cjkTokenWeight = 0.6 // CJK 表意文字：约 1.67 字符/token + + claudeImageTokens = 1600 // 单张图片/文档的粗略 token 估值（无法本地解码，按固定值计） + claudeMessageOverhead = 3 // 每条消息的结构化开销 + claudeToolOverhead = 8 // 每个工具定义的结构化开销 + claudeRequestOverhead = 8 // 请求整体固定开销（system 包装、消息分隔等） +) + +// errCountTokensInvalid 表示 count_tokens 请求体校验失败，调用方统一映射为 400 invalid_request_error。 +var errCountTokensInvalid = errors.New("invalid count_tokens request") + +// ctRequest count_tokens 的宽松解析结构：content / system / input_schema 保留为 RawMessage， +// 便于区分「字符串」与「内容块数组」两种形态后分别估算。 +type ctRequest struct { + Model string `json:"model"` + System json.RawMessage `json:"system"` + Messages []ctMessage `json:"messages"` + Tools []ctTool `json:"tools"` +} + +type ctMessage struct { + Role string `json:"role"` + Content json.RawMessage `json:"content"` +} + +type ctTool struct { + Name string `json:"name"` + Description string `json:"description"` + InputSchema json.RawMessage `json:"input_schema"` +} + +// CountClaudeTokens 本地估算 Claude Messages 请求的输入 token 数。 +// 返回 (input_tokens, error)；error 非 nil 时均为请求体校验错误。 +func CountClaudeTokens(body []byte) (int, error) { + var req ctRequest + if err := json.Unmarshal(body, &req); err != nil { + return 0, errWithMessage("invalid request body") + } + if strings.TrimSpace(req.Model) == "" { + return 0, errWithMessage("model is required") + } + if len(req.Messages) == 0 { + return 0, errWithMessage("messages is required") + } + + tokens := claudeRequestOverhead + + // system：string 或 []ClaudeContentBlock + tokens += estimateContentTokens(req.System) + + // messages + for _, msg := range req.Messages { + tokens += estimateContentTokens(msg.Content) + tokens += claudeMessageOverhead + } + + // tools：名称 + 描述 + input_schema（JSON 全文计入） + for _, tool := range req.Tools { + tokens += estimateTextTokens(tool.Name) + tokens += estimateTextTokens(tool.Description) + tokens += estimateRawTokens(tool.InputSchema) + tokens += claudeToolOverhead + } + + return tokens, nil +} + +// estimateContentTokens 估算 message.content / system 的 token 数，兼容字符串与内容块数组两种形态。 +func estimateContentTokens(raw json.RawMessage) int { + trimmed := bytesTrimSpace(raw) + if len(trimmed) == 0 || string(trimmed) == "null" { + return 0 + } + + // 形态一：纯字符串 + var s string + if json.Unmarshal(trimmed, &s) == nil { + return estimateTextTokens(s) + } + + // 形态二：内容块数组 + var blocks []dto.ClaudeContentBlock + if json.Unmarshal(trimmed, &blocks) == nil { + total := 0 + for i := range blocks { + total += estimateBlockTokens(&blocks[i]) + } + return total + } + + // 兜底：无法结构化解析时，按原始 JSON 文本估算，避免漏计 + return estimateRawTokens(trimmed) +} + +// estimateBlockTokens 估算单个内容块的 token 数（text / thinking / tool_use / tool_result / image 等）。 +func estimateBlockTokens(b *dto.ClaudeContentBlock) int { + total := 0 + if b.Text != nil { + total += estimateTextTokens(*b.Text) + } + if b.Thinking != nil { + total += estimateTextTokens(*b.Thinking) + } + if b.Name != "" { + total += estimateTextTokens(b.Name) + } + if b.Input != nil { + total += estimateAnyTokens(b.Input) // tool_use 的入参 + } + if b.Content != nil { + total += estimateAnyContent(b.Content) // tool_result 的嵌套内容（string 或 []block） + } + if b.Source != nil { + total += claudeImageTokens // 图片/文档：固定估值 + } + return total +} + +// estimateAnyContent 估算 tool_result.content：可能是字符串，也可能是内容块数组。 +func estimateAnyContent(v any) int { + if v == nil { + return 0 + } + if s, ok := v.(string); ok { + return estimateTextTokens(s) + } + raw, err := json.Marshal(v) + if err != nil { + return 0 + } + return estimateContentTokens(raw) +} + +// estimateAnyTokens 将任意结构（tool_use 入参等）序列化为 JSON 后按文本估算。 +func estimateAnyTokens(v any) int { + if v == nil { + return 0 + } + raw, err := json.Marshal(v) + if err != nil { + return 0 + } + return estimateTextTokens(string(raw)) +} + +// estimateRawTokens 按原始 JSON 字节的文本估算（input_schema 等已是 JSON 的场景）。 +func estimateRawTokens(raw json.RawMessage) int { + if len(raw) == 0 { + return 0 + } + return estimateTextTokens(string(raw)) +} + +// estimateTextTokens 按字符类别粗略估算文本 token 数：CJK 表意文字权重高，其余字符权重低。 +func estimateTextTokens(text string) int { + if text == "" { + return 0 + } + var weighted float64 + for _, r := range text { + if isCJK(r) { + weighted += cjkTokenWeight + } else { + weighted += latinTokenWeight + } + } + return int(math.Ceil(weighted)) +} + +// isCJK 判断字符是否属于常见 CJK 表意文字区段（中日韩统一表意文字及扩展、假名、谚文等）。 +func isCJK(r rune) bool { + switch { + case r >= 0x4E00 && r <= 0x9FFF, // CJK 统一表意文字 + r >= 0x3400 && r <= 0x4DBF, // CJK 扩展 A + r >= 0x20000 && r <= 0x2A6DF, // CJK 扩展 B + r >= 0xF900 && r <= 0xFAFF, // CJK 兼容表意文字 + r >= 0x3040 && r <= 0x30FF, // 平假名 + 片假名 + r >= 0xAC00 && r <= 0xD7AF, // 谚文音节 + r >= 0x3000 && r <= 0x303F: // CJK 符号和标点 + return true + } + return false +} + +// bytesTrimSpace 去除 RawMessage 两端的 JSON 空白字符。 +func bytesTrimSpace(b []byte) []byte { + return []byte(strings.TrimSpace(string(b))) +} + +// errWithMessage 构造 count_tokens 校验错误，消息文本即返回给客户端的错误信息。 +func errWithMessage(msg string) error { + return &countTokensError{message: msg} +} + +// countTokensError 携带面向客户端的错误信息，供 handler 以 Claude 原生格式回写。 +type countTokensError struct { + message string +} + +func (e *countTokensError) Error() string { return e.message } +func (e *countTokensError) Is(target error) bool { + return target == errCountTokensInvalid +} diff --git a/relay/handler/count_tokens_test.go b/relay/handler/count_tokens_test.go new file mode 100644 index 0000000..81c9db5 --- /dev/null +++ b/relay/handler/count_tokens_test.go @@ -0,0 +1,91 @@ +package handler + +import ( + "errors" + "testing" +) + +func TestCountClaudeTokens_Basic(t *testing.T) { + body := []byte(`{ + "model": "claude-3-5-sonnet", + "messages": [{"role": "user", "content": "count this prompt please"}], + "tools": [{"name": "lookup", "description": "Look up a value", "input_schema": {"type": "object", "properties": {"query": {"type": "string"}}}}] + }`) + tokens, err := CountClaudeTokens(body) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if tokens <= 0 { + t.Fatalf("expected positive input_tokens, got %d", tokens) + } +} + +func TestCountClaudeTokens_MissingMessages(t *testing.T) { + body := []byte(`{"model": "claude-3-5-sonnet"}`) + _, err := CountClaudeTokens(body) + if err == nil { + t.Fatal("expected error for missing messages") + } + if !errors.Is(err, errCountTokensInvalid) { + t.Fatalf("expected errCountTokensInvalid, got %v", err) + } + if err.Error() != "messages is required" { + t.Fatalf("unexpected message: %q", err.Error()) + } +} + +func TestCountClaudeTokens_MissingModel(t *testing.T) { + body := []byte(`{"messages": [{"role": "user", "content": "hi"}]}`) + _, err := CountClaudeTokens(body) + if err == nil || err.Error() != "model is required" { + t.Fatalf("expected model required error, got %v", err) + } +} + +func TestCountClaudeTokens_InvalidBody(t *testing.T) { + _, err := CountClaudeTokens([]byte(`not json`)) + if err == nil || !errors.Is(err, errCountTokensInvalid) { + t.Fatalf("expected invalid request error, got %v", err) + } +} + +func TestCountClaudeTokens_ContentBlocks(t *testing.T) { + // 内容块数组（含 text + image）应比同长纯文本计入更多 token（图片固定估值） + body := []byte(`{ + "model": "claude-3-5-sonnet", + "messages": [{"role": "user", "content": [ + {"type": "text", "text": "describe this"}, + {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "aGVsbG8="}} + ]}] + }`) + tokens, err := CountClaudeTokens(body) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if tokens < claudeImageTokens { + t.Fatalf("expected image tokens to dominate, got %d", tokens) + } +} + +func TestCountClaudeTokens_SystemAndCJK(t *testing.T) { + // system 文本应计入；CJK 文本 token 密度高于等长拉丁文本 + latin := []byte(`{"model":"m","system":"aaaaaaaaaa","messages":[{"role":"user","content":"x"}]}`) + cjk := []byte(`{"model":"m","system":"你好世界你好世界你好","messages":[{"role":"user","content":"x"}]}`) + lt, err := CountClaudeTokens(latin) + if err != nil { + t.Fatalf("latin err: %v", err) + } + ct, err := CountClaudeTokens(cjk) + if err != nil { + t.Fatalf("cjk err: %v", err) + } + if ct <= lt { + t.Fatalf("expected CJK (%d) to exceed latin (%d) for equal char count", ct, lt) + } +} + +func TestEstimateTextTokens_Empty(t *testing.T) { + if got := estimateTextTokens(""); got != 0 { + t.Fatalf("expected 0 for empty string, got %d", got) + } +}