llmgw/internal/store/usagescope.go

package store

import (
	"context"
	"database/sql"
	"fmt"
	"time"
)

// ScopeUsage is what one group of the tenancy tree spent.
//
// The report that was missing. Quotas could be written on an organization, a
// team or a project, and nothing could show what any of them had used: usage
// came back per key, per user and per model only. A quota you cannot see the
// consumption of is a quota nobody reviews until it refuses somebody.
type ScopeUsage struct {
	Scope Scope
	ID    string
	Name  string

	Keys       int64
	Requests   int64
	Prompt     int64
	Completion int64
	CostMicros int64
	Errors     int64
	LastAt     time.Time

	// Budget is what was set at this scope, in micro-units, and nil when this
	// level sets none.
	//
	// What is set here, not what binds. A team with no budget of its own may
	// still be bounded by its organization's, and printing the inherited figure
	// beside every team would read as though each had that much to itself —
	// which is the misunderstanding this whole change exists to remove.
	Budget       *int64
	BudgetPeriod Period
}

// Spent reports the fraction of this scope's budget used, and false when there
// is no budget here to be a fraction of.
func (s ScopeUsage) Spent() (float64, bool) {
	if s.Budget == nil || *s.Budget <= 0 {
		return 0, false
	}
	return float64(s.CostMicros) / float64(*s.Budget), true
}

// scopeColumns is the group id and name for each level, and how usage reaches
// it. Written out per level rather than built by string pasting: 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.
var scopeColumns = map[Scope]struct{ id, name, joins string }{
	ScopeOrg: {"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`},
	ScopeTeam: {"t.id", "t.name", `
		JOIN projects p ON p.id = NULLIF(u.project_id, '')
		JOIN teams t    ON t.id = p.team_id`},
	ScopeProject: {"p.id", "p.name", `
		JOIN projects p ON p.id = NULLIF(u.project_id, '')`},
}

// UsageByScope totals usage for one level of the tenancy tree.
//
// Read from usage.project_id, the project the request was spent in. It was
// joined through api_keys, which dropped every request the chat page made —
// those have no key row — so a project's spend on this page and its budget
// counter disagreed by exactly its chat use. Keys no longer move between
// projects, so where a key is and where its requests were spent are the same;
// a project moved to another team takes its history with it through the join.
func (s *DB) UsageByScope(ctx context.Context, scope Scope, since time.Time) ([]ScopeUsage, error) {
	c, ok := scopeColumns[scope]
	if !ok {
		return nil, fmt.Errorf("usage cannot be grouped by %q; try org, team or project", scope)
	}

	// Every non-aggregated column is grouped, and the budget is joined by the
	// same expression it is grouped on. PostgreSQL requires both; SQLite
	// silently picks a row for whatever is missing, which is how three of these
	// queries shipped working in tests and failing on the cluster.
	src, args := usageSince(since)
	q := `
		SELECT ` + c.id + `, ` + c.name + `,
		       COUNT(DISTINCT u.key_id),
		       SUM(u.requests), SUM(u.prompt_tokens), SUM(u.completion_tokens),
		       SUM(u.cost_micros), SUM(u.errors), MAX(u.last_ts),
		       MAX(l.budget_micros), MAX(COALESCE(l.budget_period, ''))
		FROM ` + src + ` u` + c.joins + `
		LEFT JOIN limits l ON l.scope_type = ? AND l.scope_id = ` + c.id + `
		GROUP BY ` + c.id + `, ` + c.name + `
		ORDER BY SUM(u.cost_micros) DESC, SUM(u.requests) DESC`

	rows, err := s.db.QueryContext(ctx, q, append(args, string(scope))...)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var out []ScopeUsage
	for rows.Next() {
		su := ScopeUsage{Scope: scope}
		var last int64
		var budget sql.NullInt64
		var period string
		if err := rows.Scan(&su.ID, &su.Name, &su.Keys,
			&su.Requests, &su.Prompt, &su.Completion, &su.CostMicros, &su.Errors,
			&last, &budget, &period); err != nil {
			return nil, err
		}
		su.LastAt = fromMillis(last)
		if budget.Valid {
			v := budget.Int64
			su.Budget, su.BudgetPeriod = &v, Period(period)
		}
		out = append(out, su)
	}
	return out, rows.Err()
}