llmgw/internal/store/export.go

package store

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

// ExportRow is one line of a usage export: a period, the place, person, key
// and model it is about — those asked for, the rest empty — and what was used.
type ExportRow struct {
	Period     time.Time
	Org        string
	Team       string
	Project    string
	User       string
	Key        string
	Model      string
	Requests   int64
	Errors     int64
	Limited    int64
	Prompt     int64
	Completion int64
	CostMicros int64
	Images     int64
}

// ExportDims are what an export can be broken down by.
var ExportDims = []string{"org", "team", "project", "user", "key", "model"}

// UsageExport totals usage between two moments, by period — hour, day or
// month, in UTC — and by the dimensions asked for. It reads the hourly and
// daily sums, so request records past their retention are still counted.
func (s *DB) UsageExport(ctx context.Context, from, to time.Time, period string, by []string) ([]ExportRow, error) {
	for _, d := range by {
		if !slices.Contains(ExportDims, d) {
			return nil, fmt.Errorf("%q is not one of %s", d, strings.Join(ExportDims, ", "))
		}
	}
	table := "usage_1d"
	switch period {
	case "hour":
		table = "usage_1h"
	case "day", "month":
	default:
		return nil, fmt.Errorf("period %q is not hour, day or month", period)
	}
	type raw struct {
		bucket                                              int64
		key, user, project, model                           string
		requests, errors, limited, prompt, completion, cost int64
		images                                              int64
	}
	//nolint:gosec // table is one of two literals above
	rows, err := queryAll(ctx, s.db, `SELECT bucket, key_id, user_id, project_id, served_model,
		SUM(requests), SUM(errors), SUM(limited), SUM(prompt_tokens), SUM(completion_tokens), SUM(cost_micros),
		SUM(images)
		FROM `+table+` WHERE bucket >= ? AND bucket < ?
		GROUP BY bucket, key_id, user_id, project_id, served_model`,
		[]any{millis(from), millis(to)}, func(r *sql.Rows) (raw, error) {
			var x raw
			err := r.Scan(&x.bucket, &x.key, &x.user, &x.project, &x.model,
				&x.requests, &x.errors, &x.limited, &x.prompt, &x.completion, &x.cost, &x.images)
			return x, err
		})
	if err != nil {
		return nil, err
	}

	// Names, not ids: an export is read by people and spreadsheets.
	type place struct{ org, team, project string }
	places := map[string]place{"": {project: "personal"}}
	tree, err := s.Tree(ctx)
	if err != nil {
		return nil, err
	}
	for _, o := range tree {
		for _, t := range o.Teams {
			for _, p := range t.Projects {
				places[p.ID] = place{o.Name, t.Name, p.Name}
			}
		}
	}
	emails := map[string]string{}
	users, err := s.Users(ctx, "")
	if err != nil {
		return nil, err
	}
	for _, u := range users {
		emails[u.ID] = u.Email
	}
	has := func(d string) bool { return slices.Contains(by, d) }

	totals := map[ExportRow]*ExportRow{}
	var order []ExportRow
	for _, x := range rows {
		at := fromMillis(x.bucket).UTC()
		if period == "month" {
			at = time.Date(at.Year(), at.Month(), 1, 0, 0, 0, 0, time.UTC)
		}
		pl := places[x.project]
		k := ExportRow{Period: at}
		if has("org") {
			k.Org = pl.org
		}
		if has("team") {
			k.Team = pl.team
		}
		if has("project") {
			k.Project = pl.project
			if pl.project == "" {
				k.Project = x.project
			}
		}
		if has("user") {
			k.User = emails[x.user]
			if k.User == "" {
				k.User = x.user
			}
		}
		if has("key") {
			k.Key = x.key
		}
		if has("model") {
			k.Model = x.model
		}
		t, ok := totals[k]
		if !ok {
			row := k
			t = &row
			totals[k] = t
			order = append(order, k)
		}
		t.Requests += x.requests
		t.Errors += x.errors
		t.Limited += x.limited
		t.Prompt += x.prompt
		t.Completion += x.completion
		t.CostMicros += x.cost
		t.Images += x.images
	}
	out := make([]ExportRow, 0, len(order))
	for _, k := range order {
		out = append(out, *totals[k])
	}
	slices.SortFunc(out, func(a, b ExportRow) int {
		if c := a.Period.Compare(b.Period); c != 0 {
			return c
		}
		return strings.Compare(a.Org+a.Team+a.Project+a.User+a.Key+a.Model, b.Org+b.Team+b.Project+b.User+b.Key+b.Model)
	})
	return out, nil
}