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:
@@ -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