mirror of
https://github.com/router-for-me/CLIProxyAPI.git
synced 2026-05-08 14:48:29 +08:00
ba5d8ca733
- 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.
136 lines
3.3 KiB
Go
136 lines
3.3 KiB
Go
package redisqueue
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"strings"
|
|
"time"
|
|
|
|
internallogging "github.com/router-for-me/CLIProxyAPI/v6/internal/logging"
|
|
coreusage "github.com/router-for-me/CLIProxyAPI/v6/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"
|
|
}
|
|
authType := strings.TrimSpace(record.AuthType)
|
|
if authType == "" {
|
|
authType = "unknown"
|
|
}
|
|
apiKey := strings.TrimSpace(record.APIKey)
|
|
requestID := strings.TrimSpace(internallogging.GetRequestID(ctx))
|
|
|
|
tokens := tokenStats{
|
|
InputTokens: record.Detail.InputTokens,
|
|
OutputTokens: record.Detail.OutputTokens,
|
|
ReasoningTokens: record.Detail.ReasoningTokens,
|
|
CachedTokens: record.Detail.CachedTokens,
|
|
TotalTokens: record.Detail.TotalTokens,
|
|
}
|
|
if tokens.TotalTokens == 0 {
|
|
tokens.TotalTokens = tokens.InputTokens + tokens.OutputTokens + tokens.ReasoningTokens
|
|
}
|
|
if tokens.TotalTokens == 0 {
|
|
tokens.TotalTokens = tokens.InputTokens + tokens.OutputTokens + tokens.ReasoningTokens + tokens.CachedTokens
|
|
}
|
|
|
|
failed := record.Failed
|
|
if !failed {
|
|
failed = !resolveSuccess(ctx)
|
|
}
|
|
|
|
detail := requestDetail{
|
|
Timestamp: timestamp,
|
|
LatencyMs: record.Latency.Milliseconds(),
|
|
Source: record.Source,
|
|
AuthIndex: record.AuthIndex,
|
|
Tokens: tokens,
|
|
Failed: failed,
|
|
}
|
|
|
|
payload, err := json.Marshal(queuedUsageDetail{
|
|
requestDetail: detail,
|
|
Provider: provider,
|
|
Model: modelName,
|
|
Alias: aliasName,
|
|
Endpoint: resolveEndpoint(ctx),
|
|
AuthType: authType,
|
|
APIKey: apiKey,
|
|
RequestID: requestID,
|
|
})
|
|
if err != nil {
|
|
return
|
|
}
|
|
Enqueue(payload)
|
|
}
|
|
|
|
type queuedUsageDetail struct {
|
|
requestDetail
|
|
Provider string `json:"provider"`
|
|
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"`
|
|
}
|
|
|
|
type requestDetail struct {
|
|
Timestamp time.Time `json:"timestamp"`
|
|
LatencyMs int64 `json:"latency_ms"`
|
|
Source string `json:"source"`
|
|
AuthIndex string `json:"auth_index"`
|
|
Tokens tokenStats `json:"tokens"`
|
|
Failed bool `json:"failed"`
|
|
}
|
|
|
|
type tokenStats struct {
|
|
InputTokens int64 `json:"input_tokens"`
|
|
OutputTokens int64 `json:"output_tokens"`
|
|
ReasoningTokens int64 `json:"reasoning_tokens"`
|
|
CachedTokens int64 `json:"cached_tokens"`
|
|
TotalTokens int64 `json:"total_tokens"`
|
|
}
|
|
|
|
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
|