llmgw/internal/store/usageseries.go

package store

import (
	"context"
	"fmt"
	"time"
)

// Grouping is how a usage report is broken down.
//
// Not store.Scope with one more value, although three of the four agree with
// it. Scope is where a limit hangs, and nothing hangs off a person: one person
// holds keys in several projects, so a per-user total crosses the tree rather
// than naming a level of it. Keeping the two apart means neither has to carry a
// case the other does not have.
type Grouping string

// The groupings a usage report can be asked for.
const (
	GroupByOrg     Grouping = "org"
	GroupByTeam    Grouping = "team"
	GroupByProject Grouping = "project"
	GroupByUser    Grouping = "user"
	// GroupByModel is by the model that answered, which differs from the one
	// asked for when a fallback served it.
	GroupByModel Grouping = "model"
)

// Valid reports whether this is a grouping usage can be read by.
func (g Grouping) Valid() bool {
	_, ok := groupColumns[g]
	return ok
}

// Groupings lists them in the order a page should offer them: widest first,
// then the one that is not a level at all.
func Groupings() []Grouping {
	return []Grouping{GroupByOrg, GroupByTeam, GroupByProject, GroupByUser, GroupByModel}
}

// UsageSeries is one group's usage over time.
//
// Buckets are contiguous and gap-filled: a group that stopped using the cluster
// has zeroes rather than an absence, because a line drawn straight across a gap
// says the traffic continued. It also means every series in a report shares one
// set of timestamps, which is what lets a chart sum the tail of them.
type UsageSeries struct {
	Grouping Grouping
	ID       string
	Name     string

	Buckets []Bucket

	// Totals over the whole window, so a table beside the chart does not have
	// to add the buckets up and possibly disagree with them.
	Requests   int64
	Prompt     int64
	Completion int64
	CostMicros int64
	Errors     int64
}

// Tokens is what the group consumed, prompt and completion together.
func (u UsageSeries) Tokens() int64 { return u.Prompt + u.Completion }

// groupColumns is the id and display name for each grouping, and how usage
// reaches it.
//
// Written out per grouping rather than assembled from parts: these are table
// and column names going into SQL, and the one place they must not come from
// anywhere but here is the place they are interpolated. A grouping that is not
// in this map is refused before a query is built.
//
// Read from usage.project_id, for the reason UsageByScope is: joined through
// api_keys, the chat page's requests — which have no key row — were missing
// from every grouping but the per-person one.
var groupColumns = map[Grouping]struct{ id, name, joins string }{
	GroupByOrg: {"o.id", "o.name", `
		JOIN projects p      ON p.id = NULLIF(u.project_id, '')
		JOIN teams t         ON t.id = p.team_id
		JOIN organizations o ON o.id = t.org_id`},
	GroupByTeam: {"t.id", "t.name", `
		JOIN projects p ON p.id = NULLIF(u.project_id, '')
		JOIN teams t    ON t.id = p.team_id`},
	GroupByProject: {"p.id", "p.name", `
		JOIN projects p ON p.id = NULLIF(u.project_id, '')`},
	// Who spent it, as recorded with the spend — which includes personal keys
	// and the chat page, neither of which a key's project reaches. Traffic
	// from keys with no owner groups under an empty id, which the page names
	// rather than hides.
	GroupByUser: {"u.user_id", "COALESCE(us.email, '')", `
		LEFT JOIN users us ON us.id = NULLIF(u.user_id, '')`},
	GroupByModel: {servedModel, servedModel, ""},
}

// servedModel is what answered: the served model, or the one asked for on a
// row from before the two were told apart.
const servedModel = "CASE WHEN u.served_model != '' THEN u.served_model ELSE u.model END"

// UsageBucketsBy returns each group's usage over time.
//
// One query for every group rather than one per group: the page shows them
// together, and a query per organization is a loop whose cost grows with the
// thing being reported on.
func (s *DB) UsageBucketsBy(
	ctx context.Context, g Grouping, since time.Time, width time.Duration,
) ([]UsageSeries, error) {
	return s.usageBuckets(ctx, g, "", since, width)
}

// UsageBucketsIn is UsageBucketsBy for one organization: only what was spent
// in its projects, grouped by its teams, projects or people. A person's use
// elsewhere — another organization, or on their own — is not this
// organization's to see.
func (s *DB) UsageBucketsIn(
	ctx context.Context, g Grouping, orgID string, since time.Time, width time.Duration,
) ([]UsageSeries, error) {
	if orgID == "" {
		return nil, fmt.Errorf("which organization?")
	}
	return s.usageBuckets(ctx, g, orgID, since, width)
}

func (s *DB) usageBuckets(
	ctx context.Context, g Grouping, orgID string, since time.Time, width time.Duration,
) ([]UsageSeries, error) {
	c, ok := groupColumns[g]
	if !ok {
		return nil, fmt.Errorf(
			"usage cannot be grouped by %q; try org, team, project or user", g)
	}
	if width <= 0 {
		width = time.Hour
	}
	ms := width.Milliseconds()
	inOrg := ""
	if orgID != "" {
		// By where it was spent, as recorded with the spend.
		inOrg = ` AND NULLIF(u.project_id, '') IN (
			SELECT p.id FROM projects p JOIN teams t ON t.id = p.team_id WHERE t.org_id = ?)`
	}

	// Grouped by ordinal. PostgreSQL matches a grouped expression to a selected
	// one syntactically, so "(u.ts / ?) * ?" written twice is two expressions
	// as far as it is concerned; the ordinal is the same one by construction,
	// and both dialects accept it.
	q := `
		SELECT ` + c.id + `, ` + c.name + `, (u.ts / ?) * ?,
		       COUNT(*), SUM(u.prompt_tokens), SUM(u.completion_tokens),
		       SUM(u.cost_micros),
		       SUM(CASE WHEN u.status >= 400 THEN 1 ELSE 0 END)
		FROM usage u` + c.joins + `
		WHERE u.ts >= ?` + inOrg + `
		GROUP BY 1, 2, 3
		ORDER BY 1, 3`
	from := millis(since)
	// From the rollup for slices that are whole rollup slices — every
	// window the pages offer — and from the rows for anything finer.
	if t, ok := tierForSlices(width); ok {
		q = `
		SELECT ` + c.id + `, ` + c.name + `, (u.bucket / ?) * ?,
		       SUM(u.requests), SUM(u.prompt_tokens), SUM(u.completion_tokens),
		       SUM(u.cost_micros), SUM(u.errors)
		FROM ` + t.table + ` u` + c.joins + `
		WHERE u.bucket >= ?` + inOrg + `
		GROUP BY 1, 2, 3
		ORDER BY 1, 3`
		from = t.floor(since)
	}

	args := []any{ms, ms, from}
	if orgID != "" {
		args = append(args, orgID)
	}
	rows, err := s.db.QueryContext(ctx, q, args...)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	type cell struct {
		name string
		at   map[int64]Bucket
	}
	found := map[string]*cell{}
	var order []string

	for rows.Next() {
		var id, name string
		var start int64
		var b Bucket
		if err := rows.Scan(&id, &name, &start,
			&b.Requests, &b.Prompt, &b.Completion, &b.CostMicros, &b.Errors); err != nil {
			return nil, err
		}
		b.Start = fromMillis(start)
		e, ok := found[id]
		if !ok {
			e = &cell{name: name, at: map[int64]Bucket{}}
			found[id] = e
			order = append(order, id)
		}
		e.at[start] = b
	}
	if err := rows.Err(); err != nil {
		return nil, err
	}

	first := (millis(since) / ms) * ms
	last := (millis(time.Now()) / ms) * ms

	out := make([]UsageSeries, 0, len(order))
	for _, id := range order {
		e := found[id]
		u := UsageSeries{Grouping: g, ID: id, Name: e.name}
		for t := first; t <= last; t += ms {
			b, ok := e.at[t]
			if !ok {
				b = Bucket{Start: fromMillis(t)}
			}
			u.Buckets = append(u.Buckets, b)
			u.Requests += b.Requests
			u.Prompt += b.Prompt
			u.Completion += b.Completion
			u.CostMicros += b.CostMicros
			u.Errors += b.Errors
		}
		out = append(out, u)
	}
	return out, nil
}