mirror of
https://github.com/router-for-me/CLIProxyAPI.git
synced 2026-02-03 13:00:52 +08:00
feat(translator): add token counting functionality for Gemini, Claude, and CLI - Introduced `TokenCount` handling across various Codex translators (Gemini, Claude, CLI) with respective implementations. - Added utility methods for token counting and formatting responses. - Integrated `tiktoken-go/tokenizer` library for tokenization. - Updated CodexExecutor with token counting logic to support multiple models including GPT-5 variants. - Refined go.mod and go.sum to include new dependencies. feat(runtime): add token counting functionality across executors - Implemented token counting in OpenAICompatExecutor, QwenExecutor, and IFlowExecutor. - Added utilities for token counting and response formatting using `tiktoken-go/tokenizer`. - Integrated token counting into translators for Gemini, Claude, and Gemini CLI. - Enhanced multiple model support, including GPT-5 variants, for token counting. docs: update environment variable instructions for multi-model support - Added details for setting `ANTHROPIC_DEFAULT_OPUS_MODEL`, `ANTHROPIC_DEFAULT_SONNET_MODEL`, and `ANTHROPIC_DEFAULT_HAIKU_MODEL` for version 2.x.x. - Clarified usage of `ANTHROPIC_MODEL` and `ANTHROPIC_SMALL_FAST_MODEL` for version 1.x.x. - Expanded examples for setting environment variables across different models including Gemini, GPT-5, Claude, and Qwen3.
This commit is contained in:
@@ -20,6 +20,7 @@ import (
|
||||
log "github.com/sirupsen/logrus"
|
||||
"github.com/tidwall/gjson"
|
||||
"github.com/tidwall/sjson"
|
||||
"github.com/tiktoken-go/tokenizer"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/google/uuid"
|
||||
@@ -277,7 +278,180 @@ func (e *CodexExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Au
|
||||
}
|
||||
|
||||
func (e *CodexExecutor) CountTokens(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
|
||||
return cliproxyexecutor.Response{Payload: []byte{}}, fmt.Errorf("not implemented")
|
||||
from := opts.SourceFormat
|
||||
to := sdktranslator.FromString("codex")
|
||||
body := sdktranslator.TranslateRequest(from, to, req.Model, bytes.Clone(req.Payload), false)
|
||||
|
||||
modelForCounting := req.Model
|
||||
|
||||
if util.InArray([]string{"gpt-5", "gpt-5-minimal", "gpt-5-low", "gpt-5-medium", "gpt-5-high"}, req.Model) {
|
||||
modelForCounting = "gpt-5"
|
||||
body, _ = sjson.SetBytes(body, "model", "gpt-5")
|
||||
switch req.Model {
|
||||
case "gpt-5-minimal":
|
||||
body, _ = sjson.SetBytes(body, "reasoning.effort", "minimal")
|
||||
case "gpt-5-low":
|
||||
body, _ = sjson.SetBytes(body, "reasoning.effort", "low")
|
||||
case "gpt-5-medium":
|
||||
body, _ = sjson.SetBytes(body, "reasoning.effort", "medium")
|
||||
case "gpt-5-high":
|
||||
body, _ = sjson.SetBytes(body, "reasoning.effort", "high")
|
||||
default:
|
||||
body, _ = sjson.SetBytes(body, "reasoning.effort", "low")
|
||||
}
|
||||
} else if util.InArray([]string{"gpt-5-codex", "gpt-5-codex-low", "gpt-5-codex-medium", "gpt-5-codex-high"}, req.Model) {
|
||||
modelForCounting = "gpt-5"
|
||||
body, _ = sjson.SetBytes(body, "model", "gpt-5-codex")
|
||||
switch req.Model {
|
||||
case "gpt-5-codex-low":
|
||||
body, _ = sjson.SetBytes(body, "reasoning.effort", "low")
|
||||
case "gpt-5-codex-medium":
|
||||
body, _ = sjson.SetBytes(body, "reasoning.effort", "medium")
|
||||
case "gpt-5-codex-high":
|
||||
body, _ = sjson.SetBytes(body, "reasoning.effort", "high")
|
||||
default:
|
||||
body, _ = sjson.SetBytes(body, "reasoning.effort", "low")
|
||||
}
|
||||
}
|
||||
|
||||
body, _ = sjson.DeleteBytes(body, "previous_response_id")
|
||||
body, _ = sjson.SetBytes(body, "stream", false)
|
||||
|
||||
enc, err := tokenizerForCodexModel(modelForCounting)
|
||||
if err != nil {
|
||||
return cliproxyexecutor.Response{}, fmt.Errorf("codex executor: tokenizer init failed: %w", err)
|
||||
}
|
||||
|
||||
count, err := countCodexInputTokens(enc, body)
|
||||
if err != nil {
|
||||
return cliproxyexecutor.Response{}, fmt.Errorf("codex executor: token counting failed: %w", err)
|
||||
}
|
||||
|
||||
usageJSON := fmt.Sprintf(`{"response":{"usage":{"input_tokens":%d,"output_tokens":0,"total_tokens":%d}}}`, count, count)
|
||||
translated := sdktranslator.TranslateTokenCount(ctx, to, from, count, []byte(usageJSON))
|
||||
return cliproxyexecutor.Response{Payload: []byte(translated)}, nil
|
||||
}
|
||||
|
||||
func tokenizerForCodexModel(model string) (tokenizer.Codec, error) {
|
||||
sanitized := strings.ToLower(strings.TrimSpace(model))
|
||||
switch {
|
||||
case sanitized == "":
|
||||
return tokenizer.Get(tokenizer.Cl100kBase)
|
||||
case strings.HasPrefix(sanitized, "gpt-5"):
|
||||
return tokenizer.ForModel(tokenizer.GPT5)
|
||||
case strings.HasPrefix(sanitized, "gpt-4.1"):
|
||||
return tokenizer.ForModel(tokenizer.GPT41)
|
||||
case strings.HasPrefix(sanitized, "gpt-4o"):
|
||||
return tokenizer.ForModel(tokenizer.GPT4o)
|
||||
case strings.HasPrefix(sanitized, "gpt-4"):
|
||||
return tokenizer.ForModel(tokenizer.GPT4)
|
||||
case strings.HasPrefix(sanitized, "gpt-3.5"), strings.HasPrefix(sanitized, "gpt-3"):
|
||||
return tokenizer.ForModel(tokenizer.GPT35Turbo)
|
||||
default:
|
||||
return tokenizer.Get(tokenizer.Cl100kBase)
|
||||
}
|
||||
}
|
||||
|
||||
func countCodexInputTokens(enc tokenizer.Codec, body []byte) (int64, error) {
|
||||
if enc == nil {
|
||||
return 0, fmt.Errorf("encoder is nil")
|
||||
}
|
||||
if len(body) == 0 {
|
||||
return 0, nil
|
||||
}
|
||||
|
||||
root := gjson.ParseBytes(body)
|
||||
var segments []string
|
||||
|
||||
if inst := strings.TrimSpace(root.Get("instructions").String()); inst != "" {
|
||||
segments = append(segments, inst)
|
||||
}
|
||||
|
||||
inputItems := root.Get("input")
|
||||
if inputItems.IsArray() {
|
||||
arr := inputItems.Array()
|
||||
for i := range arr {
|
||||
item := arr[i]
|
||||
switch item.Get("type").String() {
|
||||
case "message":
|
||||
content := item.Get("content")
|
||||
if content.IsArray() {
|
||||
parts := content.Array()
|
||||
for j := range parts {
|
||||
part := parts[j]
|
||||
if text := strings.TrimSpace(part.Get("text").String()); text != "" {
|
||||
segments = append(segments, text)
|
||||
}
|
||||
}
|
||||
}
|
||||
case "function_call":
|
||||
if name := strings.TrimSpace(item.Get("name").String()); name != "" {
|
||||
segments = append(segments, name)
|
||||
}
|
||||
if args := strings.TrimSpace(item.Get("arguments").String()); args != "" {
|
||||
segments = append(segments, args)
|
||||
}
|
||||
case "function_call_output":
|
||||
if out := strings.TrimSpace(item.Get("output").String()); out != "" {
|
||||
segments = append(segments, out)
|
||||
}
|
||||
default:
|
||||
if text := strings.TrimSpace(item.Get("text").String()); text != "" {
|
||||
segments = append(segments, text)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
tools := root.Get("tools")
|
||||
if tools.IsArray() {
|
||||
tarr := tools.Array()
|
||||
for i := range tarr {
|
||||
tool := tarr[i]
|
||||
if name := strings.TrimSpace(tool.Get("name").String()); name != "" {
|
||||
segments = append(segments, name)
|
||||
}
|
||||
if desc := strings.TrimSpace(tool.Get("description").String()); desc != "" {
|
||||
segments = append(segments, desc)
|
||||
}
|
||||
if params := tool.Get("parameters"); params.Exists() {
|
||||
val := params.Raw
|
||||
if params.Type == gjson.String {
|
||||
val = params.String()
|
||||
}
|
||||
if trimmed := strings.TrimSpace(val); trimmed != "" {
|
||||
segments = append(segments, trimmed)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
textFormat := root.Get("text.format")
|
||||
if textFormat.Exists() {
|
||||
if name := strings.TrimSpace(textFormat.Get("name").String()); name != "" {
|
||||
segments = append(segments, name)
|
||||
}
|
||||
if schema := textFormat.Get("schema"); schema.Exists() {
|
||||
val := schema.Raw
|
||||
if schema.Type == gjson.String {
|
||||
val = schema.String()
|
||||
}
|
||||
if trimmed := strings.TrimSpace(val); trimmed != "" {
|
||||
segments = append(segments, trimmed)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
text := strings.Join(segments, "\n")
|
||||
if text == "" {
|
||||
return 0, nil
|
||||
}
|
||||
|
||||
count, err := enc.Count(text)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return int64(count), nil
|
||||
}
|
||||
|
||||
func (e *CodexExecutor) Refresh(ctx context.Context, auth *cliproxyauth.Auth) (*cliproxyauth.Auth, error) {
|
||||
|
||||
Reference in New Issue
Block a user