feat(usage): add support for requested model alias handling
- Introduced methods for setting and retrieving model aliases in execution and usage contexts. - Enhanced `UsageReporter` and related structures to include client-requested aliases. - Updated tests to validate alias propagation and ensure correct usage reporting. - Adjusted metadata handling in CLIProxyAPI executors to address alias integration.
This commit is contained in:
@@ -22,6 +22,7 @@ import (
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/thinking"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/util"
|
||||
cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/executor"
|
||||
coreusage "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/usage"
|
||||
log "github.com/sirupsen/logrus"
|
||||
)
|
||||
|
||||
@@ -827,6 +828,7 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi
|
||||
if executor == nil {
|
||||
return nil, &Error{Code: "executor_not_found", Message: "executor not registered"}
|
||||
}
|
||||
ctx = contextWithRequestedModelAlias(ctx, opts, routeModel)
|
||||
var lastErr error
|
||||
for idx, execModel := range execModels {
|
||||
resultModel := m.stateModelForExecution(auth, routeModel, execModel, pooled)
|
||||
@@ -1319,6 +1321,7 @@ func (m *Manager) executeMixedOnce(ctx context.Context, providers []string, req
|
||||
execCtx = context.WithValue(execCtx, roundTripperContextKey{}, rt)
|
||||
execCtx = context.WithValue(execCtx, "cliproxy.roundtripper", rt)
|
||||
}
|
||||
execCtx = contextWithRequestedModelAlias(execCtx, opts, routeModel)
|
||||
|
||||
models, pooled := m.preparedExecutionModels(auth, routeModel)
|
||||
if len(models) == 0 {
|
||||
@@ -1397,6 +1400,7 @@ func (m *Manager) executeCountMixedOnce(ctx context.Context, providers []string,
|
||||
execCtx = context.WithValue(execCtx, roundTripperContextKey{}, rt)
|
||||
execCtx = context.WithValue(execCtx, "cliproxy.roundtripper", rt)
|
||||
}
|
||||
execCtx = contextWithRequestedModelAlias(execCtx, opts, routeModel)
|
||||
|
||||
models, pooled := m.preparedExecutionModels(auth, routeModel)
|
||||
if len(models) == 0 {
|
||||
@@ -1534,6 +1538,36 @@ func hasRequestedModelMetadata(meta map[string]any) bool {
|
||||
}
|
||||
}
|
||||
|
||||
func contextWithRequestedModelAlias(ctx context.Context, opts cliproxyexecutor.Options, fallback string) context.Context {
|
||||
alias := requestedModelAliasFromOptions(opts, fallback)
|
||||
return coreusage.WithRequestedModelAlias(ctx, alias)
|
||||
}
|
||||
|
||||
func requestedModelAliasFromOptions(opts cliproxyexecutor.Options, fallback string) string {
|
||||
fallback = strings.TrimSpace(fallback)
|
||||
if len(opts.Metadata) == 0 {
|
||||
return fallback
|
||||
}
|
||||
raw, ok := opts.Metadata[cliproxyexecutor.RequestedModelMetadataKey]
|
||||
if !ok || raw == nil {
|
||||
return fallback
|
||||
}
|
||||
switch value := raw.(type) {
|
||||
case string:
|
||||
if strings.TrimSpace(value) == "" {
|
||||
return fallback
|
||||
}
|
||||
return strings.TrimSpace(value)
|
||||
case []byte:
|
||||
if len(value) == 0 {
|
||||
return fallback
|
||||
}
|
||||
return strings.TrimSpace(string(value))
|
||||
default:
|
||||
return fallback
|
||||
}
|
||||
}
|
||||
|
||||
func pinnedAuthIDFromMetadata(meta map[string]any) string {
|
||||
if len(meta) == 0 {
|
||||
return ""
|
||||
@@ -3096,6 +3130,7 @@ func (m *Manager) tryAntigravityCreditsExecute(ctx context.Context, req cliproxy
|
||||
creditsCtx = context.WithValue(creditsCtx, "cliproxy.roundtripper", rt)
|
||||
}
|
||||
creditsOpts := ensureRequestedModelMetadata(opts, routeModel)
|
||||
creditsCtx = contextWithRequestedModelAlias(creditsCtx, creditsOpts, routeModel)
|
||||
publishSelectedAuthMetadata(creditsOpts.Metadata, c.auth.ID)
|
||||
models := m.executionModelCandidates(c.auth, routeModel)
|
||||
if len(models) == 0 {
|
||||
|
||||
@@ -10,20 +10,23 @@ import (
|
||||
internalconfig "github.com/router-for-me/CLIProxyAPI/v6/internal/config"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/registry"
|
||||
cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/executor"
|
||||
coreusage "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/usage"
|
||||
)
|
||||
|
||||
type aliasRoutingExecutor struct {
|
||||
id string
|
||||
|
||||
mu sync.Mutex
|
||||
executeModels []string
|
||||
mu sync.Mutex
|
||||
executeModels []string
|
||||
executeAliases []string
|
||||
}
|
||||
|
||||
func (e *aliasRoutingExecutor) Identifier() string { return e.id }
|
||||
|
||||
func (e *aliasRoutingExecutor) Execute(_ context.Context, _ *Auth, req cliproxyexecutor.Request, _ cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
|
||||
func (e *aliasRoutingExecutor) Execute(ctx context.Context, _ *Auth, req cliproxyexecutor.Request, _ cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
|
||||
e.mu.Lock()
|
||||
e.executeModels = append(e.executeModels, req.Model)
|
||||
e.executeAliases = append(e.executeAliases, coreusage.RequestedModelAliasFromContext(ctx))
|
||||
e.mu.Unlock()
|
||||
return cliproxyexecutor.Response{Payload: []byte(req.Model)}, nil
|
||||
}
|
||||
@@ -52,6 +55,14 @@ func (e *aliasRoutingExecutor) ExecuteModels() []string {
|
||||
return out
|
||||
}
|
||||
|
||||
func (e *aliasRoutingExecutor) ExecuteAliases() []string {
|
||||
e.mu.Lock()
|
||||
defer e.mu.Unlock()
|
||||
out := make([]string, len(e.executeAliases))
|
||||
copy(out, e.executeAliases)
|
||||
return out
|
||||
}
|
||||
|
||||
func TestManagerExecute_OAuthAliasBypassesBlockedRouteModel(t *testing.T) {
|
||||
const (
|
||||
provider = "antigravity"
|
||||
@@ -108,4 +119,12 @@ func TestManagerExecute_OAuthAliasBypassesBlockedRouteModel(t *testing.T) {
|
||||
if gotModels[0] != targetModel {
|
||||
t.Fatalf("execute model = %q, want %q", gotModels[0], targetModel)
|
||||
}
|
||||
|
||||
gotAliases := executor.ExecuteAliases()
|
||||
if len(gotAliases) != 1 {
|
||||
t.Fatalf("execute aliases len = %d, want 1", len(gotAliases))
|
||||
}
|
||||
if gotAliases[0] != routeModel {
|
||||
t.Fatalf("execute alias = %q, want %q", gotAliases[0], routeModel)
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user