diff --git a/docs/admin-api.openapi.json b/docs/admin-api.openapi.json index c3b8f905..869d9dd7 100644 --- a/docs/admin-api.openapi.json +++ b/docs/admin-api.openapi.json @@ -10288,6 +10288,10 @@ "CostMicros": { "type": "integer", "format": "int64" + }, + "Images": { + "type": "integer", + "format": "int64" } } }, diff --git a/internal/control/adminapi/export.go b/internal/control/adminapi/export.go index 80988381..44cbe541 100644 --- a/internal/control/adminapi/export.go +++ b/internal/control/adminapi/export.go @@ -75,11 +75,11 @@ func (a *API) usageExport(w http.ResponseWriter, r *http.Request) { func WriteExportCSV(w io.Writer, rows []store.ExportRow) { c := csv.NewWriter(w) _ = c.Write([]string{"period_start", "org", "team", "project", "user", "key", "model", - "requests", "errors", "limited", "prompt_tokens", "completion_tokens", "cost_micros", "cost"}) + "requests", "errors", "limited", "prompt_tokens", "completion_tokens", "cost_micros", "cost", "images"}) for _, r := range rows { _ = c.Write([]string{r.Period.Format(time.RFC3339), safe(r.Org), safe(r.Team), safe(r.Project), safe(r.User), r.Key, safe(r.Model), i(r.Requests), i(r.Errors), i(r.Limited), i(r.Prompt), i(r.Completion), - i(r.CostMicros), strconv.FormatFloat(float64(r.CostMicros)/1e6, 'f', 6, 64)}) + i(r.CostMicros), strconv.FormatFloat(float64(r.CostMicros)/1e6, 'f', 6, 64), i(r.Images)}) } c.Flush() } diff --git a/internal/store/export.go b/internal/store/export.go index 1058d695..7048a703 100644 --- a/internal/store/export.go +++ b/internal/store/export.go @@ -25,6 +25,7 @@ type ExportRow struct { Prompt int64 Completion int64 CostMicros int64 + Images int64 } // ExportDims are what an export can be broken down by. @@ -51,16 +52,18 @@ func (s *DB) UsageExport(ctx context.Context, from, to time.Time, period string, 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(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.requests, &x.errors, &x.limited, &x.prompt, &x.completion, &x.cost, &x.images) return x, err }) if err != nil { @@ -137,6 +140,7 @@ func (s *DB) UsageExport(ctx context.Context, from, to time.Time, period string, 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 { diff --git a/internal/store/export_test.go b/internal/store/export_test.go index 6b19bdfb..fec5e331 100644 --- a/internal/store/export_test.go +++ b/internal/store/export_test.go @@ -46,3 +46,36 @@ func TestUsageIsExportedByPeriodAndWhatWasAskedFor(t *testing.T) { t.Error("an unknown dimension was taken") } } + +func TestPicturesAreCountedInTheSumsAndTheExport(t *testing.T) { + // The request rows said how many pictures each made, and the sums did + // not: once the rows were past their retention, an export said what the + // pictures cost and not how many there were. + db := capacityStore(t) + ctx := t.Context() + at := time.Date(2026, 10, 10, 9, 0, 0, 0, time.UTC) + if err := db.WriteUsage(ctx, []Usage{ + {At: at, KeyID: "k", Model: "flux", ServedModel: "flux", PromptTokens: 9, Images: 2, CostMicros: 40000, Status: 200}, + {At: at.Add(time.Minute), KeyID: "k", Model: "flux", ServedModel: "flux", PromptTokens: 9, Images: 3, CostMicros: 60000, Status: 200}, + {At: at.Add(2 * time.Minute), KeyID: "k", Model: "chat", ServedModel: "chat", PromptTokens: 9, Status: 200}, + }); err != nil { + t.Fatal(err) + } + // Past their retention: only the sums are left. + if _, err := db.PruneUsage(ctx, at.Add(time.Hour)); err != nil { + t.Fatal(err) + } + for _, period := range []string{"hour", "day"} { + rows, err := db.UsageExport(ctx, at.Add(-24*time.Hour), at.Add(48*time.Hour), period, []string{"model"}) + if err != nil { + t.Fatal(err) + } + got := map[string]int64{} + for _, r := range rows { + got[r.Model] += r.Images + } + if got["flux"] != 5 || got["chat"] != 0 { + t.Errorf("%s: pictures by model %v, want flux 5", period, got) + } + } +} diff --git a/internal/store/migrations/postgres/0042_rollup_images.sql b/internal/store/migrations/postgres/0042_rollup_images.sql new file mode 100644 index 00000000..75005a56 --- /dev/null +++ b/internal/store/migrations/postgres/0042_rollup_images.sql @@ -0,0 +1,23 @@ +-- Ported from the SQLite migration of the same number. +-- How many pictures each slice's requests made, beside their tokens and +-- cost, so the sums and exports say it after the request rows are gone. +-- Pictures made before this were counted only in those rows: added from +-- them, for the slices since the first. +ALTER TABLE usage_5m ADD COLUMN images BIGINT NOT NULL DEFAULT 0; +ALTER TABLE usage_1h ADD COLUMN images BIGINT NOT NULL DEFAULT 0; +ALTER TABLE usage_1d ADD COLUMN images BIGINT NOT NULL DEFAULT 0; +UPDATE usage_5m SET images = COALESCE((SELECT SUM(u.images) FROM usage u + WHERE u.images > 0 AND u.ts >= usage_5m.bucket AND u.ts < usage_5m.bucket + 300000 + AND u.key_id = usage_5m.key_id AND u.user_id = usage_5m.user_id AND u.project_id = usage_5m.project_id + AND u.model = usage_5m.model AND u.served_model = usage_5m.served_model), 0) +WHERE bucket >= (SELECT COALESCE(MIN(ts), 9000000000000000000) - 300000 FROM usage WHERE images > 0); +UPDATE usage_1h SET images = COALESCE((SELECT SUM(u.images) FROM usage u + WHERE u.images > 0 AND u.ts >= usage_1h.bucket AND u.ts < usage_1h.bucket + 3600000 + AND u.key_id = usage_1h.key_id AND u.user_id = usage_1h.user_id AND u.project_id = usage_1h.project_id + AND u.model = usage_1h.model AND u.served_model = usage_1h.served_model), 0) +WHERE bucket >= (SELECT COALESCE(MIN(ts), 9000000000000000000) - 3600000 FROM usage WHERE images > 0); +UPDATE usage_1d SET images = COALESCE((SELECT SUM(u.images) FROM usage u + WHERE u.images > 0 AND u.ts >= usage_1d.bucket AND u.ts < usage_1d.bucket + 86400000 + AND u.key_id = usage_1d.key_id AND u.user_id = usage_1d.user_id AND u.project_id = usage_1d.project_id + AND u.model = usage_1d.model AND u.served_model = usage_1d.served_model), 0) +WHERE bucket >= (SELECT COALESCE(MIN(ts), 9000000000000000000) - 86400000 FROM usage WHERE images > 0); diff --git a/internal/store/migrations/sqlite/0042_rollup_images.sql b/internal/store/migrations/sqlite/0042_rollup_images.sql new file mode 100644 index 00000000..99d75b07 --- /dev/null +++ b/internal/store/migrations/sqlite/0042_rollup_images.sql @@ -0,0 +1,22 @@ +-- How many pictures each slice's requests made, beside their tokens and +-- cost, so the sums and exports say it after the request rows are gone. +-- Pictures made before this were counted only in those rows: added from +-- them, for the slices since the first. +ALTER TABLE usage_5m ADD COLUMN images INTEGER NOT NULL DEFAULT 0; +ALTER TABLE usage_1h ADD COLUMN images INTEGER NOT NULL DEFAULT 0; +ALTER TABLE usage_1d ADD COLUMN images INTEGER NOT NULL DEFAULT 0; +UPDATE usage_5m SET images = COALESCE((SELECT SUM(u.images) FROM usage u + WHERE u.images > 0 AND u.ts >= usage_5m.bucket AND u.ts < usage_5m.bucket + 300000 + AND u.key_id = usage_5m.key_id AND u.user_id = usage_5m.user_id AND u.project_id = usage_5m.project_id + AND u.model = usage_5m.model AND u.served_model = usage_5m.served_model), 0) +WHERE bucket >= (SELECT COALESCE(MIN(ts), 9000000000000000000) - 300000 FROM usage WHERE images > 0); +UPDATE usage_1h SET images = COALESCE((SELECT SUM(u.images) FROM usage u + WHERE u.images > 0 AND u.ts >= usage_1h.bucket AND u.ts < usage_1h.bucket + 3600000 + AND u.key_id = usage_1h.key_id AND u.user_id = usage_1h.user_id AND u.project_id = usage_1h.project_id + AND u.model = usage_1h.model AND u.served_model = usage_1h.served_model), 0) +WHERE bucket >= (SELECT COALESCE(MIN(ts), 9000000000000000000) - 3600000 FROM usage WHERE images > 0); +UPDATE usage_1d SET images = COALESCE((SELECT SUM(u.images) FROM usage u + WHERE u.images > 0 AND u.ts >= usage_1d.bucket AND u.ts < usage_1d.bucket + 86400000 + AND u.key_id = usage_1d.key_id AND u.user_id = usage_1d.user_id AND u.project_id = usage_1d.project_id + AND u.model = usage_1d.model AND u.served_model = usage_1d.served_model), 0) +WHERE bucket >= (SELECT COALESCE(MIN(ts), 9000000000000000000) - 86400000 FROM usage WHERE images > 0); diff --git a/internal/store/rollup.go b/internal/store/rollup.go index fa0e2964..fe0749d9 100644 --- a/internal/store/rollup.go +++ b/internal/store/rollup.go @@ -69,7 +69,7 @@ type rollupKey struct { // rollupRow is what a batch adds to one row. type rollupRow struct { requests, errors, limited, prompt int64 - completion, cost int64 + completion, cost, images int64 latency, queue, ttft, ttftN, last int64 lat, ttftH []int64 } @@ -98,6 +98,7 @@ func rollupBatch(t tier, records []Usage) (map[rollupKey]*rollupRow, []rollupKey row.prompt += r.PromptTokens row.completion += r.CompletionTokens row.cost += r.CostMicros + row.images += r.Images row.latency += r.LatencyMs row.queue += r.QueueMs row.lat[latencyBucket(r.LatencyMs)]++ @@ -116,7 +117,7 @@ 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"} + "latency_ms_sum", "queue_ms_sum", "ttft_ms_sum", "ttft_n", "last_ts", "images"} for i := range n { cols = append(cols, fmt.Sprintf("lat_%d", i)) } @@ -156,7 +157,7 @@ func writeRollup(ctx context.Context, tx conn, records []Usage) error { 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} + r.latency, r.queue, r.ttft, r.ttftN, r.last, r.images} for _, v := range r.lat { args = append(args, v) } @@ -187,11 +188,11 @@ 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` + prompt_tokens, completion_tokens, cost_micros, latency_ms_sum, last_ts, images` 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 + prompt_tokens, completion_tokens, cost_micros, latency_ms AS latency_ms_sum, ts AS last_ts, images 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 < ?