ModelRouter/internal/gateway/stats.go

package gateway

import (
	"sync"
)

// Req is one recorded gateway request (audit trail + per-key/per-model stats).
type Req struct {
	Time  int64  `json:"time"`  // unix milliseconds
	Key   string `json:"key"`   // gateway key id
	Type  string `json:"type"`  // chat | stream | image
	Model string `json:"model"` // effective model used upstream
	// Source provider/source name
	Source string `json:"source"`
	// Prompt prompt tokens
	Prompt int64 `json:"prompt_tokens"`
	// Compl completion tokens
	Compl int64 `json:"completion_tokens"`
	// LatMs total handling time ms
	LatMs int64 `json:"latency_ms"`
	OK    bool   `json:"ok"`
	// Status http status code
	Status int `json:"status"`
	// Err short error message
	Err string `json:"error,omitempty"`
}

// Stat aggregates counters for one dimension row.
type Stat struct {
	Reqs   int64 `json:"reqs"`
	OK     int64 `json:"ok"`
	Err    int64 `json:"err"`
	Tokens int64 `json:"tokens"`
	Prompt int64 `json:"prompt_tokens"`
	Compl  int64 `json:"completion_tokens"`
	LatSum int64 `json:"latency_sum_ms"`
	LatMax int64 `json:"latency_max_ms"`
}

type agrRow struct {
	Name string `json:"name"`
	Stat
}

// Stats collects per-key / per-model / per-source aggregates plus a bounded
// ring of raw request records, all guarded by one mutex.
type Stats struct {
	mu         sync.Mutex
	active     int64
	byKey      map[string]*Stat
	byModel    map[string]*Stat
	bySrc      map[string]*Stat
	byKeyModel map[string]map[string]*Stat
	byKeySrc   map[string]map[string]*Stat
	recs       []Req
	maxRecs    int
}

func NewStats(maxRecords int) *Stats {
	if maxRecords <= 0 {
		maxRecords = 3000
	}
	return &Stats{
		byKey:      map[string]*Stat{},
		byModel:    map[string]*Stat{},
		bySrc:      map[string]*Stat{},
		byKeyModel: map[string]map[string]*Stat{},
		byKeySrc:   map[string]map[string]*Stat{},
		maxRecs:    maxRecords,
	}
}

// Begin accounts an in-flight request; the returned func must be called once
// the request finished (defer ok).
func (s *Stats) Begin() func() {
	s.mu.Lock()
	s.active++
	s.mu.Unlock()
	return func() {
		s.mu.Lock()
		s.active--
		s.mu.Unlock()
	}
}

func inc(m map[string]*Stat, name string, r Req) {
	a := m[name]
	if a == nil {
		a = &Stat{}
		m[name] = a
	}
	a.Reqs++
	if r.OK {
		a.OK++
	} else {
		a.Err++
	}
	a.Tokens += r.Prompt + r.Compl
	a.Prompt += r.Prompt
	a.Compl += r.Compl
	a.LatSum += r.LatMs
	if r.LatMs > a.LatMax {
		a.LatMax = r.LatMs
	}
}

// Record appends a finished request to the aggregates and ring buffer.
func (s *Stats) Record(r Req) {
	s.mu.Lock()
	defer s.mu.Unlock()
	inc(s.byKey, r.Key, r)
	inc(s.byModel, r.Model, r)
	inc(s.bySrc, r.Source, r)
	km := s.byKeyModel[r.Key]
	if km == nil {
		km = map[string]*Stat{}
		s.byKeyModel[r.Key] = km
	}
	inc(km, r.Model, r)
	ks := s.byKeySrc[r.Key]
	if ks == nil {
		ks = map[string]*Stat{}
		s.byKeySrc[r.Key] = ks
	}
	inc(ks, r.Source, r)
	s.recs = append(s.recs, r)
	if len(s.recs) > s.maxRecs {
		s.recs = s.recs[len(s.recs)-s.maxRecs:]
	}
}

func rows(m map[string]*Stat) []StatsRow {
	out := make([]StatsRow, 0, len(m))
	for k, v := range m {
		out = append(out, StatsRow{Name: k, Stat: *v})
	}
	for i := 1; i < len(out); i++ {
		for j := i; j > 0 && out[j].Reqs > out[j-1].Reqs; j-- {
			out[j], out[j-1] = out[j-1], out[j]
		}
	}
	return out
}

// StatsRow is one aggregated row for the dashboard.
type StatsRow struct {
	Name string `json:"name"`
	Stat
}

// Snapshot returns the whole dashboard payload; when key != "" the records
// and aggregate views are restricted to that gateway key.
func (s *Stats) Snapshot(limit int, key string) map[string]interface{} {
	s.mu.Lock()
	defer s.mu.Unlock()
	if limit <= 0 {
		limit = s.maxRecs
	}
	start := 0
	if len(s.recs) > limit {
		start = len(s.recs) - limit
	}
	recs := s.recs[start:]
	if key != "" {
		filt := recs[:0]
		for _, r := range recs {
			if r.Key == key {
				filt = append(filt, r)
			}
		}
		recs = filt
	}
	var total Stat
	var byModel, byKey, bySrc map[string]*Stat
	if key == "" {
		byKey, byModel, bySrc = s.byKey, s.byModel, s.bySrc
	} else {
		byKey = map[string]*Stat{key: s.byKey[key]}
		byModel = s.byKeyModel[key]
		bySrc = s.byKeySrc[key]
	}
	for _, a := range byModel {
		total.Reqs += a.Reqs
		total.OK += a.OK
		total.Err += a.Err
		total.Tokens += a.Tokens
		total.Prompt += a.Prompt
		total.Compl += a.Compl
		total.LatSum += a.LatSum
		if a.LatMax > total.LatMax {
			total.LatMax = a.LatMax
		}
	}
	return map[string]interface{}{
		"active":    s.active,
		"total":     total,
		"by_key":    rows(byKey),
		"by_model":  rows(byModel),
		"by_source": rows(bySrc),
		"records":   append([]Req(nil), recs...),
	}
}