87 lines
2.8 KiB
Go
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, "a, &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
|
|
}
|