123 lines
4.2 KiB
Go
123 lines
4.2 KiB
Go
package service
|
|
|
|
import (
|
|
"slices"
|
|
|
|
"github.com/QuantumNous/new-api/model"
|
|
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
|
"github.com/QuantumNous/new-api/relaykit/dto"
|
|
)
|
|
|
|
const (
|
|
usageBillingPathLocal = "local"
|
|
usageBillingPathUpstream = "upstream"
|
|
usageBillingPathOpenAI = "billing-usage-openai"
|
|
usageBillingPathOpenAIEstimated = "billing-usage-openai-estimated"
|
|
usageBillingPathAnthropic = "billing-usage-anthropic"
|
|
usageBillingPathAnthropicEstimated = "billing-usage-anthropic-estimated"
|
|
usageBillingPathGemini = "billing-usage-gemini"
|
|
usageBillingPathGeminiEstimated = "billing-usage-gemini-estimated"
|
|
)
|
|
|
|
func effectiveBillingUsage(usage *dto.Usage) *dto.Usage {
|
|
if billingUsage, ok := usageFromBillingUsage(usage); ok {
|
|
return billingUsage
|
|
}
|
|
return usage
|
|
}
|
|
|
|
type claudeCacheReadAdjustment struct {
|
|
Rule string `json:"rule"`
|
|
Model string `json:"model"`
|
|
OriginalInputTokens int `json:"original_input_tokens"`
|
|
InputTokens int `json:"input_tokens"`
|
|
CacheReadTokens int `json:"cache_read_tokens"`
|
|
Status string `json:"status"`
|
|
}
|
|
|
|
// normalizeClaudeCacheReadUsage runs once after canonical upstream usage is
|
|
// selected. Only the billing copy changes; downstream responses and the
|
|
// original upstream snapshot retain their protocol values.
|
|
func normalizeClaudeCacheReadUsage(info *relaycommon.RelayInfo, usage *dto.Usage) (*dto.Usage, *claudeCacheReadAdjustment) {
|
|
if info == nil || info.ChannelMeta == nil || usage == nil ||
|
|
usage.UsageSemantic != dto.BillingUsageSemanticAnthropic ||
|
|
usage.BillingUsage == nil || usage.BillingUsage.Estimated || usage.BillingUsage.ClaudeUsage == nil {
|
|
return usage, nil
|
|
}
|
|
upstreamModel := info.GetUpstreamModelName()
|
|
if info.ResponseModel != nil && info.ResponseModel.UpstreamModel != "" {
|
|
// Some handlers replace UpstreamModelName with the response model.
|
|
// Match the name sent upstream, never the provider's response declaration.
|
|
upstreamModel = info.ResponseModel.UpstreamModel
|
|
}
|
|
if !slices.Contains(info.ChannelOtherSettings.ClaudeInputIncludesCacheReadModels, upstreamModel) {
|
|
return usage, nil
|
|
}
|
|
cached := usage.PromptTokensDetails.CachedTokens
|
|
adjustment := &claudeCacheReadAdjustment{
|
|
Rule: "claude_input_includes_cache_read_models", Model: upstreamModel,
|
|
OriginalInputTokens: usage.PromptTokens, InputTokens: usage.PromptTokens,
|
|
CacheReadTokens: cached, Status: "invalid_usage",
|
|
}
|
|
if usage.PromptTokens < 0 || cached < 0 || cached > usage.PromptTokens ||
|
|
usage.CompletionTokens < 0 || usage.InputTokens < cached || usage.TotalTokens < cached {
|
|
return usage, adjustment
|
|
}
|
|
if cached == 0 {
|
|
return usage, nil
|
|
}
|
|
adjusted := *usage
|
|
adjusted.PromptTokens -= cached
|
|
adjusted.InputTokens -= cached
|
|
adjusted.TotalTokens -= cached
|
|
adjusted.BillingUsage = dto.CloneBillingUsage(usage.BillingUsage)
|
|
adjusted.BillingUsage.ClaudeUsage.InputTokens = adjusted.PromptTokens
|
|
adjustment.InputTokens = adjusted.PromptTokens
|
|
adjustment.Status = "applied"
|
|
return &adjusted, adjustment
|
|
}
|
|
|
|
func usageBillingPathForLog(isLocalCountTokens bool, usage *dto.Usage) string {
|
|
effectiveUsage, ok := usageFromBillingUsage(usage)
|
|
if !ok {
|
|
if isLocalCountTokens {
|
|
return usageBillingPathLocal
|
|
}
|
|
return usageBillingPathUpstream
|
|
}
|
|
|
|
switch effectiveUsage.UsageSemantic {
|
|
case dto.BillingUsageSemanticOpenAI:
|
|
if usage.BillingUsage.Estimated {
|
|
return usageBillingPathOpenAIEstimated
|
|
}
|
|
return usageBillingPathOpenAI
|
|
case dto.BillingUsageSemanticAnthropic:
|
|
if usage.BillingUsage.Estimated {
|
|
return usageBillingPathAnthropicEstimated
|
|
}
|
|
return usageBillingPathAnthropic
|
|
case dto.BillingUsageSemanticGemini:
|
|
if usage.BillingUsage.Estimated {
|
|
return usageBillingPathGeminiEstimated
|
|
}
|
|
return usageBillingPathGemini
|
|
}
|
|
|
|
return usageBillingPathUpstream
|
|
}
|
|
|
|
func appendUsageBillingPathForLog(other *model.LogOther, isLocalCountTokens bool, usage *dto.Usage) {
|
|
if other == nil {
|
|
return
|
|
}
|
|
other.SetAdmin("usage_billing_path", usageBillingPathForLog(isLocalCountTokens, usage))
|
|
}
|
|
|
|
func usageFromBillingUsage(usage *dto.Usage) (*dto.Usage, bool) {
|
|
if usage == nil || usage.BillingUsage == nil {
|
|
return nil, false
|
|
}
|
|
return usage.BillingUsage.CanonicalUsage()
|
|
}
|