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 {
|
||||
|
||||
Reference in New Issue
Block a user