mirror of
https://github.com/router-for-me/CLIProxyAPI.git
synced 2026-02-03 13:00:52 +08:00
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:
@@ -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.
|
||||
|
||||
@@ -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) {
|
||||
|
||||
Reference in New Issue
Block a user