feat(executor): add CountTokens support across all executors

- Introduced `CountTokens` method to Codex, Claude, Gemini, Qwen, OpenAI-compatible, and other executors.
- Implemented `ExecuteCount` in `AuthManager` for token counting via provider round-robin.
- Updated handlers to leverage `ExecuteCountWithAuthManager` for streamlined token counting.
- Added fallback and error handling logic for token counting requests.
This commit is contained in:
Luis Pater
2025-09-23 02:27:51 +08:00
parent d32fc0400e
commit ac59023abb
11 changed files with 210 additions and 28 deletions

View File

@@ -17,9 +17,6 @@ import (
. "github.com/router-for-me/CLIProxyAPI/v6/internal/constant"
"github.com/router-for-me/CLIProxyAPI/v6/internal/interfaces"
"github.com/router-for-me/CLIProxyAPI/v6/internal/registry"
"github.com/router-for-me/CLIProxyAPI/v6/internal/runtime/executor"
coreexecutor "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/executor"
sdktranslator "github.com/router-for-me/CLIProxyAPI/v6/sdk/translator"
)
// GeminiAPIHandler contains the handlers for Gemini API endpoints.
@@ -226,35 +223,16 @@ func (h *GeminiAPIHandler) handleStreamGenerateContent(c *gin.Context, modelName
// - rawJSON: The raw JSON request body containing the content to count
func (h *GeminiAPIHandler) handleCountTokens(c *gin.Context, modelName string, rawJSON []byte) {
c.Header("Content-Type", "application/json")
alt := h.GetAlt(c)
cliCtx, cliCancel := h.GetContextWithCancel(h, c, context.Background())
defer func() { cliCancel() }()
// Execute via AuthManager with action=countTokens
req := coreexecutor.Request{
Model: modelName,
Payload: rawJSON,
Metadata: map[string]any{
"action": "countTokens",
},
}
opts := coreexecutor.Options{
Stream: false,
Alt: alt,
OriginalRequest: rawJSON,
SourceFormat: sdktranslator.FromString(h.HandlerType()),
}
resp, err := h.AuthManager.Execute(cliCtx, []string{"gemini"}, req, opts)
if err != nil {
if msg, ok := executor.UnwrapError(err); ok {
h.WriteErrorResponse(c, msg)
return
}
h.WriteErrorResponse(c, &interfaces.ErrorMessage{StatusCode: http.StatusInternalServerError, Error: err})
resp, errMsg := h.ExecuteCountWithAuthManager(cliCtx, h.HandlerType(), modelName, rawJSON, alt)
if errMsg != nil {
h.WriteErrorResponse(c, errMsg)
cliCancel(errMsg.Error)
return
}
_, _ = c.Writer.Write(resp.Payload)
_, _ = c.Writer.Write(resp)
cliCancel()
}
// handleGenerateContent handles non-streaming content generation requests for Gemini models.

View File

@@ -158,6 +158,33 @@ func (h *BaseAPIHandler) ExecuteWithAuthManager(ctx context.Context, handlerType
return cloneBytes(resp.Payload), nil
}
// ExecuteCountWithAuthManager executes a non-streaming request via the core auth manager.
// This path is the only supported execution route.
func (h *BaseAPIHandler) ExecuteCountWithAuthManager(ctx context.Context, handlerType, modelName string, rawJSON []byte, alt string) ([]byte, *interfaces.ErrorMessage) {
providers := util.GetProviderName(modelName, h.Cfg)
if len(providers) == 0 {
return nil, &interfaces.ErrorMessage{StatusCode: http.StatusBadRequest, Error: fmt.Errorf("unknown provider for model %s", modelName)}
}
req := coreexecutor.Request{
Model: modelName,
Payload: cloneBytes(rawJSON),
}
opts := coreexecutor.Options{
Stream: false,
Alt: alt,
OriginalRequest: cloneBytes(rawJSON),
SourceFormat: sdktranslator.FromString(handlerType),
}
resp, err := h.AuthManager.ExecuteCount(ctx, providers, req, opts)
if err != nil {
if msg, ok := executor.UnwrapError(err); ok {
return nil, msg
}
return nil, &interfaces.ErrorMessage{StatusCode: http.StatusInternalServerError, Error: err}
}
return cloneBytes(resp.Payload), nil
}
// ExecuteStreamWithAuthManager executes a streaming request via the core auth manager.
// This path is the only supported execution route.
func (h *BaseAPIHandler) ExecuteStreamWithAuthManager(ctx context.Context, handlerType, modelName string, rawJSON []byte, alt string) (<-chan []byte, <-chan *interfaces.ErrorMessage) {