diff options
Diffstat (limited to 'internal/stats')
| -rw-r--r-- | internal/stats/stats.go | 25 | ||||
| -rw-r--r-- | internal/stats/stats_test.go | 47 |
2 files changed, 72 insertions, 0 deletions
diff --git a/internal/stats/stats.go b/internal/stats/stats.go index 939e1aa..742a5be 100644 --- a/internal/stats/stats.go +++ b/internal/stats/stats.go @@ -75,6 +75,31 @@ type Snapshot struct { Window time.Duration } +// ScopeReqs returns the request count for a specific provider+model pair. +// Returns 0 when the provider or model is not present in the snapshot. +func (s Snapshot) ScopeReqs(provider, model string) int64 { + if pe, ok := s.Providers[provider]; ok { + if mc, ok2 := pe.Models[model]; ok2 { + return mc.Reqs + } + } + return 0 +} + +// ScopeRPM returns the requests-per-minute for a specific provider+model +// pair, derived from ScopeReqs and the snapshot's sliding window. +func (s Snapshot) ScopeRPM(provider, model string) float64 { + reqs := s.ScopeReqs(provider, model) + if reqs == 0 { + return 0 + } + mins := s.Window.Minutes() + if mins <= 0 { + mins = 0.001 + } + return float64(reqs) / mins +} + // Update appends one event and prunes old entries under lock. func Update(ctx context.Context, provider, model string, sentBytes, recvBytes int) error { dir, err := CacheDir() diff --git a/internal/stats/stats_test.go b/internal/stats/stats_test.go index 45f9e2a..a9b3d22 100644 --- a/internal/stats/stats_test.go +++ b/internal/stats/stats_test.go @@ -334,3 +334,50 @@ func TestUpdate_CancelledContext(t *testing.T) { t.Fatal("expected error from cancelled context, got nil") } } + +func TestSnapshot_ScopeReqs(t *testing.T) { + snap := Snapshot{ + Providers: map[string]ProviderEntry{ + "openai": {Models: map[string]Counters{"gpt-5.0": {Reqs: 42}}}, + }, + } + if got := snap.ScopeReqs("openai", "gpt-5.0"); got != 42 { + t.Fatalf("expected 42, got %d", got) + } + if got := snap.ScopeReqs("openai", "gpt-4.1"); got != 0 { + t.Fatalf("expected 0 for missing model, got %d", got) + } + if got := snap.ScopeReqs("anthropic", "gpt-5.0"); got != 0 { + t.Fatalf("expected 0 for missing provider, got %d", got) + } +} + +func TestSnapshot_ScopeRPM(t *testing.T) { + snap := Snapshot{ + Providers: map[string]ProviderEntry{ + "openai": {Models: map[string]Counters{"gpt-5.0": {Reqs: 60}}}, + }, + Window: time.Hour, + } + rpm := snap.ScopeRPM("openai", "gpt-5.0") + if rpm != 1.0 { + t.Fatalf("expected 1.0 rpm, got %v", rpm) + } + // Missing model should return 0 + if rpm := snap.ScopeRPM("openai", "missing"); rpm != 0 { + t.Fatalf("expected 0 rpm for missing, got %v", rpm) + } +} + +func TestSnapshot_ScopeRPM_ZeroWindow(t *testing.T) { + snap := Snapshot{ + Providers: map[string]ProviderEntry{ + "openai": {Models: map[string]Counters{"gpt-5.0": {Reqs: 10}}}, + }, + Window: 0, // edge case + } + rpm := snap.ScopeRPM("openai", "gpt-5.0") + if rpm <= 0 { + t.Fatalf("expected positive rpm even with zero window, got %v", rpm) + } +} |
