llmgw/internal/store/latency.go

package store

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

// Quantiles are a duration's percentiles, in milliseconds, read from the
// histogram kept with the usage sums. Within a bucket the value is
// interpolated; above the last bucket's edge it is that edge, and Over says
// so. N is how many durations they are read from — zero is no reading.
type Quantiles struct {
	P50, P95, P99 int64
	N             int64
	Over          bool
}

// GroupStats is one group's requests, how they ended, and how long they took.
type GroupStats struct {
	ID       string
	Requests int64
	Errors   int64
	// Limited is requests refused by a rate, token or budget limit.
	Limited int64
	Latency Quantiles
	// TTFT is time to first token, for requests that streamed one.
	TTFT Quantiles
}

// ErrorRate is the share of requests that failed, limits included, in percent.
func (g GroupStats) ErrorRate() float64 {
	if g.Requests == 0 {
		return 0
	}
	return 100 * float64(g.Errors) / float64(g.Requests)
}

// StatsBy is each group's request outcomes and latency percentiles since a
// moment, keyed by the group's id.
//
// Read from the hourly sums for a window of two days or more and from the
// five-minute ones otherwise, starting at the slice the moment falls in:
// percentiles are a shape, and a window starting up to one slice early does
// not change one. Budgets, which must be exact, are not read this way.
func (s *DB) StatsBy(ctx context.Context, g Grouping, since time.Time) (map[string]GroupStats, error) {
	c, ok := groupColumns[g]
	if !ok {
		return nil, fmt.Errorf("usage cannot be grouped by %q", g)
	}
	t := tier5m
	if time.Since(since) >= 48*time.Hour {
		t = tier1h
	}
	n := len(LatencyBounds) + 1
	var sums []string
	for i := range n {
		sums = append(sums, fmt.Sprintf("SUM(u.lat_%d)", i))
	}
	for i := range n {
		sums = append(sums, fmt.Sprintf("SUM(u.ttft_%d)", i))
	}
	rows, err := s.db.QueryContext(ctx, `
		SELECT `+c.id+`, SUM(u.requests), SUM(u.errors), SUM(u.limited), `+strings.Join(sums, ", ")+`
		FROM `+t.table+` u`+c.joins+`
		WHERE u.bucket >= ?
		GROUP BY 1`, t.floor(since))
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	out := map[string]GroupStats{}
	for rows.Next() {
		var gs GroupStats
		lat, ttft := make([]int64, n), make([]int64, n)
		dst := []any{&gs.ID, &gs.Requests, &gs.Errors, &gs.Limited}
		for i := range lat {
			dst = append(dst, &lat[i])
		}
		for i := range ttft {
			dst = append(dst, &ttft[i])
		}
		if err := rows.Scan(dst...); err != nil {
			return nil, err
		}
		gs.Latency, gs.TTFT = quantiles(lat), quantiles(ttft)
		out[gs.ID] = gs
	}
	return out, rows.Err()
}

// quantiles reads p50, p95 and p99 from counts per LatencyBounds bucket.
func quantiles(counts []int64) Quantiles {
	var q Quantiles
	for _, c := range counts {
		q.N += c
	}
	if q.N == 0 {
		return q
	}
	at := func(p float64) (int64, bool) {
		rank := p * float64(q.N)
		var seen float64
		for i, c := range counts {
			if c == 0 {
				continue
			}
			if seen+float64(c) >= rank {
				if i >= len(LatencyBounds) {
					return LatencyBounds[len(LatencyBounds)-1], true
				}
				lo := int64(0)
				if i > 0 {
					lo = LatencyBounds[i-1]
				}
				hi := LatencyBounds[i]
				frac := (rank - seen) / float64(c)
				return lo + int64(frac*float64(hi-lo)), false
			}
			seen += float64(c)
		}
		return LatencyBounds[len(LatencyBounds)-1], true
	}
	var over bool
	q.P50, _ = at(0.50)
	q.P95, _ = at(0.95)
	q.P99, over = at(0.99)
	q.Over = over
	return q
}