diff --git a/internal/modelinfo/litellm_decode_test.go b/internal/modelinfo/litellm_decode_test.go new file mode 100644 index 0000000..06a9725 --- /dev/null +++ b/internal/modelinfo/litellm_decode_test.go @@ -0,0 +1,99 @@ +package modelinfo + +import ( + "compress/gzip" + "context" + "io" + "net/http" + "net/http/httptest" + "os" + "testing" + "time" +) + +// dirtySample is the shape that broke production, minimised: a `sample_spec` +// entry whose numeric fields hold prose, plus a model that spells an integer +// window as a float. A whole-map typed decode fails on the first of these and +// discards every good row with it. +const dirtySample = `{ + "sample_spec": {"max_input_tokens": "max input tokens, if the provider specifies it", + "max_tokens": "LEGACY parameter", + "deprecation_date": "date when the model becomes deprecated in the format YYYY-MM-DD"}, + "float-window-model": {"max_input_tokens": 128000.0}, + "aws/claude-sonnet-5": {"max_input_tokens": 1000000}, + "prose-window-model": {"max_input_tokens": "quite a lot, actually"} +}` + +func TestLiteLLMSkipsMalformedEntries(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Write([]byte(dirtySample)) + })) + defer srv.Close() + l := NewLiteLLM(srv.URL, srv.Client(), time.Hour) + ctx := context.Background() + waitFor(t, l, "float-window-model") + + // A prose-filled sibling entry must not poison the whole map. + if w, ok := l.Window(ctx, "aws/claude-sonnet-5"); !ok || w != 1000000 { + t.Fatalf("window(aws/claude-sonnet-5) = %d,%v; want 1000000,true (a malformed sibling entry poisoned the decode)", w, ok) + } + // A float-spelled integer must still resolve. + if w, ok := l.Window(ctx, "float-window-model"); !ok || w != 128000 { + t.Errorf("float-spelled window = %d,%v; want 128000,true", w, ok) + } +} + +// TestLiteLLMDecodesTheRealDocument runs the decode against a checked-in snapshot +// of the actual upstream document (gzipped; see testdata). This is the regression +// test the package never had: before the per-entry decode it resolved ZERO +// entries from this exact byte sequence. +func TestLiteLLMDecodesTheRealDocument(t *testing.T) { + f, err := os.Open("testdata/litellm_prices.json.gz") + if err != nil { + t.Fatal(err) + } + defer f.Close() + zr, err := gzip.NewReader(f) + if err != nil { + t.Fatal(err) + } + body, err := io.ReadAll(zr) + if err != nil { + t.Fatal(err) + } + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Write(body) + })) + defer srv.Close() + + l := NewLiteLLM(srv.URL, srv.Client(), time.Hour) + ctx := context.Background() + waitFor(t, l, "gpt-4o") + + l.mu.Lock() + n := len(l.byKey) + l.mu.Unlock() + if n < 2900 { + t.Errorf("decoded %d keys from the real document; want >2900 (a strict whole-map decode yields 0)", n) + } + // The two models context-guru is actually deployed against must resolve a + // non-zero window, or every fraction-based trigger is dead. + for _, m := range []string{"aws/claude-sonnet-5", "aws/claude-haiku-4-5"} { + w, ok := l.Window(ctx, m) + if !ok || w == 0 { + t.Errorf("window(%s) = %d,%v; want a non-zero window", m, w, ok) + } + } + t.Logf("decoded %d window keys from the real LiteLLM document", n) +} + +func waitFor(t *testing.T, l *LiteLLM, model string) { + t.Helper() + for i := 0; i < 400; i++ { + if _, ok := l.Window(context.Background(), model); ok { + return + } + time.Sleep(5 * time.Millisecond) + } + t.Fatalf("litellm map never loaded (%s never resolved)", model) +} diff --git a/internal/modelinfo/modelinfo.go b/internal/modelinfo/modelinfo.go index 0c0ac74..d7c6079 100644 --- a/internal/modelinfo/modelinfo.go +++ b/internal/modelinfo/modelinfo.go @@ -15,6 +15,7 @@ import ( "context" "encoding/json" "io" + "log/slog" "net/http" "strings" "sync" @@ -120,18 +121,38 @@ func (l *LiteLLM) fetch(ctx context.Context) (map[string]int, error) { if err != nil { return nil, err } - var raw map[string]struct { - MaxInputTokens int `json:"max_input_tokens"` - MaxTokens int `json:"max_tokens"` - } + // Decode PER ENTRY, not into one typed map[string]struct{…}. The upstream + // document is community-maintained and not schema-clean: it carries a + // `sample_spec` documentation entry whose numeric fields hold prose + // ("deprecation_date": "date when the model becomes deprecated…"), and a + // handful of models spell an integer field as a float. encoding/json aborts the + // WHOLE map on the first type error, so a strict decode returned (nil, err) + // after successfully parsing ~2,900 good rows — which is why every context + // window resolved to "unknown" in production and no fraction-based trigger has + // ever fired. Skip the bad entries; keep the good ones; say so out loud. + var raw map[string]json.RawMessage if err := json.Unmarshal(b, &raw); err != nil { return nil, err } m := make(map[string]int, len(raw)*2) - for k, v := range raw { - w := v.MaxInputTokens + var skipped []string + for k, rv := range raw { + if k == sampleSpecKey { + continue // the document's own schema documentation, not a model + } + var v struct { + // float64, not int: a few entries spell an integer window as 128000.0, + // which an int field rejects. + MaxInputTokens float64 `json:"max_input_tokens"` + MaxTokens float64 `json:"max_tokens"` + } + if err := json.Unmarshal(rv, &v); err != nil { + skipped = append(skipped, k) + continue + } + w := int(v.MaxInputTokens) if w == 0 { - w = v.MaxTokens + w = int(v.MaxTokens) } if w == 0 { continue @@ -142,9 +163,24 @@ func (l *LiteLLM) fetch(ctx context.Context) (map[string]int, error) { m[tail] = w } } + // Degrading silently here is what hid this bug for the life of the package: an + // empty map is indistinguishable from "no model has a window" at every call + // site, because every lookup fails open. Both outcomes get a log line. + if len(m) == 0 { + slog.Warn("modelinfo: the model-window document decoded to nothing; every context window will read as unknown and fraction-based triggers will not fire", + "url", l.URL, "entries", len(raw), "skipped", len(skipped)) + } else if len(skipped) > 0 { + slog.Info("modelinfo: skipped malformed model entries", "skipped", len(skipped), + "kept", len(m), "examples", skipped[:min(3, len(skipped))]) + } return m, nil } +// sampleSpecKey is the LiteLLM document's self-documenting entry: its fields are +// prose descriptions of the schema, not values. Skipped by name so it never even +// counts as a decode failure. +const sampleSpecKey = "sample_spec" + // Window returns the model's context window from the cached LiteLLM map. func (l *LiteLLM) Window(ctx context.Context, model string) (int, bool) { l.refreshIfStale(ctx) diff --git a/internal/modelinfo/testdata/litellm_prices.json.gz b/internal/modelinfo/testdata/litellm_prices.json.gz new file mode 100644 index 0000000..dca6ffa Binary files /dev/null and b/internal/modelinfo/testdata/litellm_prices.json.gz differ