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