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
@@ -16,6 +16,9 @@ type LogView struct {
|
||||
Group string `json:"group"`
|
||||
KeyID uint `json:"key_id"`
|
||||
ModelName string `json:"model"`
|
||||
ProviderID uint `json:"provider_id"`
|
||||
ProviderType string `json:"provider_type"`
|
||||
ProviderName string `json:"provider_name"`
|
||||
StatusCode int `json:"status_code"`
|
||||
LatencyMs int64 `json:"latency_ms"`
|
||||
TokensIn int64 `json:"tokens_in"`
|
||||
@@ -34,6 +37,9 @@ func toLogView(r model.LogRecord) LogView {
|
||||
Group: r.Group,
|
||||
KeyID: r.KeyID,
|
||||
ModelName: r.ModelName,
|
||||
ProviderID: r.ProviderID,
|
||||
ProviderType: r.ProviderType,
|
||||
ProviderName: r.ProviderName,
|
||||
StatusCode: r.StatusCode,
|
||||
LatencyMs: r.LatencyMs,
|
||||
TokensIn: r.TokensIn,
|
||||
@@ -276,7 +282,7 @@ func (h *MasterHandler) ListSelfLogs(c *gin.Context) {
|
||||
c.JSON(http.StatusOK, ListLogsResponse{Total: total, Limit: limit, Offset: offset, Items: out})
|
||||
}
|
||||
|
||||
// GetSelfStats godoc
|
||||
// GetSelfLogStats godoc
|
||||
// @Summary Log stats (master)
|
||||
// @Description Aggregate request log stats for the authenticated master
|
||||
// @Tags master
|
||||
@@ -287,8 +293,8 @@ func (h *MasterHandler) ListSelfLogs(c *gin.Context) {
|
||||
// @Success 200 {object} LogStatsResponse
|
||||
// @Failure 401 {object} gin.H
|
||||
// @Failure 500 {object} gin.H
|
||||
// @Router /v1/stats [get]
|
||||
func (h *MasterHandler) GetSelfStats(c *gin.Context) {
|
||||
// @Router /v1/logs/stats [get]
|
||||
func (h *MasterHandler) GetSelfLogStats(c *gin.Context) {
|
||||
master, exists := c.Get("master")
|
||||
if !exists {
|
||||
c.JSON(http.StatusUnauthorized, gin.H{"error": "master key not found in context"})
|
||||
|
||||
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
|
||||
}
|
||||
191
internal/api/stats_handler_test.go
Normal file
191
internal/api/stats_handler_test.go
Normal file
@@ -0,0 +1,191 @@
|
||||
package api
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/alicebob/miniredis/v2"
|
||||
"github.com/ez-api/ez-api/internal/model"
|
||||
"github.com/ez-api/ez-api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/redis/go-redis/v9"
|
||||
"gorm.io/driver/sqlite"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
func TestMasterStats_AggregatesByKeyAndModel(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
dsn := fmt.Sprintf("file:%s?mode=memory&cache=shared", t.Name())
|
||||
db, err := gorm.Open(sqlite.Open(dsn), &gorm.Config{})
|
||||
if err != nil {
|
||||
t.Fatalf("open sqlite: %v", err)
|
||||
}
|
||||
if err := db.AutoMigrate(&model.Master{}, &model.Key{}, &model.LogRecord{}); err != nil {
|
||||
t.Fatalf("migrate: %v", err)
|
||||
}
|
||||
|
||||
m := &model.Master{Name: "m1", Group: "g", Status: "active", Epoch: 1, MasterKeyDigest: "d1"}
|
||||
if err := db.Create(m).Error; err != nil {
|
||||
t.Fatalf("create master: %v", err)
|
||||
}
|
||||
k1 := &model.Key{MasterID: m.ID, TokenHash: "h1", Group: "g", Status: "active", IssuedAtEpoch: 1}
|
||||
k2 := &model.Key{MasterID: m.ID, TokenHash: "h2", Group: "g", Status: "active", IssuedAtEpoch: 1}
|
||||
if err := db.Create(k1).Error; err != nil {
|
||||
t.Fatalf("create k1: %v", err)
|
||||
}
|
||||
if err := db.Create(k2).Error; err != nil {
|
||||
t.Fatalf("create k2: %v", err)
|
||||
}
|
||||
|
||||
if err := db.Create(&model.LogRecord{
|
||||
Group: "rg",
|
||||
KeyID: k1.ID,
|
||||
ModelName: "ns.m1",
|
||||
ProviderID: 10,
|
||||
ProviderType: "openai",
|
||||
ProviderName: "p1",
|
||||
StatusCode: 200,
|
||||
TokensIn: 5,
|
||||
TokensOut: 7,
|
||||
}).Error; err != nil {
|
||||
t.Fatalf("create log1: %v", err)
|
||||
}
|
||||
if err := db.Create(&model.LogRecord{
|
||||
Group: "rg",
|
||||
KeyID: k2.ID,
|
||||
ModelName: "ns.m2",
|
||||
ProviderID: 11,
|
||||
ProviderType: "anthropic",
|
||||
ProviderName: "p2",
|
||||
StatusCode: 200,
|
||||
TokensIn: 2,
|
||||
TokensOut: 3,
|
||||
}).Error; err != nil {
|
||||
t.Fatalf("create log2: %v", err)
|
||||
}
|
||||
|
||||
mr := miniredis.RunT(t)
|
||||
rdb := redis.NewClient(&redis.Options{Addr: mr.Addr()})
|
||||
masterSvc := service.NewMasterService(db)
|
||||
syncSvc := service.NewSyncService(rdb)
|
||||
h := NewMasterHandler(db, masterSvc, syncSvc)
|
||||
|
||||
withMaster := func(next gin.HandlerFunc) gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
c.Set("master", m)
|
||||
next(c)
|
||||
}
|
||||
}
|
||||
|
||||
r := gin.New()
|
||||
r.GET("/v1/stats", withMaster(h.GetSelfStats))
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet, "/v1/stats?period=all", nil)
|
||||
rr := httptest.NewRecorder()
|
||||
r.ServeHTTP(rr, req)
|
||||
|
||||
if rr.Code != http.StatusOK {
|
||||
t.Fatalf("expected 200, got %d body=%s", rr.Code, rr.Body.String())
|
||||
}
|
||||
var resp MasterUsageStatsResponse
|
||||
if err := json.Unmarshal(rr.Body.Bytes(), &resp); err != nil {
|
||||
t.Fatalf("unmarshal: %v", err)
|
||||
}
|
||||
if resp.TotalRequests != 2 || resp.TotalTokens != 17 {
|
||||
t.Fatalf("unexpected totals: %+v", resp)
|
||||
}
|
||||
if len(resp.ByKey) != 2 || len(resp.ByModel) != 2 {
|
||||
t.Fatalf("unexpected breakdown: %+v", resp)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAdminStats_AggregatesByProvider(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
dsn := fmt.Sprintf("file:%s?mode=memory&cache=shared", t.Name())
|
||||
db, err := gorm.Open(sqlite.Open(dsn), &gorm.Config{})
|
||||
if err != nil {
|
||||
t.Fatalf("open sqlite: %v", err)
|
||||
}
|
||||
if err := db.AutoMigrate(&model.Master{}, &model.Key{}, &model.LogRecord{}); err != nil {
|
||||
t.Fatalf("migrate: %v", err)
|
||||
}
|
||||
|
||||
m1 := &model.Master{Name: "m1", Group: "g", Status: "active", Epoch: 1, MasterKeyDigest: "d1"}
|
||||
m2 := &model.Master{Name: "m2", Group: "g", Status: "suspended", Epoch: 1, MasterKeyDigest: "d2"}
|
||||
if err := db.Create(m1).Error; err != nil {
|
||||
t.Fatalf("create m1: %v", err)
|
||||
}
|
||||
if err := db.Create(m2).Error; err != nil {
|
||||
t.Fatalf("create m2: %v", err)
|
||||
}
|
||||
k1 := &model.Key{MasterID: m1.ID, TokenHash: "h1", Group: "g", Status: "active", IssuedAtEpoch: 1}
|
||||
k2 := &model.Key{MasterID: m2.ID, TokenHash: "h2", Group: "g", Status: "active", IssuedAtEpoch: 1}
|
||||
if err := db.Create(k1).Error; err != nil {
|
||||
t.Fatalf("create k1: %v", err)
|
||||
}
|
||||
if err := db.Create(k2).Error; err != nil {
|
||||
t.Fatalf("create k2: %v", err)
|
||||
}
|
||||
|
||||
if err := db.Create(&model.LogRecord{
|
||||
Group: "rg",
|
||||
KeyID: k1.ID,
|
||||
ModelName: "ns.m1",
|
||||
ProviderID: 10,
|
||||
ProviderType: "openai",
|
||||
ProviderName: "p1",
|
||||
StatusCode: 200,
|
||||
TokensIn: 4,
|
||||
TokensOut: 6,
|
||||
}).Error; err != nil {
|
||||
t.Fatalf("create log1: %v", err)
|
||||
}
|
||||
if err := db.Create(&model.LogRecord{
|
||||
Group: "rg",
|
||||
KeyID: k2.ID,
|
||||
ModelName: "ns.m2",
|
||||
ProviderID: 11,
|
||||
ProviderType: "anthropic",
|
||||
ProviderName: "p2",
|
||||
StatusCode: 200,
|
||||
TokensIn: 1,
|
||||
TokensOut: 2,
|
||||
}).Error; err != nil {
|
||||
t.Fatalf("create log2: %v", err)
|
||||
}
|
||||
|
||||
mr := miniredis.RunT(t)
|
||||
rdb := redis.NewClient(&redis.Options{Addr: mr.Addr()})
|
||||
masterSvc := service.NewMasterService(db)
|
||||
syncSvc := service.NewSyncService(rdb)
|
||||
adminHandler := NewAdminHandler(db, masterSvc, syncSvc)
|
||||
|
||||
r := gin.New()
|
||||
r.GET("/admin/stats", adminHandler.GetAdminStats)
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet, "/admin/stats?period=all", nil)
|
||||
rr := httptest.NewRecorder()
|
||||
r.ServeHTTP(rr, req)
|
||||
|
||||
if rr.Code != http.StatusOK {
|
||||
t.Fatalf("expected 200, got %d body=%s", rr.Code, rr.Body.String())
|
||||
}
|
||||
var resp AdminUsageStatsResponse
|
||||
if err := json.Unmarshal(rr.Body.Bytes(), &resp); err != nil {
|
||||
t.Fatalf("unmarshal: %v", err)
|
||||
}
|
||||
if resp.TotalMasters != 2 || resp.ActiveMasters != 1 {
|
||||
t.Fatalf("unexpected master counts: %+v", resp)
|
||||
}
|
||||
if resp.TotalRequests != 2 || resp.TotalTokens != 13 {
|
||||
t.Fatalf("unexpected totals: %+v", resp)
|
||||
}
|
||||
if len(resp.ByProvider) != 2 {
|
||||
t.Fatalf("expected provider breakdown, got %+v", resp.ByProvider)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user