Files

87 lines
2.8 KiB
Go

package model
import (
"context"
"time"
"github.com/QuantumNous/new-api/common"
)
type GroupModelCost struct {
USDPerMillion *float64 `json:"usd_per_million"`
Requests int64 `json:"requests"`
Tokens float64 `json:"tokens"`
quota float64
}
// Consumption logs contain the settled charge, not the pre-consumed quota.
// Attribute by the group actually billed, never by today's channel membership.
func GetGroupModelCosts(ctx context.Context, start, end time.Time) (map[string]map[string]GroupModelCost, error) {
result := map[string]map[string]GroupModelCost{}
rows, err := LOG_DB.WithContext(ctx).Model(&Log{}).Select("group", "model_name", "quota", "prompt_tokens", "completion_tokens", "other").
Where("type = ? AND token_id > 0 AND created_at >= ? AND created_at < ?", LogTypeConsume, start.Unix(), end.Unix()).Rows()
if err != nil {
return nil, err
}
defer rows.Close()
for rows.Next() {
var group, name, raw string
var quota, prompt, completion int64
if err := rows.Scan(&group, &name, &quota, &prompt, &completion, &raw); err != nil {
return nil, err
}
if group == "" || name == "" || quota < 0 {
continue
}
tokens, valid := consumptionTokenTotal(prompt, completion, raw)
if !valid || tokens <= 0 {
continue
}
if result[group] == nil {
result[group] = map[string]GroupModelCost{}
}
value := result[group][name]
value.Requests++
value.Tokens += tokens
value.quota += float64(quota)
result[group][name] = value
}
if err := rows.Err(); err != nil {
return nil, err
}
for group, models := range result {
for name, value := range models {
if common.QuotaPerUnit > 0 && value.Tokens > 0 {
rate := value.quota / common.QuotaPerUnit / value.Tokens * 1_000_000
value.USDPerMillion = &rate
}
result[group][name] = value
}
}
return result, nil
}
func consumptionTokenTotal(prompt, completion int64, raw string) (float64, bool) {
var usage struct {
Total int64 `json:"input_tokens_total"`
Cached int64 `json:"cache_tokens"`
Created int64 `json:"cache_creation_tokens"`
Created5m int64 `json:"cache_creation_tokens_5m"`
Created1h int64 `json:"cache_creation_tokens_1h"`
Written int64 `json:"cache_write_tokens"`
Semantic string `json:"usage_semantic"`
Claude bool `json:"claude"`
}
if raw != "" && common.UnmarshalJsonStr(raw, &usage) != nil {
return 0, false
}
input := float64(max(prompt, 0))
if usage.Total > 0 {
input = float64(usage.Total)
} else if usage.Semantic == "anthropic" || (usage.Semantic == "" && usage.Claude) {
created := max(float64(usage.Written), float64(usage.Created), float64(max(usage.Created5m, 0))+float64(max(usage.Created1h, 0)), 0)
input += float64(max(usage.Cached, 0)) + created
}
return input + float64(max(completion, 0)), true
}