package sqlite import ( "time" "github.com/rose_cat707/luminary/internal/auth" "github.com/rose_cat707/luminary/internal/model" "github.com/rose_cat707/luminary/internal/settings" "github.com/rose_cat707/luminary/internal/store/cache" "gorm.io/gorm" ) type Repository struct { db *gorm.DB cache *cache.Store } func NewRepository(store *Store, c *cache.Store) *Repository { return &Repository{db: store.DB(), cache: c} } func (r *Repository) BootstrapFromDB() error { var providers []model.Provider if err := r.db.Find(&providers).Error; err != nil { return err } r.cache.LoadProviders(providers) var keys []model.ProviderKey if err := r.db.Find(&keys).Error; err != nil { return err } r.cache.LoadProviderKeys(keys) var models []model.ModelEntry if err := r.db.Find(&models).Error; err != nil { return err } r.cache.LoadModels(models) var ingress []model.IngressKey if err := r.db.Find(&ingress).Error; err != nil { return err } r.cache.LoadIngressKeys(ingress) var rules []model.IPRule if err := r.db.Find(&rules).Error; err != nil { return err } r.cache.LoadIPRules(rules) var routing []model.RoutingRule if err := r.db.Find(&routing).Error; err != nil { return err } r.cache.LoadRoutingRules(routing) r.PruneUsageRecords(model.MaxUsageLogRecords) return nil } func (r *Repository) EnsureAdmin() error { var count int64 r.db.Model(&model.AdminUser{}).Count(&count) if count > 0 { return nil } username, password := settings.LegacyAdminCredentials() hash, err := auth.HashPassword(password) if err != nil { return err } return r.db.Create(&model.AdminUser{Username: username, PasswordHash: hash}).Error } func (r *Repository) FlushUsage(deltas []cache.UsageDelta) { for _, d := range deltas { r.db.Create(&model.UsageRecord{ IngressKeyID: d.IngressKeyID, ProviderID: d.ProviderID, ProviderKeyID: d.ProviderKeyID, ClientIP: d.ClientIP, Endpoint: d.Endpoint, Model: d.Model, TokensIn: d.TokensIn, TokensOut: d.TokensOut, Cost: d.Cost, LatencyMs: d.LatencyMs, Status: d.Status, Stream: d.Stream, ErrorMsg: d.ErrorMsg, RequestBody: d.RequestBody, ResponseBody: d.ResponseBody, CreatedAt: time.Now(), }) if d.IngressKeyID > 0 { r.db.Model(&model.IngressKey{}).Where("id = ?", d.IngressKeyID).Updates(map[string]any{ "budget_used": gorm.Expr("budget_used + ?", d.BudgetUsed), "request_count": gorm.Expr("request_count + ?", d.Requests), "token_count": gorm.Expr("token_count + ?", d.Tokens), }) } if d.ProviderKeyID > 0 { r.db.Model(&model.ProviderKey{}).Where("id = ?", d.ProviderKeyID).Updates(map[string]any{ "requests_today": gorm.Expr("requests_today + ?", d.Requests), "tokens_today": gorm.Expr("tokens_today + ?", d.Tokens), }) } } r.PruneUsageRecords(model.MaxUsageLogRecords) } func (r *Repository) RecordRequestLog(rec model.UsageRecord) { if rec.CreatedAt.IsZero() { rec.CreatedAt = time.Now() } r.db.Create(&rec) r.PruneUsageRecords(model.MaxUsageLogRecords) } func (r *Repository) PruneUsageRecords(max int) { if max <= 0 { return } var count int64 r.db.Model(&model.UsageRecord{}).Count(&count) if count <= int64(max) { return } var keepIDs []uint r.db.Model(&model.UsageRecord{}).Order("created_at desc").Limit(max).Pluck("id", &keepIDs) if len(keepIDs) == 0 { return } r.db.Where("id NOT IN ?", keepIDs).Delete(&model.UsageRecord{}) }