120 lines
3.5 KiB
Go
120 lines
3.5 KiB
Go
package model
|
|
|
|
import (
|
|
"context"
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/QuantumNous/new-api/common"
|
|
)
|
|
|
|
type StatusChannel struct {
|
|
Id int
|
|
Group string
|
|
Status int
|
|
}
|
|
|
|
type GroupCacheStats struct {
|
|
Requests int64 `json:"requests"`
|
|
CachedRequests int64 `json:"cached_requests"`
|
|
InputTokens int64 `json:"input_tokens"`
|
|
CachedTokens int64 `json:"cached_tokens"`
|
|
FirstResponseMs float64 `json:"first_response_ms,omitempty"`
|
|
}
|
|
|
|
func GetStatusChannels(ctx context.Context) ([]StatusChannel, error) {
|
|
var channels []StatusChannel
|
|
err := DB.WithContext(ctx).Model(&Channel{}).Select("id", "group", "status").
|
|
Where("status IN ?", []int{common.ChannelStatusEnabled, common.ChannelStatusAutoDisabled}).
|
|
Order("id").Find(&channels).Error
|
|
return channels, err
|
|
}
|
|
|
|
// Channels retain their identity when a public group is renamed.
|
|
func GetGroupCacheStats(ctx context.Context, channels []StatusChannel, start, end time.Time) (map[string]GroupCacheStats, error) {
|
|
stats := make(map[string]GroupCacheStats)
|
|
firstResponseTotals := make(map[string]struct {
|
|
total float64
|
|
count int64
|
|
})
|
|
channelGroups := make(map[int][]string)
|
|
ids := make([]int, 0, len(channels))
|
|
for _, channel := range channels {
|
|
ids = append(ids, channel.Id)
|
|
for _, group := range strings.Split(channel.Group, ",") {
|
|
group = strings.TrimSpace(group)
|
|
if group != "" && !common.StringsContains(channelGroups[channel.Id], group) {
|
|
channelGroups[channel.Id] = append(channelGroups[channel.Id], group)
|
|
stats[group] = GroupCacheStats{}
|
|
}
|
|
}
|
|
}
|
|
if len(ids) == 0 {
|
|
return stats, nil
|
|
}
|
|
rows, err := LOG_DB.WithContext(ctx).Model(&Log{}).Select("channel_id", "prompt_tokens", "other").
|
|
Where("type = ? AND token_id > 0 AND created_at >= ? AND created_at < ?", LogTypeConsume, start.Unix(), end.Unix()).
|
|
Where("channel_id IN ?", ids).Rows()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer rows.Close()
|
|
for rows.Next() {
|
|
var channelID int
|
|
var prompt int64
|
|
var raw string
|
|
if err := rows.Scan(&channelID, &prompt, &raw); err != nil {
|
|
return nil, err
|
|
}
|
|
var usage struct {
|
|
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"`
|
|
InputTotal int64 `json:"input_tokens_total"`
|
|
Semantic string `json:"usage_semantic"`
|
|
Claude bool `json:"claude"`
|
|
FirstResponseMs float64 `json:"frt"`
|
|
}
|
|
if strings.TrimSpace(raw) != "" && common.UnmarshalJsonStr(raw, &usage) != nil {
|
|
continue
|
|
}
|
|
input := max(prompt, 0)
|
|
cached := max(usage.Cached, 0)
|
|
if usage.InputTotal > 0 {
|
|
input = usage.InputTotal
|
|
} else if usage.Semantic == "anthropic" || (usage.Semantic == "" && usage.Claude) {
|
|
input += cached + max(usage.Written, usage.Created, max(usage.Created5m, 0)+max(usage.Created1h, 0), 0)
|
|
}
|
|
if input == 0 {
|
|
continue
|
|
}
|
|
cached = min(cached, input)
|
|
for _, group := range channelGroups[channelID] {
|
|
total := stats[group]
|
|
total.Requests++
|
|
if cached > 0 {
|
|
total.CachedRequests++
|
|
}
|
|
total.InputTokens += input
|
|
total.CachedTokens += cached
|
|
stats[group] = total
|
|
if usage.FirstResponseMs > 0 {
|
|
frt := firstResponseTotals[group]
|
|
frt.total += usage.FirstResponseMs
|
|
frt.count++
|
|
firstResponseTotals[group] = frt
|
|
}
|
|
}
|
|
}
|
|
for group, total := range firstResponseTotals {
|
|
if total.count > 0 {
|
|
groupStats := stats[group]
|
|
groupStats.FirstResponseMs = total.total / float64(total.count)
|
|
stats[group] = groupStats
|
|
}
|
|
}
|
|
return stats, rows.Err()
|
|
}
|