mirror of
https://gh-proxy.org/https://github.com/klinsmaya/CLIProxyAPI.git
synced 2026-11-10 18:36:47 +08:00
feat(usage): add service tier tracking and defaults in usage reporting
- Introduced `service_tier` metadata key to capture client-requested service tiers. - Updated usage records, context propagation, and plugins to include service tier data. - Added default handling logic for cases where `service_tier` is absent. - Implemented tests for `service_tier` extraction, defaults, and updates across components.
This commit is contained in:
@@ -52,6 +52,10 @@ func (p *usageQueuePlugin) HandleUsage(ctx context.Context, record coreusage.Rec
|
||||
if reasoningEffort == "" {
|
||||
reasoningEffort = coreusage.ReasoningEffortFromContext(ctx)
|
||||
}
|
||||
serviceTier := strings.TrimSpace(record.ServiceTier)
|
||||
if serviceTier == "" {
|
||||
serviceTier = coreusage.ServiceTierFromContext(ctx)
|
||||
}
|
||||
|
||||
tokens := tokenStats{
|
||||
InputTokens: record.Detail.InputTokens,
|
||||
@@ -97,6 +101,7 @@ func (p *usageQueuePlugin) HandleUsage(ctx context.Context, record coreusage.Rec
|
||||
APIKey: apiKey,
|
||||
RequestID: requestID,
|
||||
ReasoningEffort: reasoningEffort,
|
||||
ServiceTier: serviceTier,
|
||||
})
|
||||
if err != nil {
|
||||
return
|
||||
@@ -114,6 +119,7 @@ type queuedUsageDetail struct {
|
||||
APIKey string `json:"api_key"`
|
||||
RequestID string `json:"request_id"`
|
||||
ReasoningEffort string `json:"reasoning_effort"`
|
||||
ServiceTier string `json:"service_tier"`
|
||||
}
|
||||
|
||||
type requestDetail struct {
|
||||
|
||||
@@ -33,6 +33,7 @@ func TestUsageQueuePluginPayloadIncludesStableFieldsAndSuccess(t *testing.T) {
|
||||
AuthType: "apikey",
|
||||
Source: "[email protected]",
|
||||
ReasoningEffort: "medium",
|
||||
ServiceTier: "priority",
|
||||
RequestedAt: time.Date(2026, 4, 25, 0, 0, 0, 0, time.UTC),
|
||||
Latency: 1500 * time.Millisecond,
|
||||
Detail: coreusage.Detail{
|
||||
@@ -53,6 +54,7 @@ func TestUsageQueuePluginPayloadIncludesStableFieldsAndSuccess(t *testing.T) {
|
||||
requireMissingField(t, payload, "user_api_key")
|
||||
requireStringField(t, payload, "request_id", "ctx-request-id")
|
||||
requireStringField(t, payload, "reasoning_effort", "medium")
|
||||
requireStringField(t, payload, "service_tier", "priority")
|
||||
requireHeaderField(t, payload, "response_headers", "X-Upstream-Request-Id", []string{"upstream-req-1"})
|
||||
requireHeaderField(t, payload, "response_headers", "Retry-After", []string{"30"})
|
||||
requireBoolField(t, payload, "failed", false)
|
||||
|
||||
@@ -30,6 +30,7 @@ type UsageReporter struct {
|
||||
apiKey string
|
||||
source string
|
||||
reasoning string
|
||||
serviceTier string
|
||||
requestedAt time.Time
|
||||
ttftMu sync.RWMutex
|
||||
ttft time.Duration
|
||||
@@ -53,6 +54,7 @@ func NewUsageReporter(ctx context.Context, provider, model string, auth *cliprox
|
||||
source: resolveUsageSource(auth, apiKey),
|
||||
authType: resolveUsageAuthType(auth),
|
||||
reasoning: usage.ReasoningEffortFromContext(ctx),
|
||||
serviceTier: usage.ServiceTierFromContext(ctx),
|
||||
}
|
||||
if auth != nil {
|
||||
reporter.authID = auth.ID
|
||||
@@ -78,6 +80,7 @@ func (r *UsageReporter) SetTranslatedReasoningEffort(payload []byte, format stri
|
||||
return
|
||||
}
|
||||
r.reasoning = thinking.ExtractTranslatedReasoningEffort(payload, format)
|
||||
r.serviceTier = extractServiceTierFromPayload(payload)
|
||||
}
|
||||
|
||||
func (r *UsageReporter) TrackHTTPClient(client *http.Client) *http.Client {
|
||||
@@ -239,6 +242,7 @@ func (r *UsageReporter) buildRecordForModel(model string, detail usage.Detail, f
|
||||
AuthIndex: r.authIndex,
|
||||
AuthType: r.authType,
|
||||
ReasoningEffort: r.reasoning,
|
||||
ServiceTier: r.serviceTier,
|
||||
RequestedAt: r.requestedAt,
|
||||
Latency: r.latency(),
|
||||
TTFT: r.ttftDuration(),
|
||||
@@ -248,6 +252,19 @@ func (r *UsageReporter) buildRecordForModel(model string, detail usage.Detail, f
|
||||
}
|
||||
}
|
||||
|
||||
func extractServiceTierFromPayload(payload []byte) string {
|
||||
if len(payload) == 0 {
|
||||
return usage.DefaultServiceTier
|
||||
}
|
||||
for _, path := range []string{"service_tier", "request.service_tier", "response.service_tier"} {
|
||||
serviceTier := strings.TrimSpace(gjson.GetBytes(payload, path).String())
|
||||
if serviceTier != "" {
|
||||
return serviceTier
|
||||
}
|
||||
}
|
||||
return usage.DefaultServiceTier
|
||||
}
|
||||
|
||||
func failFromErrors(errs ...error) usage.Failure {
|
||||
for _, err := range errs {
|
||||
if err == nil {
|
||||
|
||||
@@ -241,6 +241,39 @@ func TestUsageReporterBuildRecordIncludesReasoningEffort(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestUsageReporterBuildRecordIncludesServiceTier(t *testing.T) {
|
||||
ctx := usage.WithServiceTier(context.Background(), "priority")
|
||||
reporter := NewUsageReporter(ctx, "openai", "gpt-5.4", nil)
|
||||
|
||||
record := reporter.buildRecord(usage.Detail{TotalTokens: 3}, false)
|
||||
if record.ServiceTier != "priority" {
|
||||
t.Fatalf("service tier = %q, want %q", record.ServiceTier, "priority")
|
||||
}
|
||||
}
|
||||
|
||||
func TestUsageReporterSetTranslatedReasoningEffortUpdatesServiceTier(t *testing.T) {
|
||||
reporter := NewUsageReporter(context.Background(), "openai", "gpt-5.4", nil)
|
||||
|
||||
reporter.SetTranslatedReasoningEffort([]byte(`{"service_tier":"priority"}`), "openai")
|
||||
|
||||
record := reporter.buildRecord(usage.Detail{TotalTokens: 3}, false)
|
||||
if record.ServiceTier != "priority" {
|
||||
t.Fatalf("service tier = %q, want %q", record.ServiceTier, "priority")
|
||||
}
|
||||
}
|
||||
|
||||
func TestUsageReporterSetTranslatedReasoningEffortDefaultsServiceTierWhenRemoved(t *testing.T) {
|
||||
ctx := usage.WithServiceTier(context.Background(), "priority")
|
||||
reporter := NewUsageReporter(ctx, "openai", "gpt-5.4", nil)
|
||||
|
||||
reporter.SetTranslatedReasoningEffort([]byte(`{"model":"gpt-5.4"}`), "openai")
|
||||
|
||||
record := reporter.buildRecord(usage.Detail{TotalTokens: 3}, false)
|
||||
if record.ServiceTier != usage.DefaultServiceTier {
|
||||
t.Fatalf("service tier = %q, want %q", record.ServiceTier, usage.DefaultServiceTier)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUsageReporterBuildAdditionalModelRecordSkipsZeroTokens(t *testing.T) {
|
||||
reporter := &UsageReporter{
|
||||
provider: "codex",
|
||||
|
||||
@@ -20,8 +20,10 @@ import (
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/util"
|
||||
coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
|
||||
coreexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor"
|
||||
coreusage "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/usage"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/sdk/config"
|
||||
sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator"
|
||||
"github.com/tidwall/gjson"
|
||||
"golang.org/x/net/context"
|
||||
)
|
||||
|
||||
@@ -242,6 +244,21 @@ func setReasoningEffortMetadata(meta map[string]any, handlerType, model string,
|
||||
meta[coreexecutor.ReasoningEffortMetadataKey] = effort
|
||||
}
|
||||
|
||||
func setServiceTierMetadata(meta map[string]any, rawJSON []byte) {
|
||||
if meta == nil {
|
||||
return
|
||||
}
|
||||
serviceTier := coreusage.DefaultServiceTier
|
||||
node := gjson.GetBytes(rawJSON, "service_tier")
|
||||
if node.Exists() {
|
||||
value := strings.TrimSpace(node.String())
|
||||
if value != "" {
|
||||
serviceTier = value
|
||||
}
|
||||
}
|
||||
meta[coreexecutor.ServiceTierMetadataKey] = serviceTier
|
||||
}
|
||||
|
||||
// headersFromContext extracts the original HTTP request headers from the gin context
|
||||
// embedded in the provided context. This allows session affinity selectors to read
|
||||
// client headers like X-Amp-Thread-Id.
|
||||
@@ -562,6 +579,7 @@ func (h *BaseAPIHandler) executeWithAuthManager(ctx context.Context, handlerType
|
||||
reqMeta := requestExecutionMetadata(ctx)
|
||||
reqMeta[coreexecutor.RequestedModelMetadataKey] = modelName
|
||||
setReasoningEffortMetadata(reqMeta, handlerType, normalizedModel, rawJSON)
|
||||
setServiceTierMetadata(reqMeta, rawJSON)
|
||||
payload := rawJSON
|
||||
if len(payload) == 0 {
|
||||
payload = nil
|
||||
@@ -611,6 +629,7 @@ func (h *BaseAPIHandler) ExecuteCountWithAuthManager(ctx context.Context, handle
|
||||
reqMeta := requestExecutionMetadata(ctx)
|
||||
reqMeta[coreexecutor.RequestedModelMetadataKey] = modelName
|
||||
setReasoningEffortMetadata(reqMeta, handlerType, normalizedModel, rawJSON)
|
||||
setServiceTierMetadata(reqMeta, rawJSON)
|
||||
payload := rawJSON
|
||||
if len(payload) == 0 {
|
||||
payload = nil
|
||||
@@ -673,6 +692,7 @@ func (h *BaseAPIHandler) executeStreamWithAuthManager(ctx context.Context, handl
|
||||
reqMeta := requestExecutionMetadata(ctx)
|
||||
reqMeta[coreexecutor.RequestedModelMetadataKey] = modelName
|
||||
setReasoningEffortMetadata(reqMeta, handlerType, normalizedModel, rawJSON)
|
||||
setServiceTierMetadata(reqMeta, rawJSON)
|
||||
payload := rawJSON
|
||||
if len(payload) == 0 {
|
||||
payload = nil
|
||||
|
||||
@@ -38,3 +38,25 @@ func TestSetReasoningEffortMetadataSupportsOpenAIResponses(t *testing.T) {
|
||||
t.Fatalf("ReasoningEffortMetadataKey = %v, want %q", got, "medium")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSetServiceTierMetadataExtractsValue(t *testing.T) {
|
||||
meta := make(map[string]any)
|
||||
|
||||
setServiceTierMetadata(meta, []byte(`{"service_tier":"priority"}`))
|
||||
|
||||
gotServiceTier := meta[coreexecutor.ServiceTierMetadataKey]
|
||||
if gotServiceTier != "priority" {
|
||||
t.Fatalf("ServiceTierMetadataKey = %v, want %q", gotServiceTier, "priority")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSetServiceTierMetadataDefaultsWhenMissing(t *testing.T) {
|
||||
meta := make(map[string]any)
|
||||
|
||||
setServiceTierMetadata(meta, []byte(`{"model":"gpt-5.4"}`))
|
||||
|
||||
gotServiceTier := meta[coreexecutor.ServiceTierMetadataKey]
|
||||
if gotServiceTier != "default" {
|
||||
t.Fatalf("ServiceTierMetadataKey = %v, want %q", gotServiceTier, "default")
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1731,9 +1731,14 @@ func (m *Manager) prepareRequestAuth(ctx context.Context, executor ProviderExecu
|
||||
func contextWithRequestedModelAlias(ctx context.Context, opts cliproxyexecutor.Options, fallback string) context.Context {
|
||||
alias := requestedModelAliasFromOptions(opts, fallback)
|
||||
ctx = coreusage.WithRequestedModelAlias(ctx, alias)
|
||||
if effort := reasoningEffortFromOptions(opts); effort != "" {
|
||||
effort := reasoningEffortFromOptions(opts)
|
||||
if effort != "" {
|
||||
ctx = coreusage.WithReasoningEffort(ctx, effort)
|
||||
}
|
||||
serviceTier := serviceTierFromOptions(opts)
|
||||
if serviceTier != "" {
|
||||
ctx = coreusage.WithServiceTier(ctx, serviceTier)
|
||||
}
|
||||
return ctx
|
||||
}
|
||||
|
||||
@@ -1780,6 +1785,24 @@ func reasoningEffortFromOptions(opts cliproxyexecutor.Options) string {
|
||||
}
|
||||
}
|
||||
|
||||
func serviceTierFromOptions(opts cliproxyexecutor.Options) string {
|
||||
if len(opts.Metadata) == 0 {
|
||||
return ""
|
||||
}
|
||||
raw, ok := opts.Metadata[cliproxyexecutor.ServiceTierMetadataKey]
|
||||
if !ok || raw == nil {
|
||||
return ""
|
||||
}
|
||||
switch value := raw.(type) {
|
||||
case string:
|
||||
return strings.TrimSpace(value)
|
||||
case []byte:
|
||||
return strings.TrimSpace(string(value))
|
||||
default:
|
||||
return ""
|
||||
}
|
||||
}
|
||||
|
||||
func pinnedAuthIDFromMetadata(meta map[string]any) string {
|
||||
if len(meta) == 0 {
|
||||
return ""
|
||||
|
||||
@@ -13,6 +13,7 @@ func TestContextWithRequestedModelAliasIncludesReasoningEffort(t *testing.T) {
|
||||
Metadata: map[string]any{
|
||||
cliproxyexecutor.RequestedModelMetadataKey: "client-model",
|
||||
cliproxyexecutor.ReasoningEffortMetadataKey: "medium",
|
||||
cliproxyexecutor.ServiceTierMetadataKey: "priority",
|
||||
},
|
||||
}, "fallback-model")
|
||||
|
||||
@@ -22,4 +23,8 @@ func TestContextWithRequestedModelAliasIncludesReasoningEffort(t *testing.T) {
|
||||
if got := coreusage.ReasoningEffortFromContext(ctx); got != "medium" {
|
||||
t.Fatalf("reasoning effort = %q, want %q", got, "medium")
|
||||
}
|
||||
gotServiceTier := coreusage.ServiceTierFromContext(ctx)
|
||||
if gotServiceTier != "priority" {
|
||||
t.Fatalf("service tier = %q, want %q", gotServiceTier, "priority")
|
||||
}
|
||||
}
|
||||
|
||||
@@ -20,6 +20,9 @@ const DisallowFreeAuthMetadataKey = "disallow_free_auth"
|
||||
// ReasoningEffortMetadataKey stores the client-requested reasoning effort for usage logs.
|
||||
const ReasoningEffortMetadataKey = "reasoning_effort"
|
||||
|
||||
// ServiceTierMetadataKey stores the client-requested service tier for usage logs.
|
||||
const ServiceTierMetadataKey = "service_tier"
|
||||
|
||||
const (
|
||||
// PinnedAuthMetadataKey locks execution to a specific auth ID.
|
||||
PinnedAuthMetadataKey = "pinned_auth_id"
|
||||
|
||||
@@ -10,6 +10,9 @@ import (
|
||||
log "github.com/sirupsen/logrus"
|
||||
)
|
||||
|
||||
// DefaultServiceTier is used when a request does not specify service_tier.
|
||||
const DefaultServiceTier = "default"
|
||||
|
||||
// Record contains the usage statistics captured for a single provider request.
|
||||
type Record struct {
|
||||
Provider string
|
||||
@@ -22,12 +25,14 @@ type Record struct {
|
||||
Source string
|
||||
// ReasoningEffort stores the translated upstream thinking level for request event logs.
|
||||
ReasoningEffort string
|
||||
RequestedAt time.Time
|
||||
Latency time.Duration
|
||||
TTFT time.Duration
|
||||
Failed bool
|
||||
Fail Failure
|
||||
Detail Detail
|
||||
// ServiceTier stores the client-requested service tier for request event logs.
|
||||
ServiceTier string
|
||||
RequestedAt time.Time
|
||||
Latency time.Duration
|
||||
TTFT time.Duration
|
||||
Failed bool
|
||||
Fail Failure
|
||||
Detail Detail
|
||||
// ResponseHeaders stores a snapshot of upstream response headers for usage sinks.
|
||||
ResponseHeaders http.Header
|
||||
}
|
||||
@@ -51,6 +56,7 @@ type Detail struct {
|
||||
|
||||
type requestedModelAliasContextKey struct{}
|
||||
type reasoningEffortContextKey struct{}
|
||||
type serviceTierContextKey struct{}
|
||||
|
||||
// WithRequestedModelAlias stores the client-requested model name for usage sinks.
|
||||
func WithRequestedModelAlias(ctx context.Context, alias string) context.Context {
|
||||
@@ -108,6 +114,42 @@ func ReasoningEffortFromContext(ctx context.Context) string {
|
||||
}
|
||||
}
|
||||
|
||||
// WithServiceTier stores the client-requested service tier for usage sinks.
|
||||
func WithServiceTier(ctx context.Context, tier string) context.Context {
|
||||
if ctx == nil {
|
||||
ctx = context.Background()
|
||||
}
|
||||
tier = strings.TrimSpace(tier)
|
||||
if tier == "" {
|
||||
tier = DefaultServiceTier
|
||||
}
|
||||
return context.WithValue(ctx, serviceTierContextKey{}, tier)
|
||||
}
|
||||
|
||||
// ServiceTierFromContext returns the client-requested service tier stored in ctx.
|
||||
func ServiceTierFromContext(ctx context.Context) string {
|
||||
if ctx == nil {
|
||||
return DefaultServiceTier
|
||||
}
|
||||
raw := ctx.Value(serviceTierContextKey{})
|
||||
switch value := raw.(type) {
|
||||
case string:
|
||||
tier := strings.TrimSpace(value)
|
||||
if tier == "" {
|
||||
return DefaultServiceTier
|
||||
}
|
||||
return tier
|
||||
case []byte:
|
||||
tier := strings.TrimSpace(string(value))
|
||||
if tier == "" {
|
||||
return DefaultServiceTier
|
||||
}
|
||||
return tier
|
||||
default:
|
||||
return DefaultServiceTier
|
||||
}
|
||||
}
|
||||
|
||||
// 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