From bb28f2ccad3f6286a9dc4a0ad9e20f571f57533b Mon Sep 17 00:00:00 2001 From: CaIon Date: Wed, 7 Oct 2026 13:52:05 +0800 Subject: [PATCH] feat(tokenkit): move token counting into a module calibrated on agent traffic Token estimates badly undercounted agent requests. Claude tool definitions, Responses function calls and their outputs, Chat assistant tool calls, and Gemini system instructions, function calls, function responses, and tools never reached the counter; Chat tool schemas were counted as Go map syntax and Claude tool results as escaped JSON. Measured against upstream usage on 54 real Claude Code and Codex conversation windows, native prompt estimates were 31-93% low, and the Claude and Gemini completion fallbacks about 90% low because streamed tool calls were not counted. All of these now reach the counter. The completion fallbacks count tool names and arguments, plus 40 tokens per Claude tool_use block. The estimator weights are refitted per tokenizer: Claude before Opus 4.7, Claude 4.7 and later (about 30% more tokens for the same text), GPT, and Gemini. Prompt estimates now land within 1% bias and 3-4% mean error for each family. Token counting moves into tokenkit, an independent module like relaykit, so other gateways can use it. It also replaces the flat 520-token image estimate with each provider's sizing rules, checked against upstream usage on sizes outside the test table: OpenAI patch and tile rules per model and detail level (gpt-5.5 and later default to original detail, and their high detail is billed as measured, not as documented), Claude's two resolution tiers, and Gemini 3 media resolution and Gemini 2.x tiles. Chat image_url parts without detail no longer default to high in ParseContent, so they estimate as the upstream default auto; no converter sent that value upstream. tiktoken-go moves to v0.8.1, about twice as fast on ordinary text. It requires Go 1.26, so the root module and tokenkit move to Go 1.26; relaykit stays on 1.25.1. In every version BPE cost grows with the square of a pre-tokenized piece, so 100,000 Chinese characters without punctuation took 21.6 s to count during pre-consume. Count now cuts single-class runs longer than 1024 bytes, which leaves ordinary text unchanged and brings that input down to 106 ms. The encoder cache no longer grows with client-supplied model names. --- .github/workflows/ci.yml | 11 +- .github/workflows/electron-build.yml | 2 +- .github/workflows/release.yml | 6 +- AGENTS.md | 9 +- Dockerfile | 5 +- Dockerfile.dev | 5 +- common/model.go | 17 - go.mod | 9 +- go.sum | 8 +- main.go | 2 - makefile | 2 + relay/channel/claude/relay-claude.go | 3 + relay/channel/claude/relay_claude_test.go | 29 ++ relay/channel/gemini/relay-gemini.go | 10 +- .../channel/gemini/relay_gemini_usage_test.go | 103 ++++--- relay/channel/openai/relay-openai.go | 3 +- relaykit/dto/claude.go | 34 +- relaykit/dto/gemini.go | 53 +++- relaykit/dto/openai_request.go | 59 +++- relaykit/dto/token_count_meta_test.go | 100 ++++++ .../claude_messages/to_oai_chat_resp.go | 10 + service/responses_usage.go | 3 +- service/responses_usage_test.go | 5 +- service/token_counter.go | 192 ++---------- service/token_estimator.go | 230 -------------- service/tokenizer.go | 63 ---- service/usage_helpr.go | 3 +- tokenkit/README.md | 125 ++++++++ tokenkit/go.mod | 15 + tokenkit/go.sum | 14 + tokenkit/image.go | 290 ++++++++++++++++++ tokenkit/image_test.go | 124 ++++++++ tokenkit/tokenizer.go | 102 ++++++ tokenkit/tokenkit.go | 245 +++++++++++++++ tokenkit/tokenkit_test.go | 83 +++++ 35 files changed, 1404 insertions(+), 570 deletions(-) create mode 100644 relaykit/dto/token_count_meta_test.go delete mode 100644 service/token_estimator.go delete mode 100644 service/tokenizer.go create mode 100644 tokenkit/README.md create mode 100644 tokenkit/go.mod create mode 100644 tokenkit/go.sum create mode 100644 tokenkit/image.go create mode 100644 tokenkit/image_test.go create mode 100644 tokenkit/tokenizer.go create mode 100644 tokenkit/tokenkit.go create mode 100644 tokenkit/tokenkit_test.go diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index b5578dfb1c..ca34072248 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -34,6 +34,7 @@ jobs: cache-dependency-path: | go.sum relaykit/go.sum + tokenkit/go.sum - name: Create frontend embed placeholder run: | @@ -47,6 +48,10 @@ jobs: working-directory: relaykit run: go vet ./... + - name: Vet tokenkit independently + working-directory: tokenkit + run: go vet ./... + - name: Build root module run: go build ./... @@ -54,7 +59,11 @@ jobs: working-directory: relaykit run: go build ./... - - name: Test root and relaykit modules + - name: Build tokenkit independently + working-directory: tokenkit + run: go build ./... + + - name: Test root, relaykit, and tokenkit modules run: make test frontend: diff --git a/.github/workflows/electron-build.yml b/.github/workflows/electron-build.yml index f704174d7f..5df7becce6 100644 --- a/.github/workflows/electron-build.yml +++ b/.github/workflows/electron-build.yml @@ -39,7 +39,7 @@ jobs: - name: Setup Go uses: actions/setup-go@924ae3a1cded613372ab5595356fb5720e22ba16 # v6.5.0 with: - go-version: '>=1.25.1' + go-version: '>=1.26' - name: Build frontend env: diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index ecd5d24f07..fc2b911d1f 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -44,7 +44,7 @@ jobs: - name: Set up Go uses: actions/setup-go@924ae3a1cded613372ab5595356fb5720e22ba16 # v6.5.0 with: - go-version: '>=1.25.1' + go-version: '>=1.26' - name: Build Backend (amd64) run: | go mod download @@ -98,7 +98,7 @@ jobs: - name: Set up Go uses: actions/setup-go@924ae3a1cded613372ab5595356fb5720e22ba16 # v6.5.0 with: - go-version: '>=1.25.1' + go-version: '>=1.26' - name: Build Backend run: | go mod download @@ -149,7 +149,7 @@ jobs: - name: Set up Go uses: actions/setup-go@924ae3a1cded613372ab5595356fb5720e22ba16 # v6.5.0 with: - go-version: '>=1.25.1' + go-version: '>=1.26' - name: Build Backend run: | go mod download diff --git a/AGENTS.md b/AGENTS.md index c5a9793a32..4b6ce99aac 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -8,7 +8,7 @@ This is an AI API gateway/proxy built with Go. It aggregates 40+ upstream AI pro ## Tech Stack -- **Backend**: Go 1.25.1 (see each module’s `go.mod`), Gin web framework, GORM v2 ORM +- **Backend**: Go 1.26 for the root module and `tokenkit/`, Go 1.25.1 for `relaykit/` (see each module’s `go.mod`), Gin web framework, GORM v2 ORM - **Frontend**: React 19, TypeScript, Rsbuild 2, TanStack Router/Query/Table, Zustand, Base UI, Tailwind CSS 4 - **Databases**: SQLite, MySQL, PostgreSQL for the primary database (all three must be supported); a separately configured log database also supports ClickHouse - **Cache**: Redis (go-redis) + in-memory cache @@ -20,6 +20,7 @@ This is an AI API gateway/proxy built with Go. It aggregates 40+ upstream AI pro - The Go gateway handles management APIs, upstream relay, billing, and background tasks across `router/`, `middleware/`, `controller/`, `service/`, `model/`, and `relay/`. - `relaykit/` is an independent Go module for protocol DTOs and conversions; transport, authentication, database access, and billing stay in the host. +- `tokenkit/` is an independent Go module for text and image token counting and estimation; request parsing, file loading, and billing stay in the host. - JavaScript task plugins live in `plugins/tasks/`, run through `pkg/jsplugin/`, and integrate with host task polling and settlement. - `web/` is the React frontend (see `web/AGENTS.md`); `electron/` is the desktop wrapper. @@ -72,10 +73,10 @@ This is an AI API gateway/proxy built with Go. It aggregates 40+ upstream AI pro - Remove ineffective `omitempty` tags on non-pointer struct fields only after confirming the active JSON encoder preserves the same output. Do not change field types or omission behavior as part of a style cleanup; optional relay scalar fields must still follow the pointer rules below. - Format modified Go files with `gofmt` and remove unused imports after these changes. -**relaykit module independence:** The `relaykit/` Go module MUST remain independently buildable. +**relaykit and tokenkit module independence:** The `relaykit/` and `tokenkit/` Go modules MUST remain independently buildable. -- Code under `relaykit/` MUST NOT import or depend on packages from the root `new-api` module, or rely on root-only configuration, generated files, or workspace wiring. -- Any change affecting `relaykit/` or its public APIs MUST be verified with `cd relaykit && GOWORK=off go build ./...`; a successful root-module build is not sufficient. +- Code under `relaykit/` or `tokenkit/` MUST NOT import or depend on packages from the root `new-api` module, or rely on root-only configuration, generated files, or workspace wiring. `tokenkit/` also MUST NOT depend on `relaykit/`. +- Any change affecting `relaykit/`, `tokenkit/`, or their public APIs MUST be verified with `GOWORK=off go build ./...` inside that module; a successful root-module build is not sufficient. **JSON package:** In the root Go module, all JSON marshal/unmarshal operations MUST use the wrapper functions in `common/json.go`: diff --git a/Dockerfile b/Dockerfile index 399a2ab85d..7c06a42a7c 100644 --- a/Dockerfile +++ b/Dockerfile @@ -18,9 +18,10 @@ ENV GOEXPERIMENT=greenteagc WORKDIR /build ADD go.mod go.sum ./ -# relaykit is a local submodule referenced via replace; its go.mod must be -# present for go mod download to resolve the main module graph. +# relaykit and tokenkit are local submodules referenced via replace; their +# go.mod files must be present for go mod download to resolve the main module graph. ADD relaykit/go.mod ./relaykit/go.mod +ADD tokenkit/go.mod ./tokenkit/go.mod RUN go mod download COPY . . diff --git a/Dockerfile.dev b/Dockerfile.dev index 47475e34fc..e58bf79bc4 100644 --- a/Dockerfile.dev +++ b/Dockerfile.dev @@ -12,9 +12,10 @@ ENV GOEXPERIMENT=greenteagc WORKDIR /build ADD go.mod go.sum ./ -# relaykit is a local submodule referenced via replace; its go.mod must be -# present for go mod download to resolve the main module graph. +# relaykit and tokenkit are local submodules referenced via replace; their +# go.mod files must be present for go mod download to resolve the main module graph. ADD relaykit/go.mod ./relaykit/go.mod +ADD tokenkit/go.mod ./tokenkit/go.mod RUN go mod download COPY . . diff --git a/common/model.go b/common/model.go index 8a8f3202a3..0f4f117885 100644 --- a/common/model.go +++ b/common/model.go @@ -30,13 +30,6 @@ var ( "flux-", "flux.1-", } - OpenAITextModels = []string{ - "gpt-", - "o1", - "o3", - "o4", - "chatgpt", - } ) func IsOpenAIResponseOnlyModel(modelName string) bool { @@ -63,13 +56,3 @@ func IsImageGenerationModel(modelName string) bool { } return false } - -func IsOpenAITextModel(modelName string) bool { - modelName = strings.ToLower(modelName) - for _, m := range OpenAITextModels { - if strings.Contains(modelName, m) { - return true - } - } - return false -} diff --git a/go.mod b/go.mod index f2dd13d392..9d69d7ccd1 100644 --- a/go.mod +++ b/go.mod @@ -1,7 +1,7 @@ module github.com/QuantumNous/new-api // +heroku goVersion go1.18 -go 1.25.1 +go 1.26 require ( github.com/Calcium-Ion/go-epay v0.0.5 @@ -47,7 +47,7 @@ require ( github.com/thanhpk/randstr v1.0.6 github.com/tidwall/gjson v1.19.0 github.com/tidwall/sjson v1.2.5 - github.com/tiktoken-go/tokenizer v0.6.2 + github.com/tiktoken-go/tokenizer v0.8.1 // indirect github.com/waffo-com/waffo-go v1.3.2 github.com/yapingcat/gomedia v0.0.0-20240906162731-17feea57090c golang.org/x/crypto v0.52.0 @@ -72,6 +72,7 @@ require ( github.com/ClickHouse/ch-go v0.71.0 // indirect github.com/bmatcuk/doublestar/v4 v4.6.1 // indirect github.com/casbin/govaluate v1.10.0 // indirect + github.com/dlclark/regexp2/v2 v2.5.1 // indirect github.com/go-faster/city v1.0.1 // indirect github.com/go-faster/errors v0.7.1 // indirect github.com/go-jose/go-jose/v4 v4.1.4 // indirect @@ -105,7 +106,6 @@ require ( github.com/cespare/xxhash/v2 v2.3.0 // indirect github.com/cloudwego/base64x v0.1.6 // indirect github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect - github.com/dlclark/regexp2 v1.11.5 // indirect github.com/dustin/go-humanize v1.0.1 // indirect github.com/expr-lang/expr v1.17.8 github.com/fxamacker/cbor/v2 v2.9.0 @@ -169,6 +169,9 @@ require ( github.com/Calcium-Ion/moejs v0.1.0-alpha.5 github.com/ClickHouse/clickhouse-go/v2 v2.46.0 github.com/QuantumNous/new-api/relaykit v0.0.0 + github.com/QuantumNous/new-api/tokenkit v0.0.0 ) replace github.com/QuantumNous/new-api/relaykit => ./relaykit + +replace github.com/QuantumNous/new-api/tokenkit => ./tokenkit diff --git a/go.sum b/go.sum index 73b2edd4cf..5a1a1310d6 100644 --- a/go.sum +++ b/go.sum @@ -1012,8 +1012,8 @@ github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f h1:lO4WD4F/r github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f/go.mod h1:cuUVRXasLTGF7a8hSLbxyZXjz+1KgoB3wDUb6vlszIc= github.com/dgryski/go-sip13 v0.0.0-20181026042036-e10d5fee7954/go.mod h1:vAd38F8PWV+bWy6jNmig1y/TA+kYO4g3RSRF0IAv0no= github.com/distribution/distribution/v3 v3.0.0-20220526142353-ffbd94cbe269/go.mod h1:28YO/VJk9/64+sTGNuYaBjWxrXTPrj0C0XmgTIOjxX4= -github.com/dlclark/regexp2 v1.11.5 h1:Q/sSnsKerHeCkc/jSTNq1oCm7KiVgUMZRDUoRu0JQZQ= -github.com/dlclark/regexp2 v1.11.5/go.mod h1:DHkYz0B9wPfa6wondMfaivmHpzrQ3v9q8cnmRbL6yW8= +github.com/dlclark/regexp2/v2 v2.5.1 h1:E5Ug7Dh264W1ymdySmiHNcDG7fmsR307APCE5R07a20= +github.com/dlclark/regexp2/v2 v2.5.1/go.mod h1:avUrQvPaLz2DrFNHJF0taWAFFX2C1GMSSoeiqFjcBmU= github.com/dmarkham/enumer v1.5.8/go.mod h1:d10o8R3t/gROm2p3BXqTkMt2+HMuxEmWCXzorAruYak= github.com/dnaeon/go-vcr v1.0.1/go.mod h1:aBB1+wY4s93YsC3HHjMBMrwTj2R9FHDzUr9KyGc8n1E= github.com/docker/cli v0.0.0-20191017083524-a8ff7f821017/go.mod h1:JLrzqnKDaYBop7H2jaqPtU4hHvMKP+vjCwu2uszcLI8= @@ -1962,8 +1962,8 @@ github.com/tidwall/pretty v1.2.1 h1:qjsOFOWWQl+N3RsoF5/ssm1pHmJJwhjlSbZ51I6wMl4= github.com/tidwall/pretty v1.2.1/go.mod h1:ITEVvHYasfjBbM0u2Pg8T2nJnzm8xPwvNhhsoaGGjNU= github.com/tidwall/sjson v1.2.5 h1:kLy8mja+1c9jlljvWTlSazM7cKDRfJuR/bOJhcY5NcY= github.com/tidwall/sjson v1.2.5/go.mod h1:Fvgq9kS/6ociJEDnK0Fk1cpYF4FIW6ZF7LAe+6jwd28= -github.com/tiktoken-go/tokenizer v0.6.2 h1:t0GN2DvcUZSFWT/62YOgoqb10y7gSXBGs0A+4VCQK+g= -github.com/tiktoken-go/tokenizer v0.6.2/go.mod h1:6UCYI/DtOallbmL7sSy30p6YQv60qNyU/4aVigPOx6w= +github.com/tiktoken-go/tokenizer v0.8.1 h1:4obDoB6/dhdBt9xMweX4nww5cjdOq/nYF4ecwPq2+mg= +github.com/tiktoken-go/tokenizer v0.8.1/go.mod h1:eLA0t6nGvn9mDc7gt90qt7pMat+gE9ViqwQ6l9B+tA4= github.com/tklauser/go-sysconf v0.3.12 h1:0QaGUFOdQaIVdPgfITYzaTegZvdCjmYO52cSFAEVmqU= github.com/tklauser/go-sysconf v0.3.12/go.mod h1:Ho14jnntGE1fpdOqQEEaiKRpvIavV0hSfmBq8nJbHYI= github.com/tklauser/numcpus v0.6.1 h1:ng9scYS7az0Bk4OZLvrNXNSAO2Pxr1XXRAPyjhIx+Fk= diff --git a/main.go b/main.go index 26b59f9d7c..998e00e71a 100644 --- a/main.go +++ b/main.go @@ -308,8 +308,6 @@ func InitResources() error { service.InitHttpClient() - service.InitTokenEncoders() - // Initialize SQL Database err = model.InitDB() if err != nil { diff --git a/makefile b/makefile index 24048b9136..73c4845378 100644 --- a/makefile +++ b/makefile @@ -47,6 +47,8 @@ test: GOWORK=off go test $$root_packages @echo "Testing relaykit Go module..." @cd relaykit && GOWORK=off go test ./... + @echo "Testing tokenkit Go module..." + @cd tokenkit && GOWORK=off go test ./... reset-setup: @echo "Resetting local setup wizard state..." diff --git a/relay/channel/claude/relay-claude.go b/relay/channel/claude/relay-claude.go index f8fb1611b3..f44611fb6c 100644 --- a/relay/channel/claude/relay-claude.go +++ b/relay/channel/claude/relay-claude.go @@ -248,6 +248,9 @@ func HandleStreamFinalResponse(c *gin.Context, info *relaycommon.RelayInfo, clau } // 只补缺失字段，不整份覆盖——保留 message_start 已拿到的 cache 字段 fallback := service.ResponseText2Usage(c, claudeInfo.ResponseText.String(), info.UpstreamModelName, info.GetEstimatePromptTokens()) + // Each tool_use block costs about 40 output tokens beyond its name and JSON + // input (measured on Claude Haiku 4.5 and Sonnet 5 agent turns). + fallback.CompletionTokens += claudeInfo.ToolUseCount * 40 if claudeInfo.Usage.CompletionTokens == 0 || (!claudeInfo.Done && fallback.CompletionTokens > claudeInfo.Usage.CompletionTokens) { claudeInfo.Usage.CompletionTokens = fallback.CompletionTokens diff --git a/relay/channel/claude/relay_claude_test.go b/relay/channel/claude/relay_claude_test.go index 8b7c4d0d00..632b3d3d81 100644 --- a/relay/channel/claude/relay_claude_test.go +++ b/relay/channel/claude/relay_claude_test.go @@ -10,6 +10,8 @@ import ( "github.com/QuantumNous/new-api/relay/helper" "github.com/QuantumNous/new-api/relaykit/dto" "github.com/QuantumNous/new-api/relaykit/relayconvert" + "github.com/QuantumNous/new-api/relaykit/types" + "github.com/QuantumNous/new-api/tokenkit" "github.com/gin-gonic/gin" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -228,6 +230,33 @@ func TestFormatClaudeResponseInfo_ContentBlockDelta(t *testing.T) { } } +// A stream cut off during an agent tool call must still estimate the call's +// output tokens when the final usage never arrives. +func TestClaudeStreamFallbackCountsInterruptedToolUse(t *testing.T) { + claudeInfo := &ClaudeResponseInfo{Usage: &dto.Usage{}} + events := []*dto.ClaudeResponse{ + {Type: "content_block_start", ContentBlock: &dto.ClaudeMediaMessage{Type: "tool_use", Id: "toolu_1", Name: "read_file"}}, + {Type: "content_block_delta", Delta: &dto.ClaudeMediaMessage{Type: "input_json_delta", PartialJson: commonPointer(`{"path":`)}}, + {Type: "content_block_delta", Delta: &dto.ClaudeMediaMessage{Type: "input_json_delta", PartialJson: commonPointer(`"main.go"}`)}}, + } + for _, event := range events { + require.True(t, FormatClaudeResponseInfo(event, nil, claudeInfo)) + } + require.Equal(t, `read_file{"path":"main.go"}`, claudeInfo.ResponseText.String()) + + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + info := &relaycommon.RelayInfo{ + RelayFormat: types.RelayFormatClaude, + ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "claude-sonnet-5"}, + } + info.SetEstimatePromptTokens(100) + HandleStreamFinalResponse(c, info, claudeInfo) + + // Text estimate plus the per-tool_use framing the upstream bills. + assert.Equal(t, tokenkit.Estimate("claude-sonnet-5", `read_file{"path":"main.go"}`)+40, claudeInfo.Usage.CompletionTokens) + assert.Equal(t, 100, claudeInfo.Usage.PromptTokens) +} + func TestBuildOpenAIStyleUsageFromClaudeUsage(t *testing.T) { usage := &dto.Usage{ PromptTokens: 100, diff --git a/relay/channel/gemini/relay-gemini.go b/relay/channel/gemini/relay-gemini.go index f41f480976..ce7711d6c0 100644 --- a/relay/channel/gemini/relay-gemini.go +++ b/relay/channel/gemini/relay-gemini.go @@ -75,6 +75,12 @@ func geminiResponseUsageText(response *dto.GeminiChatResponse) string { if part.Text != "" { text.WriteString(part.Text) } + // Function calls are most of an agent turn's output. + if part.FunctionCall != nil { + text.WriteString(part.FunctionCall.FunctionName) + args, _ := common.Marshal(part.FunctionCall.Arguments) + text.Write(args) + } } } return text.String() @@ -206,11 +212,9 @@ func geminiStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http if part.InlineData != nil && part.InlineData.MimeType != "" { imageCount++ } - if part.Text != "" { - responseText.WriteString(part.Text) - } } } + responseText.WriteString(geminiResponseUsageText(&geminiResponse)) // 更新使用量统计 if metadata := geminiResponse.GetUsageMetadata(); dto.HasGeminiUsageMetadataTokens(metadata) { diff --git a/relay/channel/gemini/relay_gemini_usage_test.go b/relay/channel/gemini/relay_gemini_usage_test.go index fe284597e9..ee0db81df2 100644 --- a/relay/channel/gemini/relay_gemini_usage_test.go +++ b/relay/channel/gemini/relay_gemini_usage_test.go @@ -13,6 +13,7 @@ import ( relaycommon "github.com/QuantumNous/new-api/relay/common" "github.com/QuantumNous/new-api/relaykit/dto" "github.com/QuantumNous/new-api/relaykit/types" + "github.com/QuantumNous/new-api/tokenkit" "github.com/gin-gonic/gin" "github.com/stretchr/testify/require" ) @@ -494,8 +495,6 @@ func TestGeminiChatHandlerMissingUsageMetadataBuildsEstimatedBillingUsage(t *tes func TestGeminiStreamHandlerPromptOnlyUsageMetadataEstimatesCompletionTokens(t *testing.T) { gin.SetMode(gin.TestMode) - c, _ := gin.CreateTestContext(httptest.NewRecorder()) - c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil) oldStreamingTimeout := constant.StreamingTimeout constant.StreamingTimeout = 300 @@ -503,53 +502,67 @@ func TestGeminiStreamHandlerPromptOnlyUsageMetadataEstimatesCompletionTokens(t * constant.StreamingTimeout = oldStreamingTimeout }) - info := &relaycommon.RelayInfo{ - OriginModelName: "gemini-3-flash-preview", - ChannelMeta: &relaycommon.ChannelMeta{ - UpstreamModelName: "gemini-3-flash-preview", - }, - } - info.SetEstimatePromptTokens(20) + for _, tc := range []struct { + name string + part dto.GeminiPart + usageText string + }{ + {"text", dto.GeminiPart{Text: "partial streamed answer before disconnect"}, "partial streamed answer before disconnect"}, + // An agent turn that only calls a tool still produced output tokens. + {"function call", dto.GeminiPart{FunctionCall: &dto.FunctionCall{FunctionName: "read_file", Arguments: map[string]any{"path": "main.go"}}}, `read_file{"path":"main.go"}`}, + } { + t.Run(tc.name, func(t *testing.T) { + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil) - // Simulates a client aborting the stream before the final chunk: text was - // streamed but the last observed usageMetadata only carries prompt tokens. - chunk := dto.GeminiChatResponse{ - Candidates: []dto.GeminiChatCandidate{ - { - Content: dto.GeminiChatContent{ - Role: "model", - Parts: []dto.GeminiPart{ - {Text: "partial streamed answer before disconnect"}, + info := &relaycommon.RelayInfo{ + OriginModelName: "gemini-3-flash-preview", + ChannelMeta: &relaycommon.ChannelMeta{ + UpstreamModelName: "gemini-3-flash-preview", + }, + } + info.SetEstimatePromptTokens(20) + + // Simulates a client aborting the stream before the final chunk: output was + // streamed but the last observed usageMetadata only carries prompt tokens. + chunk := dto.GeminiChatResponse{ + Candidates: []dto.GeminiChatCandidate{ + { + Content: dto.GeminiChatContent{ + Role: "model", + Parts: []dto.GeminiPart{tc.part}, + }, }, }, - }, - }, - UsageMetadata: dto.GeminiUsageMetadata{ - PromptTokenCount: 151, - TotalTokenCount: 151, - }, + UsageMetadata: dto.GeminiUsageMetadata{ + PromptTokenCount: 151, + TotalTokenCount: 151, + }, + } + + chunkData, err := common.Marshal(chunk) + require.NoError(t, err) + + streamBody := []byte("data: " + string(chunkData) + "\n" + "data: [DONE]\n") + resp := &http.Response{ + Body: io.NopCloser(bytes.NewReader(streamBody)), + } + + usage, newAPIError := geminiStreamHandler(c, info, resp, func(_ string, _ *dto.GeminiChatResponse) bool { + return true + }) + require.Nil(t, newAPIError) + require.NotNil(t, usage) + require.Equal(t, 151, usage.PromptTokens) + require.Equal(t, tokenkit.Estimate("gemini-3-flash-preview", tc.usageText), usage.CompletionTokens) + require.Greater(t, usage.CompletionTokens, 0) + require.Equal(t, usage.PromptTokens+usage.CompletionTokens, usage.TotalTokens) + require.NotNil(t, usage.BillingUsage) + require.True(t, usage.BillingUsage.Estimated) + require.NotNil(t, usage.BillingUsage.GeminiUsageMetadata) + require.Equal(t, usage.CompletionTokens, usage.BillingUsage.GeminiUsageMetadata.CandidatesTokenCount) + }) } - - chunkData, err := common.Marshal(chunk) - require.NoError(t, err) - - streamBody := []byte("data: " + string(chunkData) + "\n" + "data: [DONE]\n") - resp := &http.Response{ - Body: io.NopCloser(bytes.NewReader(streamBody)), - } - - usage, newAPIError := geminiStreamHandler(c, info, resp, func(_ string, _ *dto.GeminiChatResponse) bool { - return true - }) - require.Nil(t, newAPIError) - require.NotNil(t, usage) - require.Equal(t, 151, usage.PromptTokens) - require.Greater(t, usage.CompletionTokens, 0) - require.Equal(t, usage.PromptTokens+usage.CompletionTokens, usage.TotalTokens) - require.NotNil(t, usage.BillingUsage) - require.True(t, usage.BillingUsage.Estimated) - require.NotNil(t, usage.BillingUsage.GeminiUsageMetadata) - require.Equal(t, usage.CompletionTokens, usage.BillingUsage.GeminiUsageMetadata.CandidatesTokenCount) } func TestGeminiChatHandlerPromptOnlyUsageMetadataEstimatesCompletionTokens(t *testing.T) { diff --git a/relay/channel/openai/relay-openai.go b/relay/channel/openai/relay-openai.go index 25429fc205..bc95788557 100644 --- a/relay/channel/openai/relay-openai.go +++ b/relay/channel/openai/relay-openai.go @@ -15,6 +15,7 @@ import ( "github.com/QuantumNous/new-api/relaykit/dto" "github.com/QuantumNous/new-api/relaykit/types" "github.com/QuantumNous/new-api/service" + "github.com/QuantumNous/new-api/tokenkit" "github.com/gin-gonic/gin" ) @@ -307,7 +308,7 @@ func OpenaiHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Respo completionTokens := simpleResponse.Usage.CompletionTokens if completionTokens == 0 { for _, choice := range simpleResponse.Choices { - ctkm := service.CountTextToken(choice.Message.StringContent()+choice.Message.GetReasoningContent(), info.UpstreamModelName) + ctkm := tokenkit.Count(info.UpstreamModelName, choice.Message.StringContent()+choice.Message.GetReasoningContent()) completionTokens += ctkm } } diff --git a/relaykit/dto/claude.go b/relaykit/dto/claude.go index 9195d8d85b..fb7e11574e 100644 --- a/relaykit/dto/claude.go +++ b/relaykit/dto/claude.go @@ -336,9 +336,25 @@ func (c *ClaudeRequest) GetTokenCountMeta() *types.TokenCountMeta { texts = append(texts, string(b)) } case "tool_result": - if media.Content != nil { - b, _ := kitutil.Marshal(media.Content) - texts = append(texts, string(b)) + // Count tool output as the text the model reads; JSON-encoding it + // would escape newlines and quotes and count image base64 as text. + if media.IsStringContent() { + texts = append(texts, media.GetStringContent()) + break + } + for _, part := range media.ParseMediaContent() { + if source := part.ToFileSource(); source != nil { + fileType := types.FileTypeFile + if part.Type == "image" { + fileType = types.FileTypeImage + } + fileMeta = append(fileMeta, &types.FileMeta{FileType: fileType, Source: source}) + } else if part.Type == "text" { + texts = append(texts, part.GetText()) + } else { + b, _ := kitutil.Marshal(part) + texts = append(texts, string(b)) + } } } } @@ -448,6 +464,18 @@ func ProcessTools(tools []any) ([]*Tool, []*ClaudeWebSearchTool) { normalTools = append(normalTools, &t) case ClaudeWebSearchTool: webSearchTools = append(webSearchTools, &t) + case map[string]any: + // Tools decoded from client JSON arrive as maps. + toolType, _ := t["type"].(string) + if strings.HasPrefix(toolType, "web_search") { + if webSearchTool, err := kitutil.Any2Type[ClaudeWebSearchTool](t); err == nil { + webSearchTools = append(webSearchTools, &webSearchTool) + } + continue + } + if normalTool, err := kitutil.Any2Type[Tool](t); err == nil { + normalTools = append(normalTools, &normalTool) + } default: // 未知类型，跳过 continue diff --git a/relaykit/dto/gemini.go b/relaykit/dto/gemini.go index 4997645a1f..e40877fcc3 100644 --- a/relaykit/dto/gemini.go +++ b/relaykit/dto/gemini.go @@ -2,7 +2,9 @@ package dto import ( "encoding/json" + "maps" "net/http" + "slices" "strings" kitutil "github.com/QuantumNous/new-api/relaykit/relayconvert/kitutil" @@ -74,13 +76,47 @@ func (r *GeminiChatRequest) GetTokenCountMeta() *types.TokenCountMeta { } var inputTexts []string + if r.SystemInstructions != nil { + for _, part := range r.SystemInstructions.Parts { + if part.Text != "" { + inputTexts = append(inputTexts, part.Text) + } + } + } for _, content := range r.Contents { for _, part := range content.Parts { if part.Text != "" { inputTexts = append(inputTexts, part.Text) } - if source := part.InlineData.ToFileSource(); source != nil { - mimeType := part.InlineData.MimeType + // Function calls and responses replayed in agent loops are part of the prompt. + if part.FunctionCall != nil { + args, _ := kitutil.Marshal(part.FunctionCall.Arguments) + inputTexts = append(inputTexts, part.FunctionCall.FunctionName, string(args)) + } + mediaParts := []GeminiPart{part} + if part.FunctionResponse != nil { + inputTexts = append(inputTexts, part.FunctionResponse.Name) + // Count tool output as the text the model reads, not JSON-escaped. + for _, key := range slices.Sorted(maps.Keys(part.FunctionResponse.Response)) { + if text, ok := part.FunctionResponse.Response[key].(string); ok { + inputTexts = append(inputTexts, text) + continue + } + value, _ := kitutil.Marshal(part.FunctionResponse.Response[key]) + inputTexts = append(inputTexts, string(value)) + } + // Multimodal function responses carry media such as agent screenshots. + var responseParts []GeminiPart + if kitutil.Unmarshal(part.FunctionResponse.Parts, &responseParts) == nil { + mediaParts = append(mediaParts, responseParts...) + } + } + for _, media := range mediaParts { + source := media.InlineData.ToFileSource() + if source == nil { + continue + } + mimeType := media.InlineData.MimeType var fileType types.FileType if strings.HasPrefix(mimeType, "image/") { fileType = types.FileTypeImage @@ -91,13 +127,26 @@ func (r *GeminiChatRequest) GetTokenCountMeta() *types.TokenCountMeta { } else { fileType = types.FileTypeFile } + // Image cost depends on the media resolution; a part's own level + // overrides the request-wide one. + detail := string(r.GenerationConfig.MediaResolution) + var partResolution struct { + Level string `json:"level"` + } + if kitutil.Unmarshal(media.MediaResolution, &partResolution) == nil && partResolution.Level != "" { + detail = partResolution.Level + } files = append(files, &types.FileMeta{ FileType: fileType, Source: source, + Detail: detail, }) } } } + if len(r.Tools) > 0 { + inputTexts = append(inputTexts, string(r.Tools)) + } inputText := strings.Join(inputTexts, "\n") return &types.TokenCountMeta{ diff --git a/relaykit/dto/openai_request.go b/relaykit/dto/openai_request.go index a5fab3f82d..d64624314a 100644 --- a/relaykit/dto/openai_request.go +++ b/relaykit/dto/openai_request.go @@ -251,6 +251,14 @@ func (r *GeneralOpenAIRequest) GetTokenCountMeta() *types.TokenCountMeta { } } } + // Assistant tool calls replayed in agent loops are part of the prompt. + for _, toolCall := range message.ParseToolCalls() { + if len(toolCall.Custom) > 0 { + texts = append(texts, string(toolCall.Custom)) + continue + } + texts = append(texts, toolCall.Function.Name, toolCall.Function.Arguments) + } } tools := r.Tools @@ -259,12 +267,17 @@ func (r *GeneralOpenAIRequest) GetTokenCountMeta() *types.TokenCountMeta { } for _, tool := range tools { tokenCountMeta.ToolsCount++ + if len(tool.Custom) > 0 { + texts = append(texts, string(tool.Custom)) + continue + } texts = append(texts, tool.Function.Name) if tool.Function.Description != "" { texts = append(texts, tool.Function.Description) } if tool.Function.Parameters != nil { - texts = append(texts, fmt.Sprintf("%v", tool.Function.Parameters)) + parameters, _ := kitutil.Marshal(tool.Function.Parameters) + texts = append(texts, string(parameters)) } } //toolTokens := CountTokenInput(countStr, request.Model) @@ -747,9 +760,9 @@ func (m *Message) ParseContent() []MediaContent { case ContentTypeImageURL: imageUrl := contentItem["image_url"] - temp := &MessageImageUrl{ - Detail: "high", - } + // An omitted detail stays empty: upstream defaults it to auto, which + // costs more than high on gpt-5.5 and later. + temp := &MessageImageUrl{} switch v := imageUrl.(type) { case string: temp.Url = v @@ -1107,7 +1120,12 @@ func (r *OpenAIResponsesRequest) GetTokenCountMeta() *types.TokenCountMeta { } if len(r.Instructions) > 0 { - texts = append(texts, string(r.Instructions)) + var instructions string + if kitutil.GetJsonType(r.Instructions) == "string" && kitutil.Unmarshal(r.Instructions, &instructions) == nil { + texts = append(texts, instructions) + } else { + texts = append(texts, string(r.Instructions)) + } } if len(r.Metadata) > 0 { @@ -1166,6 +1184,11 @@ type Input struct { Type string `json:"type,omitempty"` Role string `json:"role,omitempty"` Content json.RawMessage `json:"content,omitempty"` + // Tool-call items replayed in agent loops. + Name string `json:"name,omitempty"` + Arguments json.RawMessage `json:"arguments,omitempty"` + ToolInput json.RawMessage `json:"input,omitempty"` + Output json.RawMessage `json:"output,omitempty"` } type MediaInput struct { @@ -1180,7 +1203,8 @@ type MediaInput struct { // Reference implementation mirrors Message.ParseContent: // - input can be a string, treated as an input_text item // - input can be an array of objects with a `type` field -// supported types: input_text, input_image, input_file +// supported types: input_text, output_text, input_image, input_file +// - tool-call items contribute their name, arguments or input, and output as input_text func (r *OpenAIResponsesRequest) ParseInput() []MediaInput { if r.Input == nil { return nil @@ -1205,15 +1229,28 @@ func (r *OpenAIResponsesRequest) ParseInput() []MediaInput { var inputs []Input _ = kitutil.Unmarshal(r.Input, &inputs) for _, input := range inputs { - if kitutil.GetJsonType(input.Content) == "string" { + content := input.Content + switch input.Type { + case "function_call": + content = input.Arguments + case "custom_tool_call": + content = input.ToolInput + case "function_call_output", "custom_tool_call_output", "local_shell_call_output": + content = input.Output + } + if input.Name != "" { + mediaInputs = append(mediaInputs, MediaInput{Type: "input_text", Text: input.Name}) + } + + if kitutil.GetJsonType(content) == "string" { var str string - _ = kitutil.Unmarshal(input.Content, &str) + _ = kitutil.Unmarshal(content, &str) mediaInputs = append(mediaInputs, MediaInput{Type: "input_text", Text: str}) } - if kitutil.GetJsonType(input.Content) == "array" { + if kitutil.GetJsonType(content) == "array" { var array []any - _ = kitutil.Unmarshal(input.Content, &array) + _ = kitutil.Unmarshal(content, &array) for _, itemAny := range array { // Already parsed MediaContent if media, ok := itemAny.(MediaInput); ok { @@ -1232,7 +1269,7 @@ func (r *OpenAIResponsesRequest) ParseInput() []MediaInput { continue } switch typeVal { - case "input_text": + case "input_text", "output_text": text, _ := item["text"].(string) mediaInputs = append(mediaInputs, MediaInput{Type: "input_text", Text: text}) case "input_image": diff --git a/relaykit/dto/token_count_meta_test.go b/relaykit/dto/token_count_meta_test.go new file mode 100644 index 0000000000..4aca6a054d --- /dev/null +++ b/relaykit/dto/token_count_meta_test.go @@ -0,0 +1,100 @@ +package dto + +import ( + "testing" + + kitutil "github.com/QuantumNous/new-api/relaykit/relayconvert/kitutil" + "github.com/QuantumNous/new-api/relaykit/types" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// Agent loops replay tool definitions, tool calls, and tool output on every +// turn; the prompt estimate must see all of them as the model reads them. +func TestTokenCountMetaCountsAgentLoopContent(t *testing.T) { + const toolOutput = "package main\n\nfunc main() {\n\tprintln(\"hi\")\n}" + const imageData = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChwGA60e6kgAAAABJRU5ErkJggg==" + + for _, tc := range []struct { + name string + raw string + request interface{ GetTokenCountMeta() *types.TokenCountMeta } + contains []string + files int + detail string + toolCount int + }{ + { + name: "chat", + raw: `{"model":"m","messages":[ + {"role":"user","content":[{"type":"text","text":"show main.go"},{"type":"image_url","image_url":{"url":"data:image/png;base64,` + imageData + `"}}]}, + {"role":"assistant","content":null,"tool_calls":[{"id":"c1","type":"function","function":{"name":"read_file","arguments":"{\"path\":\"main.go\"}"}}]}, + {"role":"tool","tool_call_id":"c1","content":"package main\n\nfunc main() {\n\tprintln(\"hi\")\n}"}], + "tools":[{"type":"function","function":{"name":"read_file","description":"Read a file","parameters":{"type":"object","properties":{"path":{"type":"string"}}}}}]}`, + request: &GeneralOpenAIRequest{}, + contains: []string{"read_file", `{"path":"main.go"}`, toolOutput, "Read a file", `"properties":{"path":{"type":"string"}}`}, + // An omitted detail stays empty so the estimate uses the upstream default (auto). + files: 1, + toolCount: 1, + }, + { + name: "responses", + raw: `{"model":"m","instructions":"Be brief.\nUse tools.","input":[ + {"type":"message","role":"user","content":[{"type":"input_text","text":"show main.go"}]}, + {"type":"message","role":"assistant","content":[{"type":"output_text","text":"Reading it now."}]}, + {"type":"function_call","call_id":"c1","name":"read_file","arguments":"{\"path\":\"main.go\"}"}, + {"type":"function_call_output","call_id":"c1","output":"package main\n\nfunc main() {\n\tprintln(\"hi\")\n}"}, + {"type":"custom_tool_call","call_id":"c2","name":"apply_patch","input":"*** Begin Patch"}, + {"type":"custom_tool_call_output","call_id":"c2","output":[{"type":"input_text","text":"Done!"}]}]}`, + request: &OpenAIResponsesRequest{}, + contains: []string{"Be brief.\nUse tools.", "Reading it now.", "read_file", `{"path":"main.go"}`, toolOutput, "apply_patch", "*** Begin Patch", "Done!"}, + }, + { + name: "claude", + raw: `{"model":"m","max_tokens":1,"messages":[ + {"role":"user","content":"show main.go"}, + {"role":"assistant","content":[{"type":"tool_use","id":"c1","name":"read_file","input":{"path":"main.go"}}]}, + {"role":"user","content":[ + {"type":"tool_result","tool_use_id":"c1","content":"package main\n\nfunc main() {\n\tprintln(\"hi\")\n}"}, + {"type":"tool_result","tool_use_id":"c2","content":[{"type":"text","text":"screenshot taken"},{"type":"image","source":{"type":"base64","media_type":"image/png","data":"` + imageData + `"}}]}]}], + "tools":[{"name":"read_file","description":"Read a file","input_schema":{"type":"object","properties":{"path":{"type":"string"}}}}]}`, + request: &ClaudeRequest{}, + contains: []string{"read_file", `{"path":"main.go"}`, toolOutput, "screenshot taken", "Read a file", `"properties":{"path":{"type":"string"}}`}, + files: 1, + toolCount: 1, + }, + { + name: "gemini", + raw: `{"systemInstruction":{"parts":[{"text":"Be brief."}]},"contents":[ + {"role":"user","parts":[{"text":"show main.go"}]}, + {"role":"model","parts":[{"functionCall":{"name":"read_file","args":{"path":"main.go"}}}]}, + {"role":"user","parts":[{"functionResponse":{"name":"read_file","response":{"output":"package main\n\nfunc main() {\n\tprintln(\"hi\")\n}","exit_code":0}}}]}, + {"role":"user","parts":[{"functionResponse":{"name":"screenshot","response":{},"parts":[{"inlineData":{"mimeType":"image/png","data":"` + imageData + `"}}]}}]}], + "generationConfig":{"mediaResolution":"MEDIA_RESOLUTION_LOW"}, + "tools":[{"functionDeclarations":[{"name":"read_file","description":"Read a file"}]}]}`, + request: &GeminiChatRequest{}, + contains: []string{"Be brief.", "read_file", `{"path":"main.go"}`, toolOutput, "Read a file"}, + files: 1, + detail: "MEDIA_RESOLUTION_LOW", + }, + } { + t.Run(tc.name, func(t *testing.T) { + require.NoError(t, kitutil.Unmarshal([]byte(tc.raw), tc.request)) + meta := tc.request.GetTokenCountMeta() + for _, want := range tc.contains { + assert.Contains(t, meta.CombineText, want) + } + // Tool output and instructions are counted as raw text, not JSON-escaped, + // and images stay out of the text. + assert.NotContains(t, meta.CombineText, `\n`) + assert.NotContains(t, meta.CombineText, "map[") + assert.NotContains(t, meta.CombineText, imageData) + require.Len(t, meta.Files, tc.files) + for _, file := range meta.Files { + assert.Equal(t, types.FileTypeImage, file.FileType) + assert.Equal(t, tc.detail, file.Detail) + } + assert.Equal(t, tc.toolCount, meta.ToolsCount) + }) + } +} diff --git a/relaykit/relayconvert/internal/claude_messages/to_oai_chat_resp.go b/relaykit/relayconvert/internal/claude_messages/to_oai_chat_resp.go index 1e9a5c4a7e..5398c8f0d5 100644 --- a/relaykit/relayconvert/internal/claude_messages/to_oai_chat_resp.go +++ b/relaykit/relayconvert/internal/claude_messages/to_oai_chat_resp.go @@ -19,6 +19,8 @@ type ClaudeResponseInfo struct { Created int64 Model string ResponseText strings.Builder + // ToolUseCount counts streamed tool_use blocks for the output-token fallback. + ToolUseCount int Usage *dto.Usage Done bool @@ -524,6 +526,9 @@ func FormatClaudeResponseInfo(claudeResponse *dto.ClaudeResponse, oaiResponse *d if claudeResponse.Delta.Thinking != nil { claudeInfo.ResponseText.WriteString(*claudeResponse.Delta.Thinking) } + if claudeResponse.Delta.PartialJson != nil { + claudeInfo.ResponseText.WriteString(*claudeResponse.Delta.PartialJson) + } } } else if claudeResponse.Type == "message_delta" { if claudeResponse.Usage != nil { @@ -552,6 +557,11 @@ func FormatClaudeResponseInfo(claudeResponse *dto.ClaudeResponse, oaiResponse *d claudeInfo.Done = true } else if claudeResponse.Type == "content_block_start" { + // Tool calls are most of an agent turn's output; count them for usage fallback. + if block := claudeResponse.ContentBlock; block != nil && block.Type == "tool_use" { + claudeInfo.ResponseText.WriteString(block.Name) + claudeInfo.ToolUseCount++ + } } else { return false } diff --git a/service/responses_usage.go b/service/responses_usage.go index cdf02e0a95..a3339e6729 100644 --- a/service/responses_usage.go +++ b/service/responses_usage.go @@ -8,6 +8,7 @@ import ( relaycommon "github.com/QuantumNous/new-api/relay/common" "github.com/QuantumNous/new-api/relaykit/dto" "github.com/QuantumNous/new-api/relaykit/relayconvert" + "github.com/QuantumNous/new-api/tokenkit" ) // ResponsesUsageAccumulator owns the accounting facts for one Responses stream. @@ -93,7 +94,7 @@ func (a *ResponsesUsageAccumulator) Finish() *dto.Usage { } if a.usage.CompletionTokens == 0 { if output := a.outputText.String(); output != "" { - a.usage.CompletionTokens = CountTextToken(output, a.info.GetUpstreamModelName()) + a.usage.CompletionTokens = tokenkit.Count(a.info.GetUpstreamModelName(), output) } } // Upstream bills the prompt as soon as it starts generating, so a stream diff --git a/service/responses_usage_test.go b/service/responses_usage_test.go index 151435bbe3..c6be452bcd 100644 --- a/service/responses_usage_test.go +++ b/service/responses_usage_test.go @@ -7,6 +7,7 @@ import ( relaycommon "github.com/QuantumNous/new-api/relay/common" "github.com/QuantumNous/new-api/relaykit/dto" "github.com/QuantumNous/new-api/setting/operation_setting" + "github.com/QuantumNous/new-api/tokenkit" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) @@ -272,7 +273,7 @@ func TestResponsesUsageAccumulatorMissingUsageEstimation(t *testing.T) { {Type: dto.ResponsesOutputTypeItemDone, Item: &dto.ResponsesOutput{Type: dto.BuildInCallFunctionCall, Name: "shell"}}, }, wantPrompt: 100, - wantCompletion: CountTextToken(summary+arguments, model), + wantCompletion: tokenkit.Count(model, summary+arguments), }, { name: "created only then disconnect bills the prompt", @@ -302,7 +303,7 @@ func TestResponsesUsageAccumulatorMissingUsageEstimation(t *testing.T) { }}, }, wantPrompt: 100, - wantCompletion: CountTextToken("final answer", model), + wantCompletion: tokenkit.Count(model, "final answer"), }, { name: "explicit failure without usage bills nothing", diff --git a/service/token_counter.go b/service/token_counter.go index 24d1cd494f..4ab311368c 100644 --- a/service/token_counter.go +++ b/service/token_counter.go @@ -15,6 +15,7 @@ import ( constant2 "github.com/QuantumNous/new-api/relay/constant" "github.com/QuantumNous/new-api/relaykit/dto" "github.com/QuantumNous/new-api/relaykit/types" + "github.com/QuantumNous/new-api/tokenkit" "github.com/gin-gonic/gin" ) @@ -24,153 +25,24 @@ func getImageToken(c *gin.Context, fileMeta *types.FileMeta, model string, strea return 0, fmt.Errorf("image_url_is_nil") } - // Defaults for 4o/4.1/4.5 family unless overridden below - baseTokens := 85 - tileTokens := 170 - - // Model classification - lowerModel := strings.ToLower(model) - - // Special cases from existing behavior - if strings.HasPrefix(lowerModel, "glm-4") { - return 1047, nil + img := tokenkit.Image{Detail: fileMeta.Detail} + // Reading the size may download a URL image; operators can turn that off, + // and the estimate then assumes a typical image size. + measure := constant.GetMediaToken && (stream || constant.GetMediaTokenNotStream) + if !measure || !tokenkit.ImageNeedsSize(model, img.Detail) { + return tokenkit.ImageTokens(model, img), nil } - // Patch-based models (32x32 patches, capped at 1536, with multiplier) - isPatchBased := false - multiplier := 1.0 - switch { - case strings.Contains(lowerModel, "gpt-4.1-mini"): - isPatchBased = true - multiplier = 1.62 - case strings.Contains(lowerModel, "gpt-4.1-nano"): - isPatchBased = true - multiplier = 2.46 - case strings.HasPrefix(lowerModel, "o4-mini"): - isPatchBased = true - multiplier = 1.72 - case strings.HasPrefix(lowerModel, "gpt-5-mini"): - isPatchBased = true - multiplier = 1.62 - case strings.HasPrefix(lowerModel, "gpt-5-nano"): - isPatchBased = true - multiplier = 2.46 - } - - // Tile-based model tokens and bases per doc - if !isPatchBased { - if strings.HasPrefix(lowerModel, "gpt-4o-mini") { - baseTokens = 2833 - tileTokens = 5667 - } else if strings.HasPrefix(lowerModel, "gpt-5-chat-latest") || (strings.HasPrefix(lowerModel, "gpt-5") && !strings.Contains(lowerModel, "mini") && !strings.Contains(lowerModel, "nano")) { - baseTokens = 70 - tileTokens = 140 - } else if strings.HasPrefix(lowerModel, "o1") || strings.HasPrefix(lowerModel, "o3") || strings.HasPrefix(lowerModel, "o1-pro") { - baseTokens = 75 - tileTokens = 150 - } else if strings.Contains(lowerModel, "computer-use-preview") { - baseTokens = 65 - tileTokens = 129 - } else if strings.Contains(lowerModel, "4.1") || strings.Contains(lowerModel, "4o") || strings.Contains(lowerModel, "4.5") { - baseTokens = 85 - tileTokens = 170 - } - } - - // Respect existing feature flags/short-circuits - if fileMeta.Detail == "low" && !isPatchBased { - return baseTokens, nil - } - - // Whether to count image tokens at all - if !constant.GetMediaToken { - return 3 * baseTokens, nil - } - - if !constant.GetMediaTokenNotStream && !stream { - return 3 * baseTokens, nil - } - // Normalize detail - if fileMeta.Detail == "auto" || fileMeta.Detail == "" { - fileMeta.Detail = "high" - } - - // 使用统一的文件服务获取图片配置 config, format, err := GetImageConfig(c, fileMeta.Source) if err != nil { return 0, err } - if config.Width == 0 || config.Height == 0 { - // not an image, but might be a valid file - if format != "" { - // file type - return 3 * baseTokens, nil - } - return 0, errors.New(fmt.Sprintf("fail to decode image config: %s", fileMeta.GetIdentifier())) + if (config.Width == 0 || config.Height == 0) && format == "" { + return 0, fmt.Errorf("fail to decode image config: %s", fileMeta.GetIdentifier()) } - - width := config.Width - height := config.Height - logger.LogDebug(c, "image token input: format=%s, width=%d, height=%d", format, width, height) - - if isPatchBased { - // 32x32 patch-based calculation with 1536 cap and model multiplier - ceilDiv := func(a, b int) int { return (a + b - 1) / b } - rawPatchesW := ceilDiv(width, 32) - rawPatchesH := ceilDiv(height, 32) - rawPatches := rawPatchesW * rawPatchesH - if rawPatches > 1536 { - // scale down - area := float64(width * height) - r := math.Sqrt(float64(32*32*1536) / area) - wScaled := float64(width) * r - hScaled := float64(height) * r - // adjust to fit whole number of patches after scaling - adjW := math.Floor(wScaled/32.0) / (wScaled / 32.0) - adjH := math.Floor(hScaled/32.0) / (hScaled / 32.0) - adj := math.Min(adjW, adjH) - if !math.IsNaN(adj) && adj > 0 { - r = r * adj - } - wScaled = float64(width) * r - hScaled = float64(height) * r - patchesW := math.Ceil(wScaled / 32.0) - patchesH := math.Ceil(hScaled / 32.0) - imageTokens := min(int(patchesW*patchesH), 1536) - return common.QuotaRound(float64(imageTokens) * multiplier), nil - } - // below cap - imageTokens := rawPatches - return common.QuotaRound(float64(imageTokens) * multiplier), nil - } - - // Tile-based calculation for 4o/4.1/4.5/o1/o3/etc. - // Step 1: fit within 2048x2048 square - maxSide := math.Max(float64(width), float64(height)) - fitScale := 1.0 - if maxSide > 2048 { - fitScale = maxSide / 2048.0 - } - fitW := int(math.Round(float64(width) / fitScale)) - fitH := int(math.Round(float64(height) / fitScale)) - - // Step 2: scale so that shortest side is exactly 768 - minSide := math.Min(float64(fitW), float64(fitH)) - if minSide == 0 { - return baseTokens, nil - } - shortScale := 768.0 / minSide - finalW := int(math.Round(float64(fitW) * shortScale)) - finalH := int(math.Round(float64(fitH) * shortScale)) - - // Count 512px tiles - tilesW := (finalW + 512 - 1) / 512 - tilesH := (finalH + 512 - 1) / 512 - tiles := tilesW * tilesH - - logger.LogDebug(c, "image token scaled size: width=%d, height=%d, tiles=%d", finalW, finalH, tiles) - - return tiles*tileTokens + baseTokens, nil + img.Width, img.Height = config.Width, config.Height + logger.LogDebug(c, "image token input: format=%s, width=%d, height=%d", format, img.Width, img.Height) + return tokenkit.ImageTokens(model, img), nil } func EstimateRequestToken(c *gin.Context, meta *types.TokenCountMeta, info *relaycommon.RelayInfo) (int, error) { @@ -228,7 +100,7 @@ func CountRequestToken(c *gin.Context, meta *types.TokenCountMeta, info *relayco if meta.TokenType == types.TokenTypeTextNumber { tkm += utf8.RuneCountInString(meta.CombineText) } else { - tkm += CountTextToken(meta.CombineText, model) + tkm += tokenkit.Count(model, meta.CombineText) } if info.RelayFormat == types.RelayFormatOpenAI { @@ -279,15 +151,11 @@ func CountRequestToken(c *gin.Context, meta *types.TokenCountMeta, info *relayco for i, file := range meta.Files { switch file.FileType { case types.FileTypeImage: - if common.IsOpenAITextModel(model) { - token, err := getImageToken(c, file, model, info.IsStream) - if err != nil { - return 0, fmt.Errorf("error counting image token, media index[%d], identifier[%s], err: %v", i, file.GetIdentifier(), err) - } - tkm += token - } else { - tkm += 520 + token, err := getImageToken(c, file, model, info.IsStream) + if err != nil { + return 0, fmt.Errorf("error counting image token, media index[%d], identifier[%s], err: %v", i, file.GetIdentifier(), err) } + tkm += token case types.FileTypeAudio: tkm += 256 case types.FileTypeVideo: @@ -309,7 +177,7 @@ func CountTokenRealtime(info *relaycommon.RelayInfo, request dto.RealtimeEvent, switch request.Type { case dto.RealtimeEventTypeSessionUpdate: if request.Session != nil { - msgTokens := CountTextToken(request.Session.Instructions, model) + msgTokens := tokenkit.Count(model, request.Session.Instructions) textToken += msgTokens } case dto.RealtimeEventResponseAudioDelta: @@ -321,7 +189,7 @@ func CountTokenRealtime(info *relaycommon.RelayInfo, request dto.RealtimeEvent, audioToken += atk case dto.RealtimeEventResponseAudioTranscriptionDelta, dto.RealtimeEventResponseFunctionCallArgumentsDelta: // count text token - tkm := CountTextToken(request.Delta, model) + tkm := tokenkit.Count(model, request.Delta) textToken += tkm case dto.RealtimeEventInputAudioBufferAppend: // count audio token @@ -336,7 +204,7 @@ func CountTokenRealtime(info *relaycommon.RelayInfo, request dto.RealtimeEvent, case "message": for _, content := range request.Item.Content { if content.Type == "input_text" { - tokens := CountTextToken(content.Text, model) + tokens := tokenkit.Count(model, content.Text) textToken += tokens } } @@ -360,19 +228,19 @@ func CountTokenRealtime(info *relaycommon.RelayInfo, request dto.RealtimeEvent, func CountTokenInput(input any, model string) int { switch v := input.(type) { case string: - return CountTextToken(v, model) + return tokenkit.Count(model, v) case []string: var text strings.Builder for _, s := range v { text.WriteString(s) } - return CountTextToken(text.String(), model) + return tokenkit.Count(model, text.String()) case []any: var text strings.Builder for _, item := range v { text.WriteString(fmt.Sprintf("%v", item)) } - return CountTextToken(text.String(), model) + return tokenkit.Count(model, text.String()) } return CountTokenInput(fmt.Sprintf("%v", input), model) } @@ -400,17 +268,3 @@ func CountAudioTokenOutput(audioBase64 string, audioFormat string) (int, error) // duration 来自上游返回的音频元数据，饱和转换防止 int 回绕 return common.QuotaFromFloat(duration / 60 * 200 / 0.24), nil } - -// CountTextToken 统计文本的token数量，仅OpenAI模型使用tokenizer，其余模型使用估算 -func CountTextToken(text string, model string) int { - if text == "" { - return 0 - } - if common.IsOpenAITextModel(model) { - tokenEncoder := getTokenEncoder(model) - return getTokenNum(tokenEncoder, text) - } else { - // 非openai模型，使用tiktoken-go计算没有意义，使用估算节省资源 - return EstimateTokenByModel(model, text) - } -} diff --git a/service/token_estimator.go b/service/token_estimator.go deleted file mode 100644 index 9e27269ce3..0000000000 --- a/service/token_estimator.go +++ /dev/null @@ -1,230 +0,0 @@ -package service - -import ( - "math" - "strings" - "sync" - "unicode" -) - -// Provider 定义模型厂商大类 -type Provider string - -const ( - OpenAI Provider = "openai" // 代表 GPT-3.5, GPT-4, GPT-4o - Gemini Provider = "gemini" // 代表 Gemini 1.0, 1.5 Pro/Flash - Claude Provider = "claude" // 代表 Claude 3, 3.5 Sonnet - Unknown Provider = "unknown" // 兜底默认 -) - -// multipliers 定义不同厂商的计费权重 -type multipliers struct { - Word float64 // 英文单词 (每词) - Number float64 // 数字 (每连续数字串) - CJK float64 // 中日韩字符 (每字) - Symbol float64 // 普通标点符号 (每个) - MathSymbol float64 // 数学符号 (∑,∫,∂,√等，每个) - URLDelim float64 // URL分隔符 (/,:,?,&,=,#,%) - tokenizer优化好 - AtSign float64 // @符号 - 导致单词切分，消耗较高 - Emoji float64 // Emoji表情 (每个) - Newline float64 // 换行符/制表符 (每个) - Space float64 // 空格 (每个) - BasePad int // 基础起步消耗 (Start/End tokens) -} - -var ( - multipliersMap = map[Provider]multipliers{ - Gemini: { - Word: 1.15, Number: 2.8, CJK: 0.68, Symbol: 0.38, MathSymbol: 1.05, URLDelim: 1.2, AtSign: 2.5, Emoji: 1.08, Newline: 1.15, Space: 0.2, BasePad: 0, - }, - Claude: { - Word: 1.13, Number: 1.63, CJK: 1.21, Symbol: 0.4, MathSymbol: 4.52, URLDelim: 1.26, AtSign: 2.82, Emoji: 2.6, Newline: 0.89, Space: 0.39, BasePad: 0, - }, - OpenAI: { - Word: 1.02, Number: 1.55, CJK: 0.85, Symbol: 0.4, MathSymbol: 2.68, URLDelim: 1.0, AtSign: 2.0, Emoji: 2.12, Newline: 0.5, Space: 0.42, BasePad: 0, - }, - } - multipliersLock sync.RWMutex -) - -// getMultipliers 根据厂商获取权重配置 -func getMultipliers(p Provider) multipliers { - multipliersLock.RLock() - defer multipliersLock.RUnlock() - - switch p { - case Gemini: - return multipliersMap[Gemini] - case Claude: - return multipliersMap[Claude] - case OpenAI: - return multipliersMap[OpenAI] - default: - // 默认兜底 (按 OpenAI 的算) - return multipliersMap[OpenAI] - } -} - -// EstimateToken 计算 Token 数量 -func EstimateToken(provider Provider, text string) int { - m := getMultipliers(provider) - var count float64 - - // 状态机变量 - type WordType int - const ( - None WordType = iota - Latin - Number - ) - currentWordType := None - - for _, r := range text { - // 1. 处理空格和换行符 - if unicode.IsSpace(r) { - currentWordType = None - // 换行符和制表符使用Newline权重 - if r == '\n' || r == '\t' { - count += m.Newline - } else { - // 普通空格使用Space权重 - count += m.Space - } - continue - } - - // 2. 处理 CJK (中日韩) - 按字符计费 - if isCJK(r) { - currentWordType = None - count += m.CJK - continue - } - - // 3. 处理Emoji - 使用专门的Emoji权重 - if isEmoji(r) { - currentWordType = None - count += m.Emoji - continue - } - - // 4. 处理拉丁字母/数字 (英文单词) - if isLatinOrNumber(r) { - isNum := unicode.IsNumber(r) - newType := Latin - if isNum { - newType = Number - } - - // 如果之前不在单词中，或者类型发生变化（字母<->数字），则视为新token - // 注意：对于OpenAI，通常"version 3.5"会切分，"abc123xyz"有时也会切分 - // 这里简单起见，字母和数字切换时增加权重 - if currentWordType == None || currentWordType != newType { - if newType == Number { - count += m.Number - } else { - count += m.Word - } - currentWordType = newType - } - // 单词中间的字符不额外计费 - continue - } - - // 5. 处理标点符号/特殊字符 - 按类型使用不同权重 - currentWordType = None - if isMathSymbol(r) { - count += m.MathSymbol - } else if r == '@' { - count += m.AtSign - } else if isURLDelim(r) { - count += m.URLDelim - } else { - count += m.Symbol - } - } - - // 向上取整并加上基础 padding - return int(math.Ceil(count)) + m.BasePad -} - -// 辅助：判断是否为 CJK 字符 -func isCJK(r rune) bool { - return unicode.Is(unicode.Han, r) || - (r >= 0x3040 && r <= 0x30FF) || // 日文 - (r >= 0xAC00 && r <= 0xD7A3) // 韩文 -} - -// 辅助：判断是否为单词主体 (字母或数字) -func isLatinOrNumber(r rune) bool { - return unicode.IsLetter(r) || unicode.IsNumber(r) -} - -// 辅助：判断是否为Emoji字符 -func isEmoji(r rune) bool { - // Emoji的Unicode范围 - // 基本范围：0x1F300-0x1F9FF (Emoticons, Symbols, Pictographs) - // 补充范围：0x2600-0x26FF (Misc Symbols), 0x2700-0x27BF (Dingbats) - // 表情符号：0x1F600-0x1F64F (Emoticons) - // 其他：0x1F900-0x1F9FF (Supplemental Symbols and Pictographs) - return (r >= 0x1F300 && r <= 0x1F9FF) || - (r >= 0x2600 && r <= 0x26FF) || - (r >= 0x2700 && r <= 0x27BF) || - (r >= 0x1F600 && r <= 0x1F64F) || - (r >= 0x1F900 && r <= 0x1F9FF) || - (r >= 0x1FA00 && r <= 0x1FAFF) // Symbols and Pictographs Extended-A -} - -// 辅助：判断是否为数学符号 -func isMathSymbol(r rune) bool { - // 数学运算符和符号 - // 基本数学符号：∑ ∫ ∂ √ ∞ ≤ ≥ ≠ ≈ ± × ÷ - // 上下标数字：² ³ ¹ ⁴ ⁵ ⁶ ⁷ ⁸ ⁹ ⁰ - // 希腊字母等也常用于数学 - mathSymbols := "∑∫∂√∞≤≥≠≈±×÷∈∉∋∌⊂⊃⊆⊇∪∩∧∨¬∀∃∄∅∆∇∝∟∠∡∢°′″‴⁺⁻⁼⁽⁾ⁿ₀₁₂₃₄₅₆₇₈₉₊₋₌₍₎²³¹⁴⁵⁶⁷⁸⁹⁰" - for _, m := range mathSymbols { - if r == m { - return true - } - } - // Mathematical Operators (U+2200–U+22FF) - if r >= 0x2200 && r <= 0x22FF { - return true - } - // Supplemental Mathematical Operators (U+2A00–U+2AFF) - if r >= 0x2A00 && r <= 0x2AFF { - return true - } - // Mathematical Alphanumeric Symbols (U+1D400–U+1D7FF) - if r >= 0x1D400 && r <= 0x1D7FF { - return true - } - return false -} - -// 辅助：判断是否为URL分隔符（tokenizer对这些优化较好） -func isURLDelim(r rune) bool { - // URL中常见的分隔符，tokenizer通常优化处理 - urlDelims := "/:?&=;#%" - for _, d := range urlDelims { - if r == d { - return true - } - } - return false -} - -func EstimateTokenByModel(model, text string) int { - // strings.Contains(model, "gpt-4o") - if text == "" { - return 0 - } - - model = strings.ToLower(model) - if strings.Contains(model, "gemini") { - return EstimateToken(Gemini, text) - } else if strings.Contains(model, "claude") { - return EstimateToken(Claude, text) - } else { - return EstimateToken(OpenAI, text) - } -} diff --git a/service/tokenizer.go b/service/tokenizer.go deleted file mode 100644 index 9cf632b86a..0000000000 --- a/service/tokenizer.go +++ /dev/null @@ -1,63 +0,0 @@ -package service - -import ( - "sync" - - "github.com/QuantumNous/new-api/common" - "github.com/tiktoken-go/tokenizer" - "github.com/tiktoken-go/tokenizer/codec" -) - -// tokenEncoderMap won't grow after initialization -var defaultTokenEncoder tokenizer.Codec - -// tokenEncoderMap is used to store token encoders for different models -var tokenEncoderMap = make(map[string]tokenizer.Codec) - -// tokenEncoderMutex protects tokenEncoderMap for concurrent access -var tokenEncoderMutex sync.RWMutex - -func InitTokenEncoders() { - common.SysLog("initializing token encoders") - defaultTokenEncoder = codec.NewCl100kBase() - common.SysLog("token encoders initialized") -} - -func getTokenEncoder(model string) tokenizer.Codec { - // First, try to get the encoder from cache with read lock - tokenEncoderMutex.RLock() - if encoder, exists := tokenEncoderMap[model]; exists { - tokenEncoderMutex.RUnlock() - return encoder - } - tokenEncoderMutex.RUnlock() - - // If not in cache, create new encoder with write lock - tokenEncoderMutex.Lock() - defer tokenEncoderMutex.Unlock() - - // Double-check if another goroutine already created the encoder - if encoder, exists := tokenEncoderMap[model]; exists { - return encoder - } - - // Create new encoder - modelCodec, err := tokenizer.ForModel(tokenizer.Model(model)) - if err != nil { - // Cache the default encoder for this model to avoid repeated failures - tokenEncoderMap[model] = defaultTokenEncoder - return defaultTokenEncoder - } - - // Cache the new encoder - tokenEncoderMap[model] = modelCodec - return modelCodec -} - -func getTokenNum(tokenEncoder tokenizer.Codec, text string) int { - if text == "" { - return 0 - } - tkm, _ := tokenEncoder.Count(text) - return tkm -} diff --git a/service/usage_helpr.go b/service/usage_helpr.go index b5c35ee7b7..9c949bd6a7 100644 --- a/service/usage_helpr.go +++ b/service/usage_helpr.go @@ -4,6 +4,7 @@ import ( "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/constant" "github.com/QuantumNous/new-api/relaykit/dto" + "github.com/QuantumNous/new-api/tokenkit" "github.com/gin-gonic/gin" ) @@ -23,7 +24,7 @@ func ResponseText2Usage(c *gin.Context, responseText string, modeName string, pr common.SetContextKey(c, constant.ContextKeyLocalCountTokens, true) usage := &dto.Usage{} usage.PromptTokens = promptTokens - usage.CompletionTokens = EstimateTokenByModel(modeName, responseText) + usage.CompletionTokens = tokenkit.Estimate(modeName, responseText) usage.TotalTokens = usage.PromptTokens + usage.CompletionTokens return usage } diff --git a/tokenkit/README.md b/tokenkit/README.md new file mode 100644 index 0000000000..196da4116f --- /dev/null +++ b/tokenkit/README.md @@ -0,0 +1,125 @@ +# TokenKit + +TokenKit 是从 [new-api](https://github.com/QuantumNous/new-api) 中拆分出的独立 Go 模块，用来在不调用上游 API 的情况下，计算或估算大模型请求的 token 数。 + +它只做计数，不负责解析请求、下载图片或计费。调用方传入文本，或者图片的尺寸和 detail 设置，TokenKit 返回 token 数。可以脱离 new-api 主模块，嵌入其他 Go 网关、代理或客户端。 + +## 能力 + +- GPT 和 o 系列模型用 tiktoken 精确计数，GPT-4 和 GPT-3.5 走 `cl100k_base`，其余走 `o200k_base` +- Claude、Gemini 和其他模型按字符类别加权估算，Claude Opus 4.7 前后的两套 tokenizer 分开处理 +- 按各家公开的尺寸规则计算图片 token，包括 OpenAI 的 detail 级别和 Gemini 3 的 media resolution +- 不依赖 new-api 主模块、Gin、数据库或全局设置 + +## 准确度 + +文本权重是用真实 agent 会话校准的：从 Claude Code 和 Codex 的会话里取了 54 段，包括工具定义、工具调用和工具输出，和上游返回的 usage 逐条对比。 + +| 模型 | 测试模型 | prompt 偏差 / 平均误差 | 输出文本 偏差 / 平均误差 | +|---|---|---|---| +| Claude（4.7 之前） | claude-haiku-4-5 | +0.5% / 3.3% | −0.6% / 3.8% | +| Claude（4.7 及以后） | claude-sonnet-5 | +0.3% / 3.5% | −1.1% / 4.3% | +| GPT（估算） | gpt-5.6-luna | +0.9% / 2.9% | −1.8% / 5.3% | +| Gemini | gemini-3.7-flash | +0.7% / 4.2% | −1.3% / 4.7% | + +GPT 用 `Count` 时和 tiktoken 基本一致。有两处已知差异： +- 只含缩进空格的空行会多算 1 个 token，来自 tiktoken-go v0.8.1 的分词正则，1.2 MB 的 agent 文本上一共多 3 个。 +- 超过 1024 字节、中间没有分隔的连续片段（例如没有标点的长段汉字、长串空格）会被切开再计数，误差不到 0.5%。不切的话，BPE 的耗时随片段长度平方增长：10 万个连续汉字要 20 秒左右，切开后约 0.1 秒。 + +估算适合 agent 流量：代码、JSON、命令输出和中英文混排。几百 token 的短文本误差会大一些。 + +图片规则和上游实测一致：OpenAI 的 patch 和 tile 计算、Claude 两个分辨率档位、Gemini 2.x 的切块。例外和注意事项见[图片](#图片)。 + +## 安装 + +TokenKit 要求 Go 1.26 或更高版本。 + +```bash +go get github.com/QuantumNous/new-api/tokenkit@latest +``` + +## 快速开始 + +```go +package main + +import ( + "fmt" + + "github.com/QuantumNous/new-api/tokenkit" +) + +func main() { + prompt := "读取 main.go 并解释 init 函数" + + // GPT 用 tiktoken 精确计数，其他模型估算 + fmt.Println(tokenkit.Count("gpt-5.6-sol", prompt)) + fmt.Println(tokenkit.Count("claude-opus-5-5", prompt)) + + // 只要估算，不跑 tokenizer，比如流式输出中途断开时补算 completion + fmt.Println(tokenkit.Estimate("gpt-5.6-sol", prompt)) + + // 一张 1920×1080 的截图 + shot := tokenkit.Image{Width: 1920, Height: 1080} + fmt.Println(tokenkit.ImageTokens("claude-opus-5-5", shot)) // 2694 + fmt.Println(tokenkit.ImageTokens("claude-haiku-4-5", shot)) // 1564 + shot.Detail = "low" + fmt.Println(tokenkit.ImageTokens("gpt-5.6-sol", shot)) // 173 +} +``` + +## 模型分组 + +`ModelFamily` 按模型名分组，不区分大小写，允许带厂商前缀和日期后缀，例如 `us.anthropic.claude-opus-4-7-v1:0`、`claude-sonnet-4-5-20250929`。 + +| 分组 | 匹配 | 文本 | +|---|---|---| +| `FamilyOpenAI` | 名称含 `gpt-`、`o1`、`o3`、`o4`、`chatgpt` | `Count` 精确，`Estimate` 估算 | +| `FamilyClaude` | Claude 4.7 之前，例如 Haiku 4.5、Sonnet 4.6 | 估算 | +| `FamilyClaude47` | Claude 4.7 及以后，以及读不出版本号的名称，例如 Opus 4.7、Sonnet 5、Fable、Mythos | 估算 | +| `FamilyGemini` | 名称含 `gemini` | 估算 | +| `FamilyUnknown` | 其他模型 | 估算，权重未校准 | + +Claude 4.7 及以后的 tokenizer 对同样的文本多出约 30% 的 token，所以两组 Claude 用不同的权重。 + +## 图片 + +`ImageTokens(model, Image)` 返回一张图片占用的 prompt token。`Image.Width` 和 `Image.Height` 是像素尺寸，为 0 时按 1024×1024 计算。`ImageNeedsSize` 判断某个模型和 detail 设置下是否需要尺寸；返回 false 时可以不读取图片。 + +| 模型 | 计算方式 | +|---|---| +| gpt-5.2 及以后、gpt-6 | 32px patch × 1.2，按模型和 detail 缩放；不传 detail 时 gpt-5.5 及以后按 `original` 计算 | +| gpt-5、gpt-5.1、gpt-4o、gpt-4.1、o1、o3 | 512px tile，加 base tokens | +| Claude 4.7 之前 | ⌈宽/28⌉ × ⌈高/28⌉，最长边 1568，最多 1568 个 patch，另加 4 | +| Claude 4.7 及以后 | 同上，最长边 2576，最多 4784 个 patch，另加 3 | +| Gemini 3 | 与尺寸无关，按 media resolution：280 / 560 / 1120 / 2240，默认 1120 | +| Gemini 2.x | 两边都不超过 384 时 258；更大的图切成方块，每块 258，最多 12 块，另加 1 块缩略图 | +| 其他模型 | 每张 520 | + +`Image.Detail` 填 OpenAI 的 detail（`low`、`high`、`auto`、`original`），或 Gemini 的 media resolution（`MEDIA_RESOLUTION_LOW` 等）。两类值只在对应的模型上生效，例如 Gemini 会忽略 `low`。 + +注意事项： + +- 有两处按上游实际计费，而不是按官方文档。gpt-5.5 和 gpt-5.6 的 `high` 先缩到 2048×2048 以内，最短边再缩到 768，文档写的是 2500 patch 预算。gpt-6 的 `high` 先缩到 2048×2048 以内，再按 2500 patch 预算缩放，文档没有写 2048 这一步。 +- Gemini 3 用的是文档值。实测按图片宽高比低 2% 到 9%。 +- Gemini 2.x 的 4K 图会多算：3840×2160 实际 2322，这里给 3354。 +- 图片尺寸来自请求方，可能被伪造。超过 65535 像素的边会先缩小，单张最多按 30000 个 patch 计算，结果始终有上限。 + +## 开发 + +TokenKit 必须始终保持独立可构建，不依赖 new-api 主模块或 RelayKit。修改模块后，在 `tokenkit` 目录运行： + +```bash +GOWORK=off go test ./... +GOWORK=off go build ./... +``` + +`image_test.go` 里的期望值是上游实测值。改动图片规则时，用同样的方法重新测量：同一个请求带图和不带图各发一次，两次 input token 的差就是这张图的消耗。 + +## 版本与兼容性 + +TokenKit 当前使用 `v0.x` 版本，公开 API 可能在小版本中调整。估算权重和图片规则会随上游模型更新重新校准，同样的输入在不同版本之间可能得到不同的结果。 + +## 许可证 + +TokenKit 是 new-api 项目的一部分，遵循项目根目录中的 [GNU Affero General Public License v3.0](../LICENSE)。 diff --git a/tokenkit/go.mod b/tokenkit/go.mod new file mode 100644 index 0000000000..3ee5ab7598 --- /dev/null +++ b/tokenkit/go.mod @@ -0,0 +1,15 @@ +module github.com/QuantumNous/new-api/tokenkit + +go 1.26 + +require ( + github.com/stretchr/testify v1.11.1 + github.com/tiktoken-go/tokenizer v0.8.1 +) + +require ( + github.com/davecgh/go-spew v1.1.1 // indirect + github.com/dlclark/regexp2/v2 v2.5.1 // indirect + github.com/pmezard/go-difflib v1.0.0 // indirect + gopkg.in/yaml.v3 v3.0.1 // indirect +) diff --git a/tokenkit/go.sum b/tokenkit/go.sum new file mode 100644 index 0000000000..289a69f3f3 --- /dev/null +++ b/tokenkit/go.sum @@ -0,0 +1,14 @@ +github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= +github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/dlclark/regexp2/v2 v2.5.1 h1:E5Ug7Dh264W1ymdySmiHNcDG7fmsR307APCE5R07a20= +github.com/dlclark/regexp2/v2 v2.5.1/go.mod h1:avUrQvPaLz2DrFNHJF0taWAFFX2C1GMSSoeiqFjcBmU= +github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= +github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U= +github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U= +github.com/tiktoken-go/tokenizer v0.8.1 h1:4obDoB6/dhdBt9xMweX4nww5cjdOq/nYF4ecwPq2+mg= +github.com/tiktoken-go/tokenizer v0.8.1/go.mod h1:eLA0t6nGvn9mDc7gt90qt7pMat+gE9ViqwQ6l9B+tA4= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= +gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= diff --git a/tokenkit/image.go b/tokenkit/image.go new file mode 100644 index 0000000000..e89541397c --- /dev/null +++ b/tokenkit/image.go @@ -0,0 +1,290 @@ +package tokenkit + +import ( + "math" + "strings" +) + +// Image describes one input image. +type Image struct { + // Width and Height are the pixel dimensions. Zero or negative means + // unknown; where size matters, a 1024×1024 image is assumed. + Width, Height int + // Detail is the OpenAI detail level (low, high, auto, original) or the + // Gemini media resolution (MEDIA_RESOLUTION_LOW, _MEDIUM, _HIGH, + // _ULTRA_HIGH). Empty means the provider default. Each family ignores + // values meant for the other. + Detail string +} + +const ( + unknownImageSide = 1024 + // Larger sides are scaled down first; OpenAI does the same for original + // detail and every other provider rejects such images. + maxImageSide = 65535 + // OpenAI rejects images that need more patches, so no estimate exceeds it. + maxImagePatches = 30000 + // Earlier flat estimate, kept for models whose sizing rules are unknown. + unknownFamilyImageTokens = 520 +) + +// ImageTokens returns the prompt tokens one image costs on model. +func ImageTokens(model string, img Image) int { + w, h := img.Width, img.Height + if w <= 0 || h <= 0 { + w, h = unknownImageSide, unknownImageSide + } + if w > maxImageSide || h > maxImageSide { + w, h = fitWithin(w, h, maxImageSide) + } + + switch ModelFamily(model) { + case FamilyOpenAI: + return openAIImageRuleFor(model, img.Detail).tokens(w, h) + case FamilyClaude: + // Standard tier; each image adds 4 tokens of framing (measured). + w, h = claudeResizedSize(w, h, 1568, 1568) + return ceilDiv(w, 28)*ceilDiv(h, 28) + 4 + case FamilyClaude47: + // High-resolution tier; each image adds 3 tokens of framing (measured). + w, h = claudeResizedSize(w, h, 2576, 4784) + return ceilDiv(w, 28)*ceilDiv(h, 28) + 3 + case FamilyGemini: + if geminiUsesMediaResolution(model) { + return geminiMediaResolutionTokens(img.Detail) + } + return geminiTileImageTokens(w, h) + default: + return unknownFamilyImageTokens + } +} + +// ImageNeedsSize reports whether ImageTokens depends on the image dimensions +// for model and detail. When it returns false, callers can skip reading the +// image. +func ImageNeedsSize(model, detail string) bool { + switch ModelFamily(model) { + case FamilyOpenAI: + rule := openAIImageRuleFor(model, detail) + return rule.tileTokens == 0 || !rule.lowDetail + case FamilyClaude, FamilyClaude47: + return true + case FamilyGemini: + return !geminiUsesMediaResolution(model) + default: + return false + } +} + +// openAIImageRule is the sizing rule for one OpenAI model and detail level, +// from https://developers.openai.com/api/docs/guides/images-vision and +// measured upstream usage. +type openAIImageRule struct { + // Tile-based models cost baseTokens plus tileTokens per 512px tile. + baseTokens, tileTokens int + lowDetail bool + + // Patch-based models cost 32px patches times multiplier. The multiplier + // is a float32 because upstream bills that way: 1225 patches × 1.2 bills + // 1471, not 1470. + multiplier float32 + maxSide int // fit within maxSide × maxSide + shortSide int // scale the shortest side down to shortSide + budget int // resizing patch budget +} + +func openAIImageRuleFor(model, detail string) openAIImageRule { + model = strings.ToLower(model) + detail = strings.ToLower(detail) + switch detail { + case "low", "high", "original": + default: + detail = "auto" + } + major, minor, isGPT := modelVersion(model, "gpt-") + + switch { + case isGPT && major >= 6: + // gpt-6-astra: auto means original. Upstream fits high within 2048 + // before the 2,500-patch budget; the documentation omits the 2048 limit. + switch detail { + case "low": + return openAIImageRule{multiplier: 1.2, maxSide: 512} + case "high": + return openAIImageRule{multiplier: 1.2, maxSide: 2048, budget: 2500} + default: + return openAIImageRule{multiplier: 1.2} + } + case isGPT && major == 5 && minor >= 5: + // gpt-5.5 and gpt-5.6: auto means original. Upstream bills high as + // fit 2048 then shortest side 768, not the documented 2,500-patch budget. + switch detail { + case "low": + return openAIImageRule{multiplier: 1.2, maxSide: 512} + case "high": + return openAIImageRule{multiplier: 1.2, maxSide: 2048, shortSide: 768} + } + if minor == 5 { + return openAIImageRule{multiplier: 1.2, maxSide: 6000, budget: 10000} + } + return openAIImageRule{multiplier: 1.2} + case isGPT && major == 5 && minor == 4: + // gpt-5.4: auto means high. + switch detail { + case "low": + return openAIImageRule{multiplier: 1.2, maxSide: 2048, budget: 6144} + case "original": + return openAIImageRule{multiplier: 1.2, maxSide: 6000, budget: 10000} + default: + return openAIImageRule{multiplier: 1.2, maxSide: 2048, budget: 2500} + } + case isGPT && major == 5 && minor >= 2: + return openAIImageRule{multiplier: 1.2, maxSide: 2048, budget: 6144} + case isGPT && major == 5 && minor == 0 && strings.Contains(model, "nano"): + return openAIImageRule{multiplier: 1.5, budget: 1536} + case isGPT && major == 5 && minor == 0 && strings.Contains(model, "mini"): + return openAIImageRule{multiplier: 1.2, budget: 1536} + case isGPT && major == 5: + return openAIImageRule{baseTokens: 70, tileTokens: 140, lowDetail: detail == "low"} + case isGPT && major == 4 && minor == 1 && strings.Contains(model, "mini"): + return openAIImageRule{multiplier: 1.62, maxSide: 2048, budget: 6144} + case isGPT && major == 4 && minor == 1 && strings.Contains(model, "nano"): + return openAIImageRule{multiplier: 2.46, budget: 1536} + case strings.Contains(model, "4o-mini"): + return openAIImageRule{baseTokens: 2833, tileTokens: 5667, lowDetail: detail == "low"} + case strings.Contains(model, "o4-mini"): + return openAIImageRule{multiplier: 1.72, budget: 1536} + case !isGPT && (strings.Contains(model, "o1") || strings.Contains(model, "o3")): + return openAIImageRule{baseTokens: 75, tileTokens: 150, lowDetail: detail == "low"} + default: + // gpt-4o, gpt-4.1, gpt-4.5, and older vision models. + return openAIImageRule{baseTokens: 85, tileTokens: 170, lowDetail: detail == "low"} + } +} + +func (r openAIImageRule) tokens(w, h int) int { + if r.tileTokens > 0 { + if r.lowDetail { + return r.baseTokens + } + w, h = fitWithin(w, h, 2048) + w, h = shrinkShortSide(w, h, 768) + return r.baseTokens + r.tileTokens*ceilDiv(w, 512)*ceilDiv(h, 512) + } + + if r.maxSide > 0 { + w, h = fitWithin(w, h, r.maxSide) + } + patches := ceilDiv(w, 32) * ceilDiv(h, 32) + if short, long := min(w, h), max(w, h); r.shortSide > 0 && short > r.shortSide { + // Upstream keeps the long side unrounded here: 3000x2000 fits to + // 2048x1365, and 2048 x 768/1365 = 1152.3 needs 37 patch columns. + patches = ceilDiv(r.shortSide, 32) * ceilDiv(long*r.shortSide, short*32) + } + if r.budget > 0 && patches > r.budget { + shrink := math.Sqrt(float64(32*32*r.budget) / (float64(w) * float64(h))) + scaledW, scaledH := float64(w)*shrink/32, float64(h)*shrink/32 + shrink *= min(math.Floor(scaledW)/scaledW, math.Floor(scaledH)/scaledH) + w, h = max(int(float64(w)*shrink), 1), max(int(float64(h)*shrink), 1) + patches = ceilDiv(w, 32) * ceilDiv(h, 32) + } + patches = min(patches, maxImagePatches) + return int(math.Ceil(float64(patches) * float64(r.multiplier))) +} + +// claudeResizedSize is the size Claude scales an image to: the largest +// aspect-preserving size whose padded sides fit maxEdge and whose 28px patches +// fit maxTokens. Reference implementation: +// https://platform.claude.com/docs/en/build-with-claude/vision-coordinates +func claudeResizedSize(width, height, maxEdge, maxTokens int) (int, int) { + if claudeImageFits(width, height, maxEdge, maxTokens) { + return width, height + } + if height > width { + resizedH, resizedW := claudeResizedSize(height, width, maxEdge, maxTokens) + return resizedW, resizedH + } + aspectRatio := float64(width) / float64(height) + lo, hi := 1, width // lo always fits; hi never fits + for lo+1 < hi { + mid := (lo + hi) / 2 + if claudeImageFits(mid, max(int(math.RoundToEven(float64(mid)/aspectRatio)), 1), maxEdge, maxTokens) { + lo = mid + } else { + hi = mid + } + } + return lo, max(int(math.RoundToEven(float64(lo)/aspectRatio)), 1) +} + +func claudeImageFits(w, h, maxEdge, maxTokens int) bool { + return ceilDiv(w, 28)*28 <= maxEdge && ceilDiv(h, 28)*28 <= maxEdge && ceilDiv(w, 28)*ceilDiv(h, 28) <= maxTokens +} + +// geminiUsesMediaResolution reports whether a lower-cased Gemini model bills +// images by media resolution regardless of size (Gemini 3 and later, and names +// without a version). +func geminiUsesMediaResolution(model string) bool { + major, _, ok := modelVersion(strings.ToLower(model), "gemini-") + return !ok || major >= 3 +} + +// geminiMediaResolutionTokens returns the documented per-image cost for a +// Gemini media resolution. Measured costs run 2-9% lower depending on aspect. +func geminiMediaResolutionTokens(detail string) int { + level, ok := strings.CutPrefix(strings.ToLower(detail), "media_resolution_") + if !ok { + return 1120 + } + switch level { + case "low": + return 280 + case "medium": + return 560 + case "ultra_high": + return 2240 + default: + return 1120 + } +} + +// geminiTileImageTokens follows the Gemini 2.x rule measured on gemini-2.5-flash: +// images up to 384px on both sides cost one 258-token tile; larger images are +// cut into square tiles of a third less than the shortest side (256 to 768px), +// at most 12, plus one thumbnail tile. 4K images bill below this cap. +func geminiTileImageTokens(w, h int) int { + if w <= 384 && h <= 384 { + return 258 + } + tile := min(max(min(w, h)*2/3, 256), 768) + tiles := min(ceilDiv(w, tile)*ceilDiv(h, tile), 12) + return (tiles + 1) * 258 +} + +// fitWithin scales an image down so neither side exceeds limit, rounding the +// shorter side down. Smaller images are not enlarged. +func fitWithin(w, h, limit int) (int, int) { + if w <= limit && h <= limit { + return w, h + } + if w >= h { + return limit, max(int(int64(h)*int64(limit)/int64(w)), 1) + } + return max(int(int64(w)*int64(limit)/int64(h)), 1), limit +} + +// shrinkShortSide scales an image down so its shortest side is at most limit, +// rounding the other side down. +func shrinkShortSide(w, h, limit int) (int, int) { + if min(w, h) <= limit { + return w, h + } + if w <= h { + return limit, max(int(int64(h)*int64(limit)/int64(w)), 1) + } + return max(int(int64(w)*int64(limit)/int64(h)), 1), limit +} + +func ceilDiv(a, b int) int { + return (a + b - 1) / b +} diff --git a/tokenkit/image_test.go b/tokenkit/image_test.go new file mode 100644 index 0000000000..85b2eb9e6e --- /dev/null +++ b/tokenkit/image_test.go @@ -0,0 +1,124 @@ +package tokenkit + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +// Unless marked otherwise, want is the per-image prompt-token cost measured +// through upstream usage (request with the image minus the same request +// without it) on 2026-10-07. +func TestImageTokensMatchUpstreamUsage(t *testing.T) { + for _, tc := range []struct { + model string + w, h int + detail string + want int + }{ + // gpt-5.6: auto means original; high scales the shortest side to 768. + {"gpt-5.6-luna", 200, 200, "", 59}, + {"gpt-5.6-luna", 1000, 1000, "", 1229}, + {"gpt-5.6-luna", 1000, 1000, "auto", 1229}, + {"gpt-5.6-luna", 1000, 1000, "original", 1229}, + {"gpt-5.6-luna", 1000, 1000, "high", 692}, + {"gpt-5.6-luna", 1000, 1000, "low", 308}, + {"gpt-5.6-luna", 200, 200, "high", 59}, + {"gpt-5.6-luna", 2048, 2048, "", 4916}, + {"gpt-5.6-luna", 2048, 2048, "high", 692}, + {"gpt-5.6-luna", 3840, 2160, "", 9793}, + {"gpt-5.6-luna", 3840, 2160, "high", 1239}, + {"gpt-5.6-luna", 3840, 2160, "low", 173}, + {"gpt-5.6-luna", 600, 3000, "", 2144}, + {"gpt-5.6-luna", 600, 3000, "high", 999}, + {"gpt-5.6-luna", 600, 3000, "low", 77}, + {"gpt-5.6-luna", 8000, 500, "", 4801}, + {"gpt-5.6-luna", 8000, 500, "high", 308}, + {"gpt-5.6-luna", 8000, 500, "low", 20}, + {"gpt-5.6-luna", 3000, 2000, "high", 1066}, + {"gpt-5.6-luna", 1170, 2532, "high", 1498}, + {"gpt-5.6-luna", 1600, 1200, "high", 922}, + // gpt-5.5 original is limited to 6000px and 10,000 patches. + {"gpt-5.5", 8000, 500, "", 2708}, + {"gpt-5.5", 3840, 2160, "", 9793}, + {"gpt-5.5", 2048, 2048, "high", 692}, + // gpt-6-astra high fits within 2048, then applies a 2,500-patch budget. + {"gpt-6-astra", 1024, 1024, "high", 1229}, + {"gpt-6-astra", 2048, 2048, "high", 3001}, + {"gpt-6-astra", 3840, 2160, "high", 2765}, + {"gpt-6-astra", 1170, 2532, "high", 2305}, + {"gpt-6-astra", 3000, 2000, "high", 2881}, + {"gpt-6-astra", 2048, 2048, "", 4916}, + {"gpt-6-astra", 3840, 2160, "", 9793}, + {"gpt-6-astra", 600, 3000, "", 2144}, + {"gpt-6-astra", 3840, 2160, "low", 173}, + // Tile-based models follow the documented formula. + {"gpt-4o", 1000, 1000, "", 85 + 170*4}, + {"gpt-4o", 1000, 1000, "low", 85}, + {"gpt-5.1", 200, 200, "high", 70 + 140}, + // Claude: ceil(w/28) x ceil(h/28) after the tier resize, plus framing. + {"claude-haiku-4-5", 200, 200, "", 68}, + {"claude-haiku-4-5", 1000, 1000, "", 1300}, + {"claude-haiku-4-5", 1092, 1092, "", 1525}, + {"claude-haiku-4-5", 1920, 1080, "", 1564}, + {"claude-haiku-4-5", 1075, 1520, "", 1555}, + {"claude-haiku-4-5", 2000, 1500, "", 1568}, + {"claude-haiku-4-5", 3840, 2160, "", 1564}, + {"claude-haiku-4-5", 4096, 512, "", 396}, + {"claude-sonnet-5", 200, 200, "", 67}, + {"claude-sonnet-5", 1000, 1000, "", 1299}, + {"claude-sonnet-5", 1092, 1092, "", 1524}, + {"claude-sonnet-5", 1920, 1080, "", 2694}, + {"claude-sonnet-5", 1075, 1520, "", 2148}, + {"claude-sonnet-5", 2000, 1500, "", 3891}, + {"claude-sonnet-5", 3840, 2160, "", 4787}, + {"claude-sonnet-5", 4096, 512, "", 1107}, + // Gemini 2.x tiles. + {"gemini-2.5-flash", 384, 384, "", 258}, + {"gemini-2.5-flash", 385, 385, "", 1290}, + {"gemini-2.5-flash", 769, 769, "", 1290}, + {"gemini-2.5-flash", 1536, 768, "", 1806}, + {"gemini-2.5-flash", 1537, 768, "", 2322}, + {"gemini-2.5-flash", 1920, 1080, "", 1806}, + {"gemini-2.5-flash", 2000, 2000, "", 2580}, + {"gemini-2.5-flash", 2304, 768, "", 2838}, + {"gemini-2.5-flash", 3000, 1000, "", 2838}, + {"gemini-2.5-flash", 2500, 2500, "", 3354}, + {"gemini-2.5-flash", 4000, 3000, "", 3354}, + // Gemini 3 bills by media resolution (documented values; measured 2-9% lower). + // An OpenAI detail level is not a Gemini media resolution. + {"gemini-3.7-flash", 1000, 1000, "", 1120}, + {"gemini-3.7-flash", 1000, 1000, "MEDIA_RESOLUTION_LOW", 280}, + {"gemini-3.7-flash", 1000, 1000, "media_resolution_medium", 560}, + {"gemini-3.7-flash", 3840, 2160, "MEDIA_RESOLUTION_ULTRA_HIGH", 2240}, + {"gemini-3.7-flash", 1000, 1000, "low", 1120}, + // Unknown sizes assume 1024x1024. + {"gpt-5.6-sol", 0, 0, "", 1229}, + {"claude-opus-4-7", 0, 0, "", 37*37 + 3}, + {"gemini-2.5-flash", -1, 0, "", 1290}, + {"deepseek-v4", 1000, 1000, "", 520}, + } { + got := ImageTokens(tc.model, Image{Width: tc.w, Height: tc.h, Detail: tc.detail}) + assert.Equal(t, tc.want, got, "%s %dx%d detail=%q", tc.model, tc.w, tc.h, tc.detail) + } +} + +// Image dimensions come from client-supplied headers; the cost stays bounded. +func TestImageTokensBoundForgedSizes(t *testing.T) { + const huge = 1 << 30 + assert.Equal(t, 36001, ImageTokens("gpt-5.6-sol", Image{Width: huge, Height: huge})) + // The largest square within the visual-token budget: 39x39 and 69x69 patches. + assert.Equal(t, 39*39+4, ImageTokens("claude-sonnet-4-6", Image{Width: huge, Height: huge})) + assert.Equal(t, 69*69+3, ImageTokens("claude-opus-5-5", Image{Width: huge, Height: huge})) + assert.Equal(t, 13*258, ImageTokens("gemini-2.5-pro", Image{Width: huge, Height: huge})) +} + +func TestImageNeedsSize(t *testing.T) { + assert.True(t, ImageNeedsSize("gpt-5.6-sol", "low")) + assert.True(t, ImageNeedsSize("gpt-4o", "high")) + assert.False(t, ImageNeedsSize("gpt-4o", "low")) + assert.True(t, ImageNeedsSize("claude-haiku-4-5", "")) + assert.True(t, ImageNeedsSize("gemini-2.5-flash", "")) + assert.False(t, ImageNeedsSize("gemini-3.7-flash", "")) + assert.False(t, ImageNeedsSize("deepseek-v4", "")) +} diff --git a/tokenkit/tokenizer.go b/tokenkit/tokenizer.go new file mode 100644 index 0000000000..adb1fee619 --- /dev/null +++ b/tokenkit/tokenizer.go @@ -0,0 +1,102 @@ +package tokenkit + +import ( + "strings" + "sync" + "unicode" + "unicode/utf8" + + "github.com/tiktoken-go/tokenizer" + "github.com/tiktoken-go/tokenizer/codec" +) + +var ( + o200kBase = sync.OnceValue(func() tokenizer.Codec { return codec.NewO200kBase() }) + cl100kBase = sync.OnceValue(func() tokenizer.Codec { return codec.NewCl100kBase() }) +) + +// maxRunBytes bounds a run of letters, whitespace, or symbols passed to the +// encoder in one piece. BPE cost grows with the square of a pre-tokenized +// piece: 100,000 Chinese characters without punctuation take about 20 s in one +// piece. Cutting longer runs keeps counting linear; ordinary text has no such +// runs, and a cut changes the count only inside the run (under 0.5% for CJK). +const maxRunBytes = 1024 + +// Count returns the token count of text for model. GPT and o-series models +// are counted with their tiktoken encoding; other models use Estimate. +func Count(model, text string) int { + if text == "" { + return 0 + } + model = strings.ToLower(model) + if ModelFamily(model) != FamilyOpenAI { + return Estimate(model, text) + } + encoder := o200kBase() + if usesCl100k(model) { + encoder = cl100kBase() + } + + total, start := 0, 0 + runClass, runStart := runNone, 0 + for i, r := range text { + class := runClassOf(r) + if class != runClass { + runClass, runStart = class, i + continue + } + if class != runNumber && i-runStart >= maxRunBytes { + n, _ := encoder.Count(text[start:i]) + total += n + start, runStart = i, i + } + } + n, _ := encoder.Count(text[start:]) + return total + n +} + +type runClass int + +const ( + runNone runClass = iota + runLetter + runNumber + runSpace + runSymbol +) + +// runClassOf groups runes the way tiktoken's pre-tokenizer joins them into +// pieces. Digits are split into groups of three by the encoder itself. +func runClassOf(r rune) runClass { + if r < utf8.RuneSelf { + switch { + case 'a' <= r && r <= 'z' || 'A' <= r && r <= 'Z': + return runLetter + case '0' <= r && r <= '9': + return runNumber + case r == ' ' || '\t' <= r && r <= '\r': + return runSpace + default: + return runSymbol + } + } + switch { + case unicode.IsLetter(r) || unicode.IsMark(r): + return runLetter + case unicode.IsNumber(r): + return runNumber + case unicode.IsSpace(r): + return runSpace + default: + return runSymbol + } +} + +// usesCl100k reports whether an OpenAI model predates o200k_base: GPT-4 (but +// not GPT-4o, GPT-4.1, or GPT-4.5) and GPT-3.5. Every later model, including +// GPT-5, GPT-6, and the o-series, uses o200k_base. +func usesCl100k(model string) bool { + model = strings.TrimPrefix(model, "ft:") + return model == "gpt-4" || strings.HasPrefix(model, "gpt-4-") || + strings.HasPrefix(model, "gpt-3.5") || strings.HasPrefix(model, "gpt-35") +} diff --git a/tokenkit/tokenkit.go b/tokenkit/tokenkit.go new file mode 100644 index 0000000000..8544090a28 --- /dev/null +++ b/tokenkit/tokenkit.go @@ -0,0 +1,245 @@ +// Package tokenkit estimates LLM token counts for text and images without +// calling an upstream API. +// +// Text is counted exactly with tiktoken for OpenAI GPT and o-series models and +// estimated with per-tokenizer weights for Claude, Gemini, and other models. +// Image costs follow each provider's published sizing rules. The weights were +// calibrated against upstream usage on real agent traffic (Claude Code and +// Codex sessions). +package tokenkit + +import ( + "math" + "strconv" + "strings" + "unicode" +) + +// Family is a tokenizer family. Models in one family tokenize text alike. +type Family string + +const ( + FamilyOpenAI Family = "openai" // GPT and o-series models + FamilyClaude Family = "claude" // Claude models before Claude Opus 4.7 + FamilyClaude47 Family = "claude4.7" // Claude Opus 4.7 and later, including Claude 5, Fable, and Mythos + FamilyGemini Family = "gemini" // Gemini models + FamilyUnknown Family = "unknown" // everything else +) + +// ModelFamily classifies a model name, case-insensitively. Vendor prefixes and +// date suffixes are allowed (for example us.anthropic.claude-opus-4-7-v1:0). +func ModelFamily(model string) Family { + model = strings.ToLower(model) + switch { + case strings.Contains(model, "gemini"): + return FamilyGemini + case strings.Contains(model, "claude"): + if claudeUsesOpus47Tokenizer(model) { + return FamilyClaude47 + } + return FamilyClaude + case isOpenAITextModel(model): + return FamilyOpenAI + default: + return FamilyUnknown + } +} + +func isOpenAITextModel(model string) bool { + for _, marker := range []string{"gpt-", "o1", "o3", "o4", "chatgpt"} { + if strings.Contains(model, marker) { + return true + } + } + return false +} + +// claudeUsesOpus47Tokenizer reports whether a lower-cased Claude model name uses +// the tokenizer and high-resolution image tier introduced with Claude Opus 4.7. +// Claude 4.7 and later use them; names without a version are treated as new. +func claudeUsesOpus47Tokenizer(model string) bool { + major, minor, ok := modelVersion(model, "claude-") + if !ok { + return true + } + return major > 4 || (major == 4 && minor >= 7) +} + +// modelVersion reads the first one- or two-digit version numbers after prefix, +// skipping family names (claude-opus-4-7, claude-3-5-sonnet, gemini-3.7-flash). +// A date such as 20250929 ends the version. ok is false when no version is found. +func modelVersion(model, prefix string) (major, minor int, ok bool) { + _, name, found := strings.Cut(model, prefix) + if !found { + return 0, 0, false + } + version := make([]int, 0, 2) + for part := range strings.FieldsFuncSeq(name, func(r rune) bool { return r == '-' || r == '.' }) { + digits := len(part) - len(strings.TrimLeft(part, "0123456789")) + if digits == 0 || digits > 2 { + if len(version) > 0 { + break + } + continue + } + n, _ := strconv.Atoi(part[:digits]) + version = append(version, n) + if len(version) == 2 || digits < len(part) { + break + } + } + switch len(version) { + case 0: + return 0, 0, false + case 1: + return version[0], 0, true + default: + return version[0], version[1], true + } +} + +// weights are the estimated tokens per text feature. +type weights struct { + Word float64 // Latin word (each run of letters) + Number float64 // each run of digits + CJK float64 // each Chinese, Japanese, or Korean character + Symbol float64 // each ordinary punctuation mark + MathSymbol float64 // each math symbol (∑, ∫, ∂, √, ...) + URLDelim float64 // each URL delimiter (/ : ? & = ; # %) + AtSign float64 // each @, which splits words + Emoji float64 // each emoji + Newline float64 // each newline or tab + Space float64 // each other whitespace character +} + +// Word, Number, CJK, Symbol, URLDelim, Newline, and Space are calibrated +// against upstream counts on agent traffic. MathSymbol, AtSign, and Emoji make +// up under 0.3% of that traffic and keep their earlier values. FamilyUnknown +// keeps the earlier OpenAI weights; no other vendor was calibrated. +var familyWeights = map[Family]weights{ + FamilyGemini: { + Word: 0.89, Number: 2.42, CJK: 0.93, Symbol: 0.65, MathSymbol: 1.05, URLDelim: 1.3, AtSign: 2.5, Emoji: 1.08, Newline: 1.28, Space: 0, + }, + FamilyClaude: { + Word: 1.24, Number: 1.53, CJK: 1.38, Symbol: 0.66, MathSymbol: 4.52, URLDelim: 0.74, AtSign: 2.82, Emoji: 2.6, Newline: 1.19, Space: 0.14, + }, + FamilyClaude47: { + Word: 1.52, Number: 1.77, CJK: 1.4, Symbol: 0.89, MathSymbol: 4.52, URLDelim: 1.45, AtSign: 2.82, Emoji: 2.6, Newline: 1.38, Space: 0.21, + }, + FamilyOpenAI: { + Word: 0.75, Number: 1.86, CJK: 0.96, Symbol: 0.6, MathSymbol: 2.68, URLDelim: 0.77, AtSign: 2.0, Emoji: 2.12, Newline: 0.95, Space: 0.17, + }, + FamilyUnknown: { + Word: 1.02, Number: 1.55, CJK: 0.85, Symbol: 0.4, MathSymbol: 2.68, URLDelim: 1.0, AtSign: 2.0, Emoji: 2.12, Newline: 0.5, Space: 0.42, + }, +} + +// Estimate returns the estimated token count of text for model, using the +// weights of the model's family. It never runs a tokenizer; use Count for an +// exact count where one is available. +func Estimate(model, text string) int { + if text == "" { + return 0 + } + return EstimateFamily(ModelFamily(model), text) +} + +// EstimateFamily returns the estimated token count of text for a tokenizer +// family. Unrecognized families use the FamilyUnknown weights. +func EstimateFamily(family Family, text string) int { + w, ok := familyWeights[family] + if !ok { + w = familyWeights[FamilyUnknown] + } + var count float64 + + type wordType int + const ( + none wordType = iota + latin + number + ) + current := none + + for _, r := range text { + if unicode.IsSpace(r) { + current = none + if r == '\n' || r == '\t' { + count += w.Newline + } else { + count += w.Space + } + continue + } + + if isCJK(r) { + current = none + count += w.CJK + continue + } + + if isEmoji(r) { + current = none + count += w.Emoji + continue + } + + // A run of letters or digits counts once; switching between letters + // and digits starts a new run. + if unicode.IsLetter(r) || unicode.IsNumber(r) { + next := latin + if unicode.IsNumber(r) { + next = number + } + if current != next { + if next == number { + count += w.Number + } else { + count += w.Word + } + current = next + } + continue + } + + current = none + switch { + case isMathSymbol(r): + count += w.MathSymbol + case r == '@': + count += w.AtSign + case isURLDelim(r): + count += w.URLDelim + default: + count += w.Symbol + } + } + + return int(math.Ceil(count)) +} + +func isCJK(r rune) bool { + return unicode.Is(unicode.Han, r) || + (r >= 0x3040 && r <= 0x30FF) || // Japanese kana + (r >= 0xAC00 && r <= 0xD7A3) // Korean Hangul syllables +} + +func isEmoji(r rune) bool { + return (r >= 0x1F300 && r <= 0x1F9FF) || // pictographs, emoticons, supplemental symbols + (r >= 0x2600 && r <= 0x26FF) || // miscellaneous symbols + (r >= 0x2700 && r <= 0x27BF) || // dingbats + (r >= 0x1FA00 && r <= 0x1FAFF) // symbols and pictographs extended-A +} + +func isMathSymbol(r rune) bool { + if strings.ContainsRune("∑∫∂√∞≤≥≠≈±×÷∈∉∋∌⊂⊃⊆⊇∪∩∧∨¬∀∃∄∅∆∇∝∟∠∡∢°′″‴⁺⁻⁼⁽⁾ⁿ₀₁₂₃₄₅₆₇₈₉₊₋₌₍₎²³¹⁴⁵⁶⁷⁸⁹⁰", r) { + return true + } + return (r >= 0x2200 && r <= 0x22FF) || // mathematical operators + (r >= 0x2A00 && r <= 0x2AFF) || // supplemental mathematical operators + (r >= 0x1D400 && r <= 0x1D7FF) // mathematical alphanumeric symbols +} + +func isURLDelim(r rune) bool { + return strings.ContainsRune("/:?&=;#%", r) +} diff --git a/tokenkit/tokenkit_test.go b/tokenkit/tokenkit_test.go new file mode 100644 index 0000000000..23eb9a90bd --- /dev/null +++ b/tokenkit/tokenkit_test.go @@ -0,0 +1,83 @@ +package tokenkit + +import ( + "math/rand" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "github.com/tiktoken-go/tokenizer/codec" +) + +func TestModelFamily(t *testing.T) { + for _, tc := range []struct { + model string + family Family + }{ + {"claude-3-5-sonnet-20240620", FamilyClaude}, + {"claude-sonnet-4-20250514", FamilyClaude}, + {"claude-sonnet-4-5-20250929", FamilyClaude}, + {"claude-haiku-4-5", FamilyClaude}, + {"claude-sonnet-4-6", FamilyClaude}, + // Claude Opus 4.7 introduced the tokenizer every later Claude model uses. + {"claude-opus-4-7", FamilyClaude47}, + {"claude-opus-4-7@20260416", FamilyClaude47}, + {"us.anthropic.claude-opus-4-7-v1:0", FamilyClaude47}, + {"claude-sonnet-5", FamilyClaude47}, + {"claude-fable-5-1", FamilyClaude47}, + {"Claude-Opus-5-5", FamilyClaude47}, + {"claude-mythos-preview", FamilyClaude47}, + {"gemini-3.7-flash", FamilyGemini}, + {"gpt-5.6-sol", FamilyOpenAI}, + {"gpt-6-astra", FamilyOpenAI}, + {"o3-mini", FamilyOpenAI}, + {"deepseek-v4", FamilyUnknown}, + } { + assert.Equal(t, tc.family, ModelFamily(tc.model), tc.model) + } +} + +func TestEstimateUsesFamilyWeights(t *testing.T) { + const text = "func main() {\n\tfmt.Println(\"统计 token\", 42)\n}\nSee https://example.com/a?b=1" + require.NotEqual(t, EstimateFamily(FamilyClaude, text), EstimateFamily(FamilyClaude47, text)) + + assert.Equal(t, EstimateFamily(FamilyClaude, text), Estimate("claude-haiku-4-5", text)) + assert.Equal(t, EstimateFamily(FamilyClaude47, text), Estimate("claude-sonnet-5", text)) + assert.Equal(t, EstimateFamily(FamilyUnknown, text), Estimate("deepseek-v4", text)) + assert.Equal(t, EstimateFamily(FamilyUnknown, text), EstimateFamily("other", text)) + assert.Zero(t, Estimate("claude-sonnet-5", "")) +} + +func TestCountUsesTiktokenForOpenAIModels(t *testing.T) { + const text = "大语言模型按 token 计费；Large language models bill by token. 🙂" + o200k, _ := codec.NewO200kBase().Count(text) + cl100k, _ := codec.NewCl100kBase().Count(text) + require.NotEqual(t, o200k, cl100k) + + assert.Equal(t, o200k, Count("gpt-5.6-sol", text)) + assert.Equal(t, o200k, Count("gpt-6-astra", text)) + assert.Equal(t, o200k, Count("gpt-4o-mini", text)) + assert.Equal(t, cl100k, Count("gpt-4", text)) + assert.Equal(t, cl100k, Count("gpt-3.5-turbo", text)) + assert.Equal(t, Estimate("claude-sonnet-5", text), Count("claude-sonnet-5", text)) +} + +// Long runs without separators are cut before encoding so BPE stays linear; +// ordinary text is counted exactly and cut runs stay within 0.5%. +func TestCountCutsOnlyLongRuns(t *testing.T) { + encoder := codec.NewO200kBase() + ordinary := strings.Repeat("func main() {\n\tfmt.Println(\"统计 token，按模型计费。\", 42)\n}\n| a | b |\n|---|---|\n", 200) + want, _ := encoder.Count(ordinary) + assert.Equal(t, want, Count("gpt-5.6-sol", ordinary)) + + rng := rand.New(rand.NewSource(1)) + var cjk strings.Builder + for range 5000 { + cjk.WriteRune(rune(0x4E00 + rng.Intn(0x9FA5-0x4E00))) + } + for _, text := range []string{cjk.String(), strings.Repeat("token", 2000), "a" + strings.Repeat(" ", 10000) + "b", strings.Repeat("=", 10000)} { + exact, _ := encoder.Count(text) + assert.InEpsilon(t, exact, Count("gpt-5.6-sol", text), 0.005) + } +}