175 lines · 5.6 KB · Go
RawDownload
package 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 := helps.TranslateRequestWithAPIKeyModelCompatibility(ctx, opts.Headers, e.cfg, from, to, baseModel, req.Payload, false, helps.APIKeyModelIsCompat(req)) 	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, helps.IsNativeCodexRequest(req.Payload, opts)) 	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}