mirror of
https://github.com/EZ-Api/ez-api.git
synced 2026-01-13 17:47:51 +00:00
364 lines
11 KiB
Go
364 lines
11 KiB
Go
package api
|
|
|
|
import (
|
|
"fmt"
|
|
"net/http"
|
|
"sort"
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/ez-api/ez-api/internal/model"
|
|
"github.com/gin-gonic/gin"
|
|
"gorm.io/gorm"
|
|
)
|
|
|
|
type KeyUsageStat struct {
|
|
KeyID uint `json:"key_id"`
|
|
Requests int64 `json:"requests"`
|
|
Tokens int64 `json:"tokens"`
|
|
}
|
|
|
|
type ModelUsageStat struct {
|
|
Model string `json:"model"`
|
|
Requests int64 `json:"requests"`
|
|
Tokens int64 `json:"tokens"`
|
|
}
|
|
|
|
type MasterUsageStatsResponse struct {
|
|
Period string `json:"period,omitempty"`
|
|
TotalRequests int64 `json:"total_requests"`
|
|
TotalTokens int64 `json:"total_tokens"`
|
|
ByKey []KeyUsageStat `json:"by_key"`
|
|
ByModel []ModelUsageStat `json:"by_model"`
|
|
}
|
|
|
|
// GetSelfStats godoc
|
|
// @Summary Usage stats (master)
|
|
// @Description Aggregate request stats for the authenticated master
|
|
// @Tags master
|
|
// @Produce json
|
|
// @Security MasterAuth
|
|
// @Param period query string false "today|week|month|all"
|
|
// @Param since query int false "unix seconds"
|
|
// @Param until query int false "unix seconds"
|
|
// @Success 200 {object} MasterUsageStatsResponse
|
|
// @Failure 400 {object} gin.H
|
|
// @Failure 401 {object} gin.H
|
|
// @Failure 500 {object} gin.H
|
|
// @Router /v1/stats [get]
|
|
func (h *MasterHandler) GetSelfStats(c *gin.Context) {
|
|
master, exists := c.Get("master")
|
|
if !exists {
|
|
c.JSON(http.StatusUnauthorized, gin.H{"error": "master key not found in context"})
|
|
return
|
|
}
|
|
m := master.(*model.Master)
|
|
|
|
rng, err := parseStatsRange(c)
|
|
if err != nil {
|
|
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
|
return
|
|
}
|
|
|
|
logDB := h.logDBConn()
|
|
base := h.logBaseQuery()
|
|
if logDB == h.db {
|
|
base = base.Joins("JOIN keys ON keys.id = log_records.key_id").
|
|
Where("keys.master_id = ?", m.ID)
|
|
} else {
|
|
var keyIDs []uint
|
|
if err := h.db.Model(&model.Key{}).
|
|
Where("master_id = ?", m.ID).
|
|
Pluck("id", &keyIDs).Error; err != nil {
|
|
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to load master keys", "details": err.Error()})
|
|
return
|
|
}
|
|
if len(keyIDs) == 0 {
|
|
base = base.Where("1 = 0")
|
|
} else {
|
|
base = base.Where("log_records.key_id IN ?", keyIDs)
|
|
}
|
|
}
|
|
base = applyStatsRange(base, rng)
|
|
|
|
totalRequests, totalTokens, err := aggregateTotals(base)
|
|
if err != nil {
|
|
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to aggregate stats", "details": err.Error()})
|
|
return
|
|
}
|
|
|
|
var byKey []KeyUsageStat
|
|
if err := base.Session(&gorm.Session{}).
|
|
Select("log_records.key_id as key_id, COUNT(*) as requests, COALESCE(SUM(log_records.tokens_in + log_records.tokens_out),0) as tokens").
|
|
Group("log_records.key_id").
|
|
Scan(&byKey).Error; err != nil {
|
|
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to group by key", "details": err.Error()})
|
|
return
|
|
}
|
|
|
|
var byModel []ModelUsageStat
|
|
if err := base.Session(&gorm.Session{}).
|
|
Select("log_records.model_name as model, COUNT(*) as requests, COALESCE(SUM(log_records.tokens_in + log_records.tokens_out),0) as tokens").
|
|
Group("log_records.model_name").
|
|
Scan(&byModel).Error; err != nil {
|
|
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to group by model", "details": err.Error()})
|
|
return
|
|
}
|
|
|
|
c.JSON(http.StatusOK, MasterUsageStatsResponse{
|
|
Period: rng.Period,
|
|
TotalRequests: totalRequests,
|
|
TotalTokens: totalTokens,
|
|
ByKey: byKey,
|
|
ByModel: byModel,
|
|
})
|
|
}
|
|
|
|
type MasterUsageAgg struct {
|
|
MasterID uint `json:"master_id"`
|
|
Requests int64 `json:"requests"`
|
|
Tokens int64 `json:"tokens"`
|
|
}
|
|
|
|
type ProviderUsageAgg struct {
|
|
ProviderID uint `json:"provider_id"`
|
|
ProviderType string `json:"provider_type"`
|
|
ProviderName string `json:"provider_name"`
|
|
Requests int64 `json:"requests"`
|
|
Tokens int64 `json:"tokens"`
|
|
}
|
|
|
|
type AdminUsageStatsResponse struct {
|
|
Period string `json:"period,omitempty"`
|
|
TotalMasters int64 `json:"total_masters"`
|
|
ActiveMasters int64 `json:"active_masters"`
|
|
TotalRequests int64 `json:"total_requests"`
|
|
TotalTokens int64 `json:"total_tokens"`
|
|
ByMaster []MasterUsageAgg `json:"by_master"`
|
|
ByProvider []ProviderUsageAgg `json:"by_provider"`
|
|
}
|
|
|
|
// GetAdminStats godoc
|
|
// @Summary Usage stats (admin)
|
|
// @Description Aggregate request stats across all masters
|
|
// @Tags admin
|
|
// @Produce json
|
|
// @Security AdminAuth
|
|
// @Param period query string false "today|week|month|all"
|
|
// @Param since query int false "unix seconds"
|
|
// @Param until query int false "unix seconds"
|
|
// @Success 200 {object} AdminUsageStatsResponse
|
|
// @Failure 400 {object} gin.H
|
|
// @Failure 500 {object} gin.H
|
|
// @Router /admin/stats [get]
|
|
func (h *AdminHandler) GetAdminStats(c *gin.Context) {
|
|
rng, err := parseStatsRange(c)
|
|
if err != nil {
|
|
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
|
return
|
|
}
|
|
|
|
var totalMasters int64
|
|
if err := h.db.Model(&model.Master{}).Count(&totalMasters).Error; err != nil {
|
|
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to count masters", "details": err.Error()})
|
|
return
|
|
}
|
|
var activeMasters int64
|
|
if err := h.db.Model(&model.Master{}).Where("status = ?", "active").Count(&activeMasters).Error; err != nil {
|
|
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to count active masters", "details": err.Error()})
|
|
return
|
|
}
|
|
|
|
logDB := h.logDBConn()
|
|
base := h.logBaseQuery()
|
|
base = applyStatsRange(base, rng)
|
|
|
|
totalRequests, totalTokens, err := aggregateTotals(base)
|
|
if err != nil {
|
|
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to aggregate stats", "details": err.Error()})
|
|
return
|
|
}
|
|
|
|
var byMaster []MasterUsageAgg
|
|
if logDB == h.db {
|
|
if err := base.Session(&gorm.Session{}).
|
|
Joins("JOIN keys ON keys.id = log_records.key_id").
|
|
Select("keys.master_id as master_id, COUNT(*) as requests, COALESCE(SUM(log_records.tokens_in + log_records.tokens_out),0) as tokens").
|
|
Group("keys.master_id").
|
|
Scan(&byMaster).Error; err != nil {
|
|
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to group by master", "details": err.Error()})
|
|
return
|
|
}
|
|
} else {
|
|
var err error
|
|
byMaster, err = aggregateByMasterFromLogs(base, h.db)
|
|
if err != nil {
|
|
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to group by master", "details": err.Error()})
|
|
return
|
|
}
|
|
}
|
|
|
|
var byProvider []ProviderUsageAgg
|
|
if err := base.Session(&gorm.Session{}).
|
|
Select("log_records.provider_id as provider_id, log_records.provider_type as provider_type, log_records.provider_name as provider_name, COUNT(*) as requests, COALESCE(SUM(log_records.tokens_in + log_records.tokens_out),0) as tokens").
|
|
Group("log_records.provider_id, log_records.provider_type, log_records.provider_name").
|
|
Scan(&byProvider).Error; err != nil {
|
|
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to group by provider", "details": err.Error()})
|
|
return
|
|
}
|
|
|
|
c.JSON(http.StatusOK, AdminUsageStatsResponse{
|
|
Period: rng.Period,
|
|
TotalMasters: totalMasters,
|
|
ActiveMasters: activeMasters,
|
|
TotalRequests: totalRequests,
|
|
TotalTokens: totalTokens,
|
|
ByMaster: byMaster,
|
|
ByProvider: byProvider,
|
|
})
|
|
}
|
|
|
|
type statsRange struct {
|
|
Since *time.Time
|
|
Until *time.Time
|
|
Period string
|
|
}
|
|
|
|
func parseStatsRange(c *gin.Context) (statsRange, error) {
|
|
period := strings.ToLower(strings.TrimSpace(c.Query("period")))
|
|
if period != "" {
|
|
if period == "all" {
|
|
return statsRange{Period: period}, nil
|
|
}
|
|
start, now := periodWindow(period)
|
|
if start.IsZero() {
|
|
return statsRange{}, fmt.Errorf("invalid period")
|
|
}
|
|
return statsRange{Since: &start, Until: &now, Period: period}, nil
|
|
}
|
|
|
|
var since *time.Time
|
|
if t, ok := parseUnixSeconds(c.Query("since")); ok {
|
|
since = &t
|
|
}
|
|
var until *time.Time
|
|
if t, ok := parseUnixSeconds(c.Query("until")); ok {
|
|
until = &t
|
|
}
|
|
return statsRange{Since: since, Until: until}, nil
|
|
}
|
|
|
|
func periodWindow(period string) (time.Time, time.Time) {
|
|
now := time.Now().UTC()
|
|
startOfDay := time.Date(now.Year(), now.Month(), now.Day(), 0, 0, 0, 0, time.UTC)
|
|
switch period {
|
|
case "today":
|
|
return startOfDay, now
|
|
case "week":
|
|
weekday := int(startOfDay.Weekday())
|
|
if weekday == 0 {
|
|
weekday = 7
|
|
}
|
|
start := startOfDay.AddDate(0, 0, -(weekday - 1))
|
|
return start, now
|
|
case "month":
|
|
start := time.Date(now.Year(), now.Month(), 1, 0, 0, 0, 0, time.UTC)
|
|
return start, now
|
|
default:
|
|
return time.Time{}, time.Time{}
|
|
}
|
|
}
|
|
|
|
func applyStatsRange(q *gorm.DB, rng statsRange) *gorm.DB {
|
|
if rng.Since != nil {
|
|
q = q.Where("log_records.created_at >= ?", *rng.Since)
|
|
}
|
|
if rng.Until != nil {
|
|
q = q.Where("log_records.created_at <= ?", *rng.Until)
|
|
}
|
|
return q
|
|
}
|
|
|
|
func aggregateTotals(q *gorm.DB) (int64, int64, error) {
|
|
var totalRequests int64
|
|
if err := q.Session(&gorm.Session{}).Count(&totalRequests).Error; err != nil {
|
|
return 0, 0, err
|
|
}
|
|
type totals struct {
|
|
Tokens int64
|
|
}
|
|
var t totals
|
|
if err := q.Session(&gorm.Session{}).
|
|
Select("COALESCE(SUM(log_records.tokens_in + log_records.tokens_out),0) as tokens").
|
|
Scan(&t).Error; err != nil {
|
|
return 0, 0, err
|
|
}
|
|
return totalRequests, t.Tokens, nil
|
|
}
|
|
|
|
func aggregateByMasterFromLogs(base *gorm.DB, mainDB *gorm.DB) ([]MasterUsageAgg, error) {
|
|
if base == nil || mainDB == nil {
|
|
return nil, nil
|
|
}
|
|
var byKey []KeyUsageStat
|
|
if err := base.Session(&gorm.Session{}).
|
|
Select("log_records.key_id as key_id, COUNT(*) as requests, COALESCE(SUM(log_records.tokens_in + log_records.tokens_out),0) as tokens").
|
|
Group("log_records.key_id").
|
|
Scan(&byKey).Error; err != nil {
|
|
return nil, err
|
|
}
|
|
if len(byKey) == 0 {
|
|
return []MasterUsageAgg{}, nil
|
|
}
|
|
|
|
keyIDs := make([]uint, 0, len(byKey))
|
|
for _, row := range byKey {
|
|
if row.KeyID > 0 {
|
|
keyIDs = append(keyIDs, row.KeyID)
|
|
}
|
|
}
|
|
if len(keyIDs) == 0 {
|
|
return []MasterUsageAgg{}, nil
|
|
}
|
|
|
|
type keyMaster struct {
|
|
ID uint
|
|
MasterID uint
|
|
}
|
|
var keyMasters []keyMaster
|
|
if err := mainDB.Model(&model.Key{}).
|
|
Select("id, master_id").
|
|
Where("id IN ?", keyIDs).
|
|
Scan(&keyMasters).Error; err != nil {
|
|
return nil, err
|
|
}
|
|
keyToMaster := make(map[uint]uint, len(keyMasters))
|
|
for _, row := range keyMasters {
|
|
keyToMaster[row.ID] = row.MasterID
|
|
}
|
|
|
|
agg := make(map[uint]*MasterUsageAgg)
|
|
for _, row := range byKey {
|
|
masterID := keyToMaster[row.KeyID]
|
|
if masterID == 0 {
|
|
continue
|
|
}
|
|
entry, ok := agg[masterID]
|
|
if !ok {
|
|
entry = &MasterUsageAgg{MasterID: masterID}
|
|
agg[masterID] = entry
|
|
}
|
|
entry.Requests += row.Requests
|
|
entry.Tokens += row.Tokens
|
|
}
|
|
|
|
out := make([]MasterUsageAgg, 0, len(agg))
|
|
for _, row := range agg {
|
|
out = append(out, *row)
|
|
}
|
|
sort.Slice(out, func(i, j int) bool {
|
|
return out[i].MasterID < out[j].MasterID
|
|
})
|
|
return out, nil
|
|
}
|