mirror of
https://github.com/EZ-Api/ez-api.git
synced 2026-01-13 17:47:51 +00:00
feat(stats): add usage stats and quota reset
This commit is contained in:
zenfun
270
internal/api/stats_handler.go
Normal file
270
internal/api/stats_handler.go
Normal file
@@ -0,0 +1,270 @@
|
||||
package api
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net/http"
|
||||
"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
|
||||
}
|
||||
|
||||
base := h.db.Model(&model.LogRecord{}).
|
||||
Joins("JOIN keys ON keys.id = log_records.key_id").
|
||||
Where("keys.master_id = ?", m.ID)
|
||||
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
|
||||
}
|
||||
|
||||
base := h.db.Model(&model.LogRecord{})
|
||||
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 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
|
||||
}
|
||||
|
||||
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
|
||||
}
|
||||
Reference in New Issue
Block a user