Files
NianAIGC/backend/internal/usage/usage.go

215 lines
7.4 KiB
Go

package usage
import (
"context"
"errors"
"fmt"
"sort"
"time"
)
type Preset string
const (
PresetToday Preset = "today"
Preset7Days Preset = "7d"
Preset30Days Preset = "30d"
PresetMonth Preset = "month"
)
var ErrForbidden = errors.New("usage scope forbidden")
var ErrInvalidDateRange = errors.New("invalid usage date range")
var ErrDateRangeTooLong = errors.New("usage date range exceeds 3660 days")
const UnassignedOrganizationID = "__unassigned__"
var shanghai = time.FixedZone("Asia/Shanghai", 8*60*60)
type DateRange struct {
From string `json:"from"`
To string `json:"to"`
StartDate string `json:"startDate"`
EndDate string `json:"endDate"`
Label string `json:"label"`
DayCount int `json:"dayCount"`
}
func PresetRange(preset Preset, now time.Time) (DateRange, error) {
today := now.In(shanghai)
start := time.Date(today.Year(), today.Month(), today.Day(), 0, 0, 0, 0, shanghai)
switch preset {
case PresetToday:
case Preset7Days:
start = start.AddDate(0, 0, -6)
case Preset30Days:
start = start.AddDate(0, 0, -29)
case PresetMonth:
start = time.Date(today.Year(), today.Month(), 1, 0, 0, 0, 0, shanghai)
default:
return DateRange{}, fmt.Errorf("unknown usage preset %q", preset)
}
end := time.Date(today.Year(), today.Month(), today.Day(), 0, 0, 0, 0, shanghai)
labels := map[Preset]string{PresetToday: "今天", Preset7Days: "近 7 天", Preset30Days: "近 30 天", PresetMonth: "本月"}
return DateRange{From: start.Format(time.RFC3339), To: end.AddDate(0, 0, 1).Format(time.RFC3339), StartDate: start.Format("2006-01-02"), EndDate: end.Format("2006-01-02"), Label: labels[preset], DayCount: int(end.Sub(start).Hours()/24) + 1}, nil
}
type Event struct {
ID, OwnerID, JobID, Source, Capability, Provider, ReqKey, AccountUsername, AccountDisplayName, TenantID, OrganizationID, OrganizationName, EstimatedUnit, Currency, CreatedAt string
Quantity int
ChargedAmountFen *int64
}
type Filters struct{ OwnerID, OrganizationID, Capability, Provider, From, To string }
type Requester struct{ AccountID, OrganizationID, Role string }
type Repository interface {
Insert(Event) (Event, bool, error)
List(Filters) ([]Event, error)
}
type OrganizationOptionSource interface {
ListOrganizationOptions(context.Context, Requester) ([]Option, error)
}
type OrganizationOptionSourceFunc func(context.Context, Requester) ([]Option, error)
func (fn OrganizationOptionSourceFunc) ListOrganizationOptions(ctx context.Context, requester Requester) ([]Option, error) {
return fn(ctx, requester)
}
type Service struct {
Repository Repository
OrganizationOptions OrganizationOptionSource
}
type PersonalRequest struct {
AccountID string
Preset Preset
Now time.Time
}
type AdminRequest struct {
Requester
StartDate, EndDate, OwnerID, Capability, Provider string
Now time.Time
RedactAccounts bool
}
type CountItem struct {
Key string `json:"key"`
Label string `json:"label"`
Count int `json:"count"`
}
type Record struct {
ID string `json:"id"`
JobID string `json:"jobId"`
OwnerID string `json:"ownerId"`
AccountName string `json:"accountName"`
AccountUsername string `json:"accountUsername,omitempty"`
OrganizationID string `json:"organizationId"`
OrganizationName string `json:"organizationName"`
Capability string `json:"capability"`
CapabilityLabel string `json:"capabilityLabel"`
Provider string `json:"provider,omitempty"`
ProviderLabel string `json:"providerLabel"`
ReqKey string `json:"reqKey,omitempty"`
CreatedAt string `json:"createdAt"`
}
type TrendPoint struct {
Date string `json:"date"`
Label string `json:"label"`
Count int `json:"count"`
}
type Summary struct {
Total int `json:"total"`
ActiveAccounts int `json:"activeAccounts"`
ActiveOrganizations int `json:"activeOrganizations"`
AveragePerDay float64 `json:"averagePerDay"`
}
type organizationRow struct {
OrganizationID string `json:"organizationId"`
OrganizationName string `json:"organizationName"`
Count int `json:"count"`
AccountCount int `json:"accountCount"`
LastUsedAt string `json:"lastUsedAt,omitempty"`
accounts map[string]bool
}
type accountRow struct {
OwnerID string `json:"ownerId"`
AccountName string `json:"accountName"`
AccountUsername string `json:"accountUsername,omitempty"`
OrganizationID string `json:"organizationId"`
OrganizationName string `json:"organizationName"`
Count int `json:"count"`
LastUsedAt string `json:"lastUsedAt,omitempty"`
}
type Options struct {
Organizations []Option `json:"organizations"`
Accounts []Option `json:"accounts"`
Capabilities []Option `json:"capabilities"`
Providers []Option `json:"providers"`
}
type Option struct {
Value string `json:"value"`
Label string `json:"label"`
}
type PersonalReport struct {
Preset Preset `json:"preset"`
Range DateRange `json:"range"`
Total int `json:"total"`
ByCapability []CountItem `json:"byCapability"`
Recent []Record `json:"recent"`
}
type AdminReport struct {
Range DateRange `json:"range"`
Summary Summary `json:"summary"`
Trend []TrendPoint `json:"trend"`
ByCapability []CountItem `json:"byCapability"`
ByProvider []CountItem `json:"byProvider"`
Organizations any `json:"organizations"`
Accounts any `json:"accounts"`
Recent any `json:"recent"`
Options Options `json:"options"`
}
type Reporter interface {
Personal(context.Context, PersonalRequest) (PersonalReport, error)
Admin(context.Context, AdminRequest) (AdminReport, error)
}
func (s Service) Record(event Event) (*Event, error) {
if event.Source == "api" || event.Provider == "mock" {
return nil, nil
}
saved, _, err := s.Repository.Insert(event)
if err != nil {
return nil, err
}
return &saved, nil
}
func (s Service) Report(who Requester, filters Filters) ([]Event, error) {
switch who.Role {
case "super_admin":
case "organization_admin":
if filters.OrganizationID != "" && filters.OrganizationID != who.OrganizationID {
return nil, ErrForbidden
}
filters.OrganizationID = who.OrganizationID
case "user":
filters.OwnerID = who.AccountID
default:
filters.OwnerID = who.AccountID
}
events, err := s.Repository.List(filters)
if err != nil {
return nil, err
}
seen := map[string]bool{}
out := make([]Event, 0, len(events))
for _, e := range events {
if e.Source == "api" || e.Provider == "mock" || seen[e.JobID] || filters.OwnerID != "" && e.OwnerID != filters.OwnerID || filters.OrganizationID != "" && e.OrganizationID != filters.OrganizationID || filters.Capability != "" && e.Capability != filters.Capability || filters.Provider != "" && e.Provider != filters.Provider {
continue
}
seen[e.JobID] = true
if who.Role == "organization_admin" {
e.AccountUsername = ""
}
out = append(out, e)
}
sort.SliceStable(out, func(i, j int) bool { return out[i].CreatedAt > out[j].CreatedAt })
return out, nil
}