Add projects

backend/internal/runtime/executor/codex_executor_tokens.go

@@ -0,0 +1,175 @@
+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)
+
+	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
+}