diff --git a/internal/router/bandit/bandit.go b/internal/router/bandit/bandit.go index a2ee6188f..75e88cec5 100644 --- a/internal/router/bandit/bandit.go +++ b/internal/router/bandit/bandit.go @@ -109,8 +109,13 @@ func (b *Router) Route(ctx context.Context, req router.Request) (router.Decision argmaxModel := dec.Model pick := candidates[0] b.annotate(&dec, pick.model, pick.provider, b.propensity(clusterIDs, candidates, pick.model), pick.sample) - if pick.model != argmaxModel && md.EffectiveKnobsHash != 0 { - md.EffectiveKnobsHash = mixModel(md.EffectiveKnobsHash, pick.model) + if pick.model != argmaxModel { + if md.EffectiveKnobsHash != 0 { + md.EffectiveKnobsHash = mixModel(md.EffectiveKnobsHash, pick.model) + } + // The scorer's runner-up was computed against the argmax and may now + // equal the served peer; recompute it against the served model. + repairBandPair(&md.PairedModel, &md.PairedProvider, &md.PairedScore, md.CandidateScores, candidates, pick.model) } return dec, nil } @@ -185,6 +190,40 @@ func (b *Router) annotate(dec *router.Decision, model, provider string, propensi } } +// repairBandPair recomputes the band pair's runner-up against the served +// model. Picks the highest-scoring servable peer other than served (ties +// broken by name); clears the pair if none remains. Providers come from the +// live candidate list so peers resolved only via Decision.Provider stay eligible. +func repairBandPair(outModel, outProvider *string, outScore *float32, scores map[string]float32, cands []candidate, served string) { + providers := make(map[string]string, len(cands)) + for _, c := range cands { + providers[c.model] = c.provider + } + models := make([]string, 0, len(scores)) + for m := range scores { + models = append(models, m) + } + sort.Strings(models) + + bestModel, bestProvider := "", "" + var bestScore float32 + for _, m := range models { + if m == served { + continue + } + sc := scores[m] + if bestModel != "" && sc <= bestScore { + continue + } + provider, ok := providers[m] + if !ok || provider == "" { + continue + } + bestModel, bestProvider, bestScore = m, provider, sc + } + *outModel, *outProvider, *outScore = bestModel, bestProvider, bestScore +} + func mixModel(h uint64, model string) uint64 { f := fnv.New64a() var b [8]byte diff --git a/internal/router/bandit/bandit_test.go b/internal/router/bandit/bandit_test.go index 4eeda5b79..03b82086e 100644 --- a/internal/router/bandit/bandit_test.go +++ b/internal/router/bandit/bandit_test.go @@ -131,8 +131,12 @@ func TestRoute_PropensityReflectsPosteriorOverlap(t *testing.T) { } func TestRoute_KeepsArgmaxWhenSameModel(t *testing.T) { + // Keeping the argmax must leave the scorer's PairedModel untouched. scores := map[string]float32{"claude-haiku-4-5": 0.9, "claude-sonnet-4-6": 0.5} inner := clusterDecision(scores, []int{0}, "claude-haiku-4-5", "anthropic") + inner.Metadata.PairedModel = "claude-sonnet-4-6" + inner.Metadata.PairedProvider = "anthropic" + inner.Metadata.PairedScore = 0.5 post, err := LoadPosterior("testdata/ts_posterior.json") if err != nil { t.Fatal(err) @@ -149,6 +153,53 @@ func TestRoute_KeepsArgmaxWhenSameModel(t *testing.T) { if dec.Metadata.EffectiveKnobsHash != 99 { t.Fatal("keeping argmax must preserve cache key") } + if dec.Metadata.PairedModel != "claude-sonnet-4-6" { + t.Fatalf("argmax keep must preserve scorer runner-up sonnet, got %q", dec.Metadata.PairedModel) + } + if dec.Metadata.PairedScore != 0.5 { + t.Fatalf("argmax keep must preserve scorer paired score, got %v", dec.Metadata.PairedScore) + } +} + +func TestRoute_RepairsBandPairWhenServingRunnerUp(t *testing.T) { + // Serving the former runner-up must recompute PairedModel to the prior argmax. + scores := map[string]float32{ + "claude-sonnet-4-6": 0.90, + "claude-haiku-4-5": 0.85, + } + inner := clusterDecision(scores, []int{0}, "claude-sonnet-4-6", "anthropic") + inner.Metadata.PairedModel = "claude-haiku-4-5" + inner.Metadata.PairedProvider = "anthropic" + inner.Metadata.PairedScore = 0.85 + post := &Posterior{cells: map[int]map[string]Arm{ + 0: { + "claude-haiku-4-5": {Mean: 0.9, Variance: 0}, + "claude-sonnet-4-6": {Mean: 0.1, Variance: 0}, + }, + }} + b := New(&fakeInner{dec: inner}, post) + b.norm = func() float64 { return 0 } + b.trials = 0 + + dec, err := b.Route(context.Background(), router.Request{}) + if err != nil { + t.Fatal(err) + } + if dec.Model != "claude-haiku-4-5" { + t.Fatalf("expected Thompson pick haiku (runner-up), got %q", dec.Model) + } + if dec.Metadata.PairedModel == dec.Model { + t.Fatalf("band pair collapsed: Model and PairedModel are both %q", dec.Model) + } + if dec.Metadata.PairedModel != "claude-sonnet-4-6" { + t.Fatalf("expected runner-up recomputed to sonnet, got %q", dec.Metadata.PairedModel) + } + if dec.Metadata.PairedProvider != "anthropic" { + t.Fatalf("expected paired provider anthropic, got %q", dec.Metadata.PairedProvider) + } + if dec.Metadata.PairedScore != 0.90 { + t.Fatalf("expected paired score 0.90, got %v", dec.Metadata.PairedScore) + } } func TestRoute_UsesPerRequestProviderBinding(t *testing.T) {