Something went wrong. Try again.
sloproxy
Something went wrong. Try again.
5.6 kB · 175 lines
Go
at commit e8e39526
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176package executor
import ( "context" "fmt" "strings"
"github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps" "github.com/router-for-me/CLIProxyAPI/v7/internal/thinking" cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator" "github.com/tidwall/gjson" "github.com/tidwall/sjson" "github.com/tiktoken-go/tokenizer")
func (e *CodexExecutor) CountTokens(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { baseModel := thinking.ParseSuffix(req.Model).ModelName
from := opts.SourceFormat responseFormat := cliproxyexecutor.ResponseFormatOrSource(opts) to := sdktranslator.FromString("codex") body := sdktranslator.TranslateRequest(from, to, baseModel, req.Payload, false)
body, err := helps.ApplyRequestThinking(body, req, opts, from.String(), to.String(), e.Identifier()) if err != nil { return cliproxyexecutor.Response{}, err }
body = helps.SetStringIfDifferent(body, "model", baseModel) body, _ = sjson.DeleteBytes(body, "previous_response_id") body, _ = sjson.DeleteBytes(body, "generate") body, _ = sjson.DeleteBytes(body, "prompt_cache_retention") body, _ = sjson.DeleteBytes(body, "safety_identifier") body, _ = sjson.DeleteBytes(body, "stream_options") body = helps.SetBoolIfDifferent(body, "stream", false) body = normalizeCodexInstructions(body)
enc, err := tokenizerForCodexModel(baseModel) 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, responseFormat, count, []byte(usageJSON)) return cliproxyexecutor.Response{Payload: 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}