darroyo
4c9ed3c24b
Complete rewrite from Python/FastAPI to Go/Gin: - Go backend: auth (API keys + bcrypt), llama.cpp subprocess manager, hot-swap multi-model, rate limiting, quota system, webhooks - Vue 3 SPA admin panel (src/) with Tailwind CSS - Deployment: Docker multi-stage, docker-compose, nginx, systemd - GORM/SQLite models: ApiKey, Model, UsageLog, Quota, Webhook - REST API: /api/v1/admin/* (keys, models, chat, usage, health) - Embedded frontend via go:embed (build output at web/dist/) Removed legacy Python artifacts (app/, tests/, pyproject.toml, etc.)
120 lines
2.7 KiB
Go
120 lines
2.7 KiB
Go
package quota
|
|
|
|
import (
|
|
"errors"
|
|
"log/slog"
|
|
"time"
|
|
|
|
"github.com/google/uuid"
|
|
"gorm.io/gorm"
|
|
|
|
"github.com/llamalink/llamalink/internal/db"
|
|
)
|
|
|
|
var (
|
|
ErrQuotaExceeded = errors.New("quota exceeded for this period")
|
|
)
|
|
|
|
type Service struct {
|
|
db *gorm.DB
|
|
}
|
|
|
|
func NewService(db *gorm.DB) *Service {
|
|
return &Service{db: db}
|
|
}
|
|
|
|
func (s *Service) CheckQuota(apiKeyID uuid.UUID, modelScope string) (bool, string, error) {
|
|
now := time.Now().UTC()
|
|
|
|
var quota db.Quota
|
|
err := s.db.Where(
|
|
"api_key_id = ? AND period_start <= ? AND period_end > ?",
|
|
apiKeyID, now, now,
|
|
).First("a).Error
|
|
|
|
if err != nil {
|
|
if errors.Is(err, gorm.ErrRecordNotFound) {
|
|
return true, "", nil // No quota configured
|
|
}
|
|
return false, "", err
|
|
}
|
|
|
|
// Check model-specific scope if set
|
|
if quota.ModelScope != nil && *quota.ModelScope != "" && modelScope != "" {
|
|
if *quota.ModelScope != modelScope {
|
|
return true, "", nil // Quota doesn't apply to this model
|
|
}
|
|
}
|
|
|
|
if quota.TokensUsed >= quota.TokensLimit {
|
|
return false, "quota exceeded", nil
|
|
}
|
|
|
|
// Warning at 90%
|
|
if quota.TokensLimit > 0 {
|
|
usage := float64(quota.TokensUsed) / float64(quota.TokensLimit)
|
|
if usage >= 0.9 {
|
|
slog.Warn("quota usage warning",
|
|
"api_key_id", apiKeyID,
|
|
"usage_percent", usage*100,
|
|
"tokens_used", quota.TokensUsed,
|
|
"tokens_limit", quota.TokensLimit,
|
|
)
|
|
}
|
|
}
|
|
|
|
return true, "", nil
|
|
}
|
|
|
|
func (s *Service) ConsumeQuota(apiKeyID uuid.UUID, modelScope string, tokens int) error {
|
|
now := time.Now().UTC()
|
|
periodStart := time.Date(now.Year(), now.Month(), 1, 0, 0, 0, 0, time.UTC)
|
|
periodEnd := periodStart.AddDate(0, 1, 0)
|
|
|
|
return s.db.Transaction(func(tx *gorm.DB) error {
|
|
var quota db.Quota
|
|
err := tx.Where(
|
|
"api_key_id = ? AND period_start = ? AND period_end = ?",
|
|
apiKeyID, periodStart, periodEnd,
|
|
).First("a).Error
|
|
|
|
if err != nil {
|
|
if errors.Is(err, gorm.ErrRecordNotFound) {
|
|
return nil // No quota configured
|
|
}
|
|
return err
|
|
}
|
|
|
|
// Check model scope
|
|
if quota.ModelScope != nil && *quota.ModelScope != "" && modelScope != "" {
|
|
if *quota.ModelScope != modelScope {
|
|
return nil // Quota doesn't apply
|
|
}
|
|
}
|
|
|
|
newUsed := quota.TokensUsed + tokens
|
|
if quota.TokensLimit > 0 && newUsed > quota.TokensLimit {
|
|
return ErrQuotaExceeded
|
|
}
|
|
|
|
return tx.Model("a).Update("tokens_used", newUsed).Error
|
|
})
|
|
}
|
|
|
|
func (s *Service) GetUsage(apiKeyID uuid.UUID, periodStart, periodEnd time.Time) (used int, limit int, err error) {
|
|
var quota db.Quota
|
|
err = s.db.Where(
|
|
"api_key_id = ? AND period_start = ? AND period_end = ?",
|
|
apiKeyID, periodStart, periodEnd,
|
|
).First("a).Error
|
|
|
|
if err != nil {
|
|
if errors.Is(err, gorm.ErrRecordNotFound) {
|
|
return 0, 0, nil
|
|
}
|
|
return 0, 0, err
|
|
}
|
|
|
|
return quota.TokensUsed, quota.TokensLimit, nil
|
|
}
|