llmgw/internal/store/selfusage.go

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?")
	}
	src, args := usageSince(since)
	places, err := queryAll(ctx, s.db, `
		SELECT COALESCE(NULLIF(u.project_id, ''), ''),
		       COALESCE(o.name || ' / ' || t.name || ' / ' || p.name, ''),
		       SUM(u.requests), SUM(u.prompt_tokens + u.completion_tokens), SUM(u.cost_micros)
		FROM `+src+` 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 = ?
		GROUP BY 1, 2
		ORDER BY 4 DESC`, append(args, userID),
		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()
	q, from := `
		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`, millis(since)
	if t, ok := tierForSlices(width); ok {
		q, from = `
		SELECT (bucket / ?) * ?, SUM(requests), SUM(errors),
		       SUM(prompt_tokens), SUM(completion_tokens), SUM(cost_micros)
		FROM `+t.table+` WHERE user_id = ? AND bucket >= ?
		GROUP BY 1 ORDER BY 1`, t.floor(since)
	}
	got, err := queryAll(ctx, s.db, q, []any{ms, ms, userID, from},
		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
	src, srcArgs := usageSince(since)
	err := s.db.QueryRowContext(ctx,
		`SELECT SUM(u.cost_micros) FROM `+src+` u WHERE `+where,
		append(srcArgs, args...)...).Scan(&total)
	return total.Int64, err
}