Something went wrong. Try again.
sloproxy
Something went wrong. Try again.
6.8 kB · 212 lines
Go
at commit 291cfb87
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213package redisqueue
import ( "context" "encoding/json" "net/http" "strings" "time"
internallogging "github.com/router-for-me/CLIProxyAPI/v7/internal/logging" coreusage "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/usage")
func init() { coreusage.RegisterPlugin(&usageQueuePlugin{})}
type usageQueuePlugin struct{}
func (p *usageQueuePlugin) HandleUsage(ctx context.Context, record coreusage.Record) { if p == nil { return } if !Enabled() || !UsageStatisticsEnabled() { return }
timestamp := record.RequestedAt if timestamp.IsZero() { timestamp = time.Now() }
modelName := strings.TrimSpace(record.Model) if modelName == "" { modelName = "unknown" } aliasName := strings.TrimSpace(record.Alias) if aliasName == "" { aliasName = modelName } provider := strings.TrimSpace(record.Provider) if provider == "" { provider = "unknown" } executorType := strings.TrimSpace(record.ExecutorType) if executorType == "" { executorType = "unknown" } authType := strings.TrimSpace(record.AuthType) if authType == "" { authType = "unknown" } apiKey := strings.TrimSpace(record.APIKey) requestID := strings.TrimSpace(internallogging.GetRequestID(ctx)) reasoningEffort := strings.TrimSpace(record.ReasoningEffort) if reasoningEffort == "" { reasoningEffort = coreusage.ReasoningEffortFromContext(ctx) } serviceTier := strings.TrimSpace(record.ServiceTier) if serviceTier == "" { serviceTier = strings.TrimSpace(record.RequestServiceTier) } if serviceTier == "" { serviceTier = coreusage.ServiceTierFromContext(ctx) } responseServiceTier := strings.TrimSpace(record.ResponseServiceTier) clientRequestMetadata := internallogging.GetClientRequestMetadata(ctx)
usageDetail := coreusage.EnsureTokenBreakdownForProvider(record.Detail, record.Provider, record.ExecutorType) tokens := tokenStats{ InputTokens: usageDetail.InputTokens, OutputTokens: usageDetail.OutputTokens, ReasoningTokens: usageDetail.ReasoningTokens, CachedTokens: usageDetail.CachedTokens, CacheReadTokens: usageDetail.CacheReadTokens, CacheReadTokensPresent: true, CacheCreationTokens: usageDetail.CacheCreationTokens, TotalTokens: usageDetail.TotalTokens, }
failed := record.Failed if !failed { failed = !resolveSuccess(ctx) } fail := resolveFail(ctx, record, failed)
stream := record.Stream if !stream { stream = coreusage.StreamFromContext(ctx) }
detail := requestDetail{ Timestamp: timestamp, LatencyMs: record.Latency.Milliseconds(), TTFTMs: record.TTFT.Milliseconds(), Source: record.Source, AuthIndex: record.AuthIndex, AccessTokenHash: record.AccessTokenSHA256, ClientIP: clientRequestMetadata.ClientIP, XForwardedFor: clientRequestMetadata.XForwardedFor, UserAgent: clientRequestMetadata.UserAgent, Tokens: tokens, Failed: failed, Generate: coreusage.GenerateEnabled(record.Generate), Stream: stream, Fail: fail, ResponseHeaders: record.ResponseHeaders, }
payload, err := json.Marshal(queuedUsageDetail{ requestDetail: detail, AccountingVersion: coreusage.TokenAccountingSchemaVersion, TokenBreakdown: usageDetail.TokenBreakdown, Provider: provider, ExecutorType: executorType, Model: modelName, Alias: aliasName, Endpoint: resolveEndpoint(ctx), AuthType: authType, APIKey: apiKey, RequestID: requestID, ReasoningEffort: reasoningEffort, ServiceTier: serviceTier, ResponseServiceTier: responseServiceTier, }) if err != nil { return } Enqueue(payload)}
type queuedUsageDetail struct { requestDetail AccountingVersion int `json:"accounting_version"` TokenBreakdown coreusage.TokenBreakdown `json:"token_breakdown"` Provider string `json:"provider"` ExecutorType string `json:"executor_type"` Model string `json:"model"` Alias string `json:"alias"` Endpoint string `json:"endpoint"` AuthType string `json:"auth_type"` APIKey string `json:"api_key"` RequestID string `json:"request_id"` ReasoningEffort string `json:"reasoning_effort"` ServiceTier string `json:"service_tier"` ResponseServiceTier string `json:"response_service_tier,omitempty"`}
type requestDetail struct { Timestamp time.Time `json:"timestamp"` LatencyMs int64 `json:"latency_ms"` TTFTMs int64 `json:"ttft_ms"` Source string `json:"source"` AuthIndex string `json:"auth_index"` AccessTokenHash string `json:"access_token_sha256,omitempty"` ClientIP string `json:"client_ip"` XForwardedFor string `json:"x_forwarded_for"` UserAgent string `json:"user_agent"` Tokens tokenStats `json:"tokens"` Failed bool `json:"failed"` Generate bool `json:"generate"` Stream bool `json:"stream"` Fail failDetail `json:"fail"` ResponseHeaders http.Header `json:"response_headers,omitempty"`}
type tokenStats struct { InputTokens int64 `json:"input_tokens"` OutputTokens int64 `json:"output_tokens"` ReasoningTokens int64 `json:"reasoning_tokens"` CachedTokens int64 `json:"cached_tokens"` CacheReadTokens int64 `json:"cache_read_tokens"` CacheReadTokensPresent bool `json:"cache_read_tokens_present"` CacheCreationTokens int64 `json:"cache_creation_tokens"` TotalTokens int64 `json:"total_tokens"`}
type failDetail struct { StatusCode int `json:"status_code"` Body string `json:"body"`}
func resolveFail(ctx context.Context, record coreusage.Record, failed bool) failDetail { fail := failDetail{ StatusCode: record.Fail.StatusCode, Body: strings.TrimSpace(record.Fail.Body), } if !failed { return failDetail{StatusCode: 200} } if fail.StatusCode <= 0 { fail.StatusCode = internallogging.GetResponseStatus(ctx) } if fail.StatusCode <= 0 { fail.StatusCode = 500 } return fail}
func resolveSuccess(ctx context.Context) bool { status := internallogging.GetResponseStatus(ctx) if status == 0 { return true } return status < httpStatusBadRequest}
func resolveEndpoint(ctx context.Context) string { return strings.TrimSpace(internallogging.GetEndpoint(ctx))}
const httpStatusBadRequest = 400