llmgw/internal/store/rollup.go

package store

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

// tier is one table usage is summed into as it is written: slices of a fixed
// width, per key, person, project and model.
type tier struct {
	table string
	ms    int64
}

var (
	tier5m = tier{"usage_5m", int64(5 * time.Minute / time.Millisecond)}
	tier1h = tier{"usage_1h", int64(time.Hour / time.Millisecond)}
	tier1d = tier{"usage_1d", int64(24 * time.Hour / time.Millisecond)}
	// tiers, finest first.
	tiers = []tier{tier5m, tier1h, tier1d}
)

// floor is the start of the slice a moment falls in. A chart read from a
// tier starts its first slice there, as a chart read from the rows always did.
func (t tier) floor(at time.Time) int64 { return (millis(at) / t.ms) * t.ms }

// ceil is the start of the first slice wholly after a moment.
func (t tier) ceil(at time.Time) int64 {
	return ((millis(at) + t.ms - 1) / t.ms) * t.ms
}

// tierForSlices is the coarsest tier whose slices make up a chart slice of
// width exactly, and false when none does — a slice finer than five minutes,
// or not a whole number of them, is read from the rows.
func tierForSlices(width time.Duration) (tier, bool) {
	ms := width.Milliseconds()
	for i := len(tiers) - 1; i >= 0; i-- {
		if ms >= tiers[i].ms && ms%tiers[i].ms == 0 {
			return tiers[i], true
		}
	}
	return tier{}, false
}

// LatencyBounds are the upper edges, in milliseconds, of the histogram buckets
// latency and time to first token are counted in; one more bucket holds
// everything above the last. Fixed, because buckets are only summable across
// slices, keys and teams when every row has the same ones.
var LatencyBounds = []int64{100, 250, 500, 1000, 2500, 5000, 10000, 30000, 60000, 120000}

// latencyBucket is the histogram bucket a duration is counted in.
func latencyBucket(ms int64) int {
	for i, b := range LatencyBounds {
		if ms <= b {
			return i
		}
	}
	return len(LatencyBounds)
}

// rollupKey is one row of a tier: a slice, and who spent what on what.
type rollupKey struct {
	bucket                                       int64
	keyID, userID, projectID, model, servedModel string
}

// rollupRow is what a batch adds to one row.
type rollupRow struct {
	requests, errors, limited, prompt int64
	completion, cost                  int64
	latency, queue, ttft, ttftN, last int64
	lat, ttftH                        []int64
}

// rollupBatch sums a batch of records into the rows of one tier they add to.
func rollupBatch(t tier, records []Usage) (map[rollupKey]*rollupRow, []rollupKey) {
	rows := map[rollupKey]*rollupRow{}
	var order []rollupKey
	n := len(LatencyBounds) + 1
	for _, r := range records {
		ts := millis(r.At)
		k := rollupKey{(ts / t.ms) * t.ms, r.KeyID, r.UserID, r.ProjectID, r.Model, r.ServedModel}
		row, ok := rows[k]
		if !ok {
			row = &rollupRow{lat: make([]int64, n), ttftH: make([]int64, n)}
			rows[k] = row
			order = append(order, k)
		}
		row.requests++
		if r.Status >= 400 {
			row.errors++
		}
		if r.Status == 429 {
			row.limited++
		}
		row.prompt += r.PromptTokens
		row.completion += r.CompletionTokens
		row.cost += r.CostMicros
		row.latency += r.LatencyMs
		row.queue += r.QueueMs
		row.lat[latencyBucket(r.LatencyMs)]++
		if r.TTFTMs > 0 {
			row.ttft += r.TTFTMs
			row.ttftN++
			row.ttftH[latencyBucket(r.TTFTMs)]++
		}
		row.last = max(row.last, ts)
	}
	return rows, order
}

// rollupColumns are a tier's columns, in the order rollupUpsert takes them.
var rollupColumns = func() []string {
	n := len(LatencyBounds) + 1
	cols := []string{"bucket", "key_id", "user_id", "project_id", "model", "served_model",
		"requests", "errors", "limited", "prompt_tokens", "completion_tokens", "cost_micros",
		"latency_ms_sum", "queue_ms_sum", "ttft_ms_sum", "ttft_n", "last_ts"}
	for i := range n {
		cols = append(cols, fmt.Sprintf("lat_%d", i))
	}
	for i := range n {
		cols = append(cols, fmt.Sprintf("ttft_%d", i))
	}
	return cols
}()

// rollupUpsert is the statement that adds a batch's sums to a tier.
func rollupUpsert(table string) string {
	var sets []string
	for _, c := range rollupColumns[6:] {
		if c == "last_ts" {
			sets = append(sets, "last_ts = CASE WHEN excluded.last_ts > "+table+".last_ts "+
				"THEN excluded.last_ts ELSE "+table+".last_ts END")
			continue
		}
		sets = append(sets, c+" = "+table+"."+c+" + excluded."+c)
	}
	return `INSERT INTO ` + table + ` (` + strings.Join(rollupColumns, ", ") + `)
		VALUES (` + strings.TrimSuffix(strings.Repeat("?, ", len(rollupColumns)), ", ") + `)
		ON CONFLICT (bucket, key_id, user_id, project_id, model, served_model)
		DO UPDATE SET ` + strings.Join(sets, ", ")
}

// writeRollup adds a batch to every tier, inside the transaction that writes
// its rows, so the sums never disagree with them.
func writeRollup(ctx context.Context, tx conn, records []Usage) error {
	for _, t := range tiers {
		rows, order := rollupBatch(t, records)
		stmt, err := tx.PrepareContext(ctx, rollupUpsert(t.table))
		if err != nil {
			return err
		}
		for _, k := range order {
			r := rows[k]
			args := []any{k.bucket, k.keyID, k.userID, k.projectID, k.model, k.servedModel,
				r.requests, r.errors, r.limited, r.prompt, r.completion, r.cost,
				r.latency, r.queue, r.ttft, r.ttftN, r.last}
			for _, v := range r.lat {
				args = append(args, v)
			}
			for _, v := range r.ttftH {
				args = append(args, v)
			}
			if _, err := stmt.ExecContext(ctx, args...); err != nil {
				_ = stmt.Close()
				return fmt.Errorf("adding to %s: %w", t.table, err)
			}
		}
		if err := stmt.Close(); err != nil {
			return err
		}
	}
	return nil
}

// usageSince is every request since a moment, exactly, as a table with the
// columns a tier has (but the histograms): the rows up to the next
// five-minute slice, those slices up to the next hour, hours up to the next
// day, and days from there. Totals are read from this — a budget against it
// must be exact, and so is everything else, at the cost of a few rows read
// from each of four tables rather than a window rounded to a slice.
//
// Its placeholders come first in any query it is put into.
func usageSince(since time.Time) (string, []any) {
	from := millis(since)
	to5m, to1h, to1d := tier5m.ceil(since), tier1h.ceil(since), tier1d.ceil(since)
	const cols = `key_id, user_id, project_id, model, served_model, requests, errors, limited,
		prompt_tokens, completion_tokens, cost_micros, latency_ms_sum, last_ts`
	return `(SELECT key_id, user_id, project_id, model, served_model, 1 AS requests,
			CASE WHEN status >= 400 THEN 1 ELSE 0 END AS errors,
			CASE WHEN status = 429 THEN 1 ELSE 0 END AS limited,
			prompt_tokens, completion_tokens, cost_micros, latency_ms AS latency_ms_sum, ts AS last_ts
		FROM usage WHERE ts >= ? AND ts < ?
		UNION ALL SELECT ` + cols + ` FROM usage_5m WHERE bucket >= ? AND bucket < ?
		UNION ALL SELECT ` + cols + ` FROM usage_1h WHERE bucket >= ? AND bucket < ?
		UNION ALL SELECT ` + cols + ` FROM usage_1d WHERE bucket >= ?)`,
		[]any{from, to5m, to5m, to1h, to1h, to1d, to1d}
}

// PruneUsage deletes request rows older than a moment. The tiers keep their
// sums, so charts, totals and budgets are unchanged; what goes is the record
// of each request — `ctl request` on one that old finds nothing.
func (s *DB) PruneUsage(ctx context.Context, before time.Time) (int64, error) {
	res, err := s.db.ExecContext(ctx, `DELETE FROM usage WHERE ts < ?`, millis(before))
	if err != nil {
		return 0, err
	}
	return res.RowsAffected()
}