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:
Luis Pater
2026-05-05 01:47:53 +08:00
parent 28b4b19e7e
commit ba5d8ca733
8 changed files with 125 additions and 6 deletions
+32
View File
@@ -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)