package store import ( "context" "database/sql" "fmt" "strings" "time" ) // PlaceUsage is one person's use in one place: a project, or on their own. type PlaceUsage struct { // ProjectID is empty for personal use. ProjectID string // Path names the project; empty for personal use. Path string Requests int64 Tokens int64 CostMicros int64 } // UsageOfUser is one person's own use since a time: where it went, and how it // ran over time in slices of width. What a person sees of themselves — every // place they spent, and nothing anybody else did. func (s *DB) UsageOfUser( ctx context.Context, userID string, since time.Time, width time.Duration, ) ([]PlaceUsage, []Bucket, error) { if userID == "" { return nil, nil, fmt.Errorf("whose usage?") } places, err := queryAll(ctx, s.db, ` SELECT COALESCE(NULLIF(u.project_id, ''), ''), COALESCE(o.name || ' / ' || t.name || ' / ' || p.name, ''), COUNT(*), SUM(u.prompt_tokens + u.completion_tokens), SUM(u.cost_micros) FROM usage u LEFT JOIN projects p ON p.id = NULLIF(u.project_id, '') LEFT JOIN teams t ON t.id = p.team_id LEFT JOIN organizations o ON o.id = t.org_id WHERE u.user_id = ? AND u.ts >= ? GROUP BY 1, 2 ORDER BY 4 DESC`, []any{userID, millis(since)}, func(rows *sql.Rows) (PlaceUsage, error) { var pu PlaceUsage return pu, rows.Scan(&pu.ProjectID, &pu.Path, &pu.Requests, &pu.Tokens, &pu.CostMicros) }) if err != nil { return nil, nil, err } if width <= 0 { width = 24 * time.Hour } ms := width.Milliseconds() got, err := queryAll(ctx, s.db, ` SELECT (ts / ?) * ?, COUNT(*), SUM(CASE WHEN status >= 400 THEN 1 ELSE 0 END), SUM(prompt_tokens), SUM(completion_tokens), SUM(cost_micros) FROM usage WHERE user_id = ? AND ts >= ? GROUP BY 1 ORDER BY 1`, []any{ms, ms, userID, millis(since)}, func(rows *sql.Rows) (Bucket, error) { var b Bucket var start int64 err := rows.Scan(&start, &b.Requests, &b.Errors, &b.Prompt, &b.Completion, &b.CostMicros) b.Start = fromMillis(start) return b, err }) if err != nil { return nil, nil, err } // Gap-filled, for the reason UsageBuckets is: a day with nothing in it is // a day, not a hole. at := map[int64]Bucket{} for _, b := range got { at[millis(b.Start)] = b } var buckets []Bucket for t := (millis(since) / ms) * ms; t <= millis(time.Now()); t += ms { b, ok := at[t] if !ok { b = Bucket{Start: fromMillis(t)} } buckets = append(buckets, b) } return places, buckets, nil } // SpendOfPool is what has been spent since a time from the pool a limit // counts against — the figure its budget is compared with. // // The pool follows the level: a project, team or organization's budget is // spent by everyone beneath it; a person's by every key they own and the chat // page; a per-member limit by one person within one scope, whose ScopeID is // "scopeID/userID". A key's own limit is not answered here, and gives zero. func (s *DB) SpendOfPool(ctx context.Context, l Limits, since time.Time) (int64, error) { under := func(scope Scope) (string, error) { switch scope { case ScopeProject: return `NULLIF(u.project_id, '') = ?`, nil case ScopeTeam: return `NULLIF(u.project_id, '') IN (SELECT id FROM projects WHERE team_id = ?)`, nil case ScopeOrg: return `NULLIF(u.project_id, '') IN (SELECT p.id FROM projects p JOIN teams t ON t.id = p.team_id WHERE t.org_id = ?)`, nil } return "", fmt.Errorf("no pool for %q", scope) } var where string var args []any switch l.Scope { case ScopeUser: where, args = `u.user_id = ?`, []any{l.ScopeID} case ScopeOrg, ScopeTeam, ScopeProject: w, err := under(l.Scope) if err != nil { return 0, err } where, args = w, []any{l.ScopeID} case ScopeOrgMember, ScopeTeamMember, ScopeProjectMember: scopeID, userID, ok := strings.Cut(l.ScopeID, "/") if !ok { return 0, fmt.Errorf("a per-member pool is scope/user, not %q", l.ScopeID) } level := map[Scope]Scope{ScopeOrgMember: ScopeOrg, ScopeTeamMember: ScopeTeam, ScopeProjectMember: ScopeProject}[l.Scope] w, err := under(level) if err != nil { return 0, err } where, args = `u.user_id = ? AND `+w, []any{userID, scopeID} default: return 0, nil } var total sql.NullInt64 err := s.db.QueryRowContext(ctx, `SELECT SUM(u.cost_micros) FROM usage u WHERE `+where+` AND u.ts >= ?`, append(args, millis(since))...).Scan(&total) return total.Int64, err }