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:
@@ -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",
|
||||
|
||||
Reference in New Issue
Block a user