mirror of
https://github.com/router-for-me/CLIProxyAPI.git
synced 2026-06-07 22:48:33 +08:00
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:
@@ -33,6 +33,10 @@ func (p *usageQueuePlugin) HandleUsage(ctx context.Context, record coreusage.Rec
|
||||
if modelName == "" {
|
||||
modelName = "unknown"
|
||||
}
|
||||
aliasName := strings.TrimSpace(record.Alias)
|
||||
if aliasName == "" {
|
||||
aliasName = modelName
|
||||
}
|
||||
provider := strings.TrimSpace(record.Provider)
|
||||
if provider == "" {
|
||||
provider = "unknown"
|
||||
@@ -76,6 +80,7 @@ func (p *usageQueuePlugin) HandleUsage(ctx context.Context, record coreusage.Rec
|
||||
requestDetail: detail,
|
||||
Provider: provider,
|
||||
Model: modelName,
|
||||
Alias: aliasName,
|
||||
Endpoint: resolveEndpoint(ctx),
|
||||
AuthType: authType,
|
||||
APIKey: apiKey,
|
||||
@@ -91,6 +96,7 @@ type queuedUsageDetail struct {
|
||||
requestDetail
|
||||
Provider string `json:"provider"`
|
||||
Model string `json:"model"`
|
||||
Alias string `json:"alias"`
|
||||
Endpoint string `json:"endpoint"`
|
||||
AuthType string `json:"auth_type"`
|
||||
APIKey string `json:"api_key"`
|
||||
|
||||
@@ -24,6 +24,7 @@ func TestUsageQueuePluginPayloadIncludesStableFieldsAndSuccess(t *testing.T) {
|
||||
plugin.HandleUsage(ctx, coreusage.Record{
|
||||
Provider: "openai",
|
||||
Model: "gpt-5.4",
|
||||
Alias: "client-gpt",
|
||||
APIKey: "test-key",
|
||||
AuthIndex: "0",
|
||||
AuthType: "apikey",
|
||||
@@ -40,6 +41,7 @@ func TestUsageQueuePluginPayloadIncludesStableFieldsAndSuccess(t *testing.T) {
|
||||
payload := popSinglePayload(t)
|
||||
requireStringField(t, payload, "provider", "openai")
|
||||
requireStringField(t, payload, "model", "gpt-5.4")
|
||||
requireStringField(t, payload, "alias", "client-gpt")
|
||||
requireStringField(t, payload, "endpoint", "POST /v1/chat/completions")
|
||||
requireStringField(t, payload, "auth_type", "apikey")
|
||||
requireStringField(t, payload, "request_id", "ctx-request-id")
|
||||
@@ -58,6 +60,7 @@ func TestUsageQueuePluginPayloadIncludesStableFieldsAndFailureAndGinRequestID(t
|
||||
plugin.HandleUsage(ctx, coreusage.Record{
|
||||
Provider: "openai",
|
||||
Model: "gpt-5.4-mini",
|
||||
Alias: "client-mini",
|
||||
APIKey: "test-key",
|
||||
AuthIndex: "0",
|
||||
AuthType: "apikey",
|
||||
@@ -74,6 +77,7 @@ func TestUsageQueuePluginPayloadIncludesStableFieldsAndFailureAndGinRequestID(t
|
||||
payload := popSinglePayload(t)
|
||||
requireStringField(t, payload, "provider", "openai")
|
||||
requireStringField(t, payload, "model", "gpt-5.4-mini")
|
||||
requireStringField(t, payload, "alias", "client-mini")
|
||||
requireStringField(t, payload, "endpoint", "GET /v1/responses")
|
||||
requireStringField(t, payload, "auth_type", "apikey")
|
||||
requireStringField(t, payload, "request_id", "gin-request-id")
|
||||
@@ -102,6 +106,7 @@ func TestUsageQueuePluginAsyncIgnoresRecycledGinContext(t *testing.T) {
|
||||
mgr.Publish(ctx, coreusage.Record{
|
||||
Provider: "openai",
|
||||
Model: "gpt-5.4",
|
||||
Alias: "client-gpt",
|
||||
APIKey: "test-key",
|
||||
AuthIndex: "0",
|
||||
AuthType: "apikey",
|
||||
@@ -117,6 +122,7 @@ func TestUsageQueuePluginAsyncIgnoresRecycledGinContext(t *testing.T) {
|
||||
|
||||
payload := waitForSinglePayload(t, 2*time.Second)
|
||||
requireStringField(t, payload, "endpoint", "POST /v1/chat/completions")
|
||||
requireStringField(t, payload, "alias", "client-gpt")
|
||||
requireStringField(t, payload, "request_id", "ctx-request-id")
|
||||
requireBoolField(t, payload, "failed", true)
|
||||
})
|
||||
|
||||
@@ -18,6 +18,7 @@ import (
|
||||
type UsageReporter struct {
|
||||
provider string
|
||||
model string
|
||||
alias string
|
||||
authID string
|
||||
authIndex string
|
||||
authType string
|
||||
@@ -29,9 +30,14 @@ type UsageReporter struct {
|
||||
|
||||
func NewUsageReporter(ctx context.Context, provider, model string, auth *cliproxyauth.Auth) *UsageReporter {
|
||||
apiKey := APIKeyFromContext(ctx)
|
||||
alias := usage.RequestedModelAliasFromContext(ctx)
|
||||
if alias == "" {
|
||||
alias = model
|
||||
}
|
||||
reporter := &UsageReporter{
|
||||
provider: provider,
|
||||
model: model,
|
||||
alias: strings.TrimSpace(alias),
|
||||
requestedAt: time.Now(),
|
||||
apiKey: apiKey,
|
||||
source: resolveUsageSource(auth, apiKey),
|
||||
@@ -139,6 +145,7 @@ func (r *UsageReporter) buildRecordForModel(model string, detail usage.Detail, f
|
||||
return usage.Record{
|
||||
Provider: r.provider,
|
||||
Model: model,
|
||||
Alias: r.alias,
|
||||
Source: r.source,
|
||||
APIKey: r.apiKey,
|
||||
AuthID: r.authID,
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
package helps
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
@@ -107,6 +108,19 @@ func TestUsageReporterBuildRecordIncludesLatency(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestUsageReporterBuildRecordIncludesRequestedModelAlias(t *testing.T) {
|
||||
ctx := usage.WithRequestedModelAlias(context.Background(), "client-gpt")
|
||||
reporter := NewUsageReporter(ctx, "openai", "gpt-5.4", nil)
|
||||
|
||||
record := reporter.buildRecord(usage.Detail{TotalTokens: 3}, false)
|
||||
if record.Model != "gpt-5.4" {
|
||||
t.Fatalf("model = %q, want %q", record.Model, "gpt-5.4")
|
||||
}
|
||||
if record.Alias != "client-gpt" {
|
||||
t.Fatalf("alias = %q, want %q", record.Alias, "client-gpt")
|
||||
}
|
||||
}
|
||||
|
||||
func TestUsageReporterBuildAdditionalModelRecordSkipsZeroTokens(t *testing.T) {
|
||||
reporter := &UsageReporter{
|
||||
provider: "codex",
|
||||
|
||||
@@ -539,7 +539,7 @@ func (h *BaseAPIHandler) ExecuteWithAuthManager(ctx context.Context, handlerType
|
||||
return nil, nil, errMsg
|
||||
}
|
||||
reqMeta := requestExecutionMetadata(ctx)
|
||||
reqMeta[coreexecutor.RequestedModelMetadataKey] = normalizedModel
|
||||
reqMeta[coreexecutor.RequestedModelMetadataKey] = modelName
|
||||
payload := rawJSON
|
||||
if len(payload) == 0 {
|
||||
payload = nil
|
||||
@@ -587,7 +587,7 @@ func (h *BaseAPIHandler) ExecuteCountWithAuthManager(ctx context.Context, handle
|
||||
return nil, nil, errMsg
|
||||
}
|
||||
reqMeta := requestExecutionMetadata(ctx)
|
||||
reqMeta[coreexecutor.RequestedModelMetadataKey] = normalizedModel
|
||||
reqMeta[coreexecutor.RequestedModelMetadataKey] = modelName
|
||||
payload := rawJSON
|
||||
if len(payload) == 0 {
|
||||
payload = nil
|
||||
@@ -639,7 +639,7 @@ func (h *BaseAPIHandler) ExecuteStreamWithAuthManager(ctx context.Context, handl
|
||||
return nil, nil, errChan
|
||||
}
|
||||
reqMeta := requestExecutionMetadata(ctx)
|
||||
reqMeta[coreexecutor.RequestedModelMetadataKey] = normalizedModel
|
||||
reqMeta[coreexecutor.RequestedModelMetadataKey] = modelName
|
||||
payload := rawJSON
|
||||
if len(payload) == 0 {
|
||||
payload = nil
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -2,6 +2,7 @@ package usage
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
@@ -12,6 +13,7 @@ import (
|
||||
type Record struct {
|
||||
Provider string
|
||||
Model string
|
||||
Alias string
|
||||
APIKey string
|
||||
AuthID string
|
||||
AuthIndex string
|
||||
@@ -32,6 +34,36 @@ type Detail struct {
|
||||
TotalTokens int64
|
||||
}
|
||||
|
||||
type requestedModelAliasContextKey struct{}
|
||||
|
||||
// WithRequestedModelAlias stores the client-requested model name for usage sinks.
|
||||
func WithRequestedModelAlias(ctx context.Context, alias string) context.Context {
|
||||
if ctx == nil {
|
||||
ctx = context.Background()
|
||||
}
|
||||
alias = strings.TrimSpace(alias)
|
||||
if alias == "" {
|
||||
return ctx
|
||||
}
|
||||
return context.WithValue(ctx, requestedModelAliasContextKey{}, alias)
|
||||
}
|
||||
|
||||
// RequestedModelAliasFromContext returns the client-requested model name stored in ctx.
|
||||
func RequestedModelAliasFromContext(ctx context.Context) string {
|
||||
if ctx == nil {
|
||||
return ""
|
||||
}
|
||||
raw := ctx.Value(requestedModelAliasContextKey{})
|
||||
switch value := raw.(type) {
|
||||
case string:
|
||||
return strings.TrimSpace(value)
|
||||
case []byte:
|
||||
return strings.TrimSpace(string(value))
|
||||
default:
|
||||
return ""
|
||||
}
|
||||
}
|
||||
|
||||
// Plugin consumes usage records emitted by the proxy runtime.
|
||||
type Plugin interface {
|
||||
HandleUsage(ctx context.Context, record Record)
|
||||
|
||||
Reference in New Issue
Block a user