diff --git a/Makefile b/Makefile index 156ee62..c91e6b8 100644 --- a/Makefile +++ b/Makefile @@ -40,8 +40,12 @@ yaegi-check: # pruning rule and naming-bridge mechanism. Not run by CI (network # fetch, and the generated table is committed); an operator re-runs this # by hand when LiteLLM's own pricing drifts. +# PRICING_MAX_AGE_MONTHS drops entries older than N months (0 = keep all); +# PRICING_KEEP is a comma-separated list of bare ids exempt from that rule. +PRICING_MAX_AGE_MONTHS ?= 12 +PRICING_KEEP ?= pricing-sync: - cd tools/pricing-sync && GOWORK=off go run . $(CURDIR) + cd tools/pricing-sync && GOWORK=off go run . -max-age-months=$(PRICING_MAX_AGE_MONTHS) -keep=$(PRICING_KEEP) $(CURDIR) gofmt -w pricing_data_gen.go INTEGRATION_COMPOSE := integration/docker-compose.yml diff --git a/tools/pricing-sync/main.go b/tools/pricing-sync/main.go index 8cdc63f..8b4c56e 100644 --- a/tools/pricing-sync/main.go +++ b/tools/pricing-sync/main.go @@ -60,9 +60,11 @@ // 4. Freshness: an entry is dropped when LiteLLM's own deprecation_date // has already passed, or when a full calendar date embedded in its // bare id ("claude-opus-4-1-20250805", "gpt-4o-2024-11-20") is -// older than maxModelAgeMonths. LiteLLM carries no release date, so -// an id with no such date and no deprecation_date is always kept: -// an undated alias is never dropped on a guess. +// older than the -max-age-months flag (default 12; 0 turns the whole +// rule off). LiteLLM carries no release date, so an id with no such +// date and no deprecation_date is always kept: an undated alias is +// never dropped on a guess. An id in alwaysKeep, or in the -keep +// flag's comma-separated list, is exempt. // // # Bare-id collision resolution // @@ -101,6 +103,7 @@ import ( "crypto/sha256" "encoding/hex" "encoding/json" + "flag" "fmt" "go/format" "io" @@ -136,20 +139,37 @@ const microUSDPerUSDPerToken = 1e12 // this specific, known case. const sampleSpecKey = "sample_spec" -// maxModelAgeMonths is the freshness rule's age cutoff for a date embedded -// in a bare model id. -const maxModelAgeMonths = 12 +// defaultMaxAgeMonths is the -max-age-months flag's default. +const defaultMaxAgeMonths = 12 + +// alwaysKeep lists bare ids the freshness rule never drops, because they +// are still served after their age or deprecation date says otherwise. +var alwaysKeep = map[string]bool{ + "claude-sonnet-4-5-20250929": true, +} + +// pruneOptions is the freshness rule's configuration. maxAgeMonths <= 0 +// turns the rule off. +type pruneOptions struct { + now time.Time + keep map[string]bool + maxAgeMonths int +} // idDatePattern finds a full calendar date in a model id, as YYYYMMDD or // YYYY-MM-DD, delimited by the id's ends or a "-" or "_". var idDatePattern = regexp.MustCompile(`(?:^|[-_])(20\d{2})-?(0[1-9]|1[0-2])-?(0[1-9]|[12]\d|3[01])(?:$|[-_])`) // isStale reports whether an entry fails the freshness rule: its -// deprecation_date (YYYY-MM-DD) is before now, or a date embedded in -// bareID is more than maxModelAgeMonths before now. -func isStale(bareID, deprecationDate string, now time.Time) bool { +// deprecation_date (YYYY-MM-DD) is before opts.now, or a date embedded in +// bareID is more than opts.maxAgeMonths before it. Ids in opts.keep and +// alwaysKeep never fail. +func isStale(bareID, deprecationDate string, opts pruneOptions) bool { + if opts.maxAgeMonths <= 0 || opts.keep[bareID] || alwaysKeep[bareID] { + return false + } if deprecationDate != "" { - if d, err := time.Parse("2006-01-02", deprecationDate); err == nil && d.Before(now) { + if d, err := time.Parse("2006-01-02", deprecationDate); err == nil && d.Before(opts.now) { return true } } @@ -161,7 +181,7 @@ func isStale(bareID, deprecationDate string, now time.Time) bool { if err != nil { return false } - return d.Before(now.AddDate(0, -maxModelAgeMonths, 0)) + return d.Before(opts.now.AddDate(0, -opts.maxAgeMonths, 0)) } // allowedProviders is the pruning rule's provider allowlist — see the @@ -290,9 +310,18 @@ func main() { } func run() error { + maxAge := flag.Int("max-age-months", defaultMaxAgeMonths, "drop entries older than this many months; 0 disables the freshness rule") + keepList := flag.String("keep", "", "comma-separated bare ids the freshness rule never drops") + flag.Parse() repoRoot := "." - if len(os.Args) > 1 { - repoRoot = os.Args[1] + if flag.NArg() > 0 { + repoRoot = flag.Arg(0) + } + keep := map[string]bool{} + for _, id := range strings.Split(*keepList, ",") { + if id = strings.TrimSpace(id); id != "" { + keep[id] = true + } } outPath := filepath.Join(repoRoot, "pricing_data_gen.go") @@ -304,7 +333,7 @@ func run() error { return fmt.Errorf("fetch %s: %w", litellmSourceURL, err) } - table, totalUpstream, err := parseAndPrune(data, time.Now()) + table, totalUpstream, err := parseAndPrune(data, pruneOptions{now: time.Now(), maxAgeMonths: *maxAge, keep: keep}) if err != nil { return fmt.Errorf("parse and prune: %w", err) } @@ -355,7 +384,7 @@ func fetchLiteLLMPrices(ctx context.Context, url string) ([]byte, error) { // chooseWinner's deterministic bare-id collision resolution. It returns // the pruned table and the total number of top-level keys data carried // (including ones this function skips), for the summary line run prints. -func parseAndPrune(data []byte, now time.Time) (map[string]builtinEntry, int, error) { +func parseAndPrune(data []byte, opts pruneOptions) (map[string]builtinEntry, int, error) { var raw map[string]json.RawMessage if err := json.Unmarshal(data, &raw); err != nil { return nil, 0, fmt.Errorf("decode top-level object: %w", err) @@ -382,7 +411,7 @@ func parseAndPrune(data []byte, now time.Time) (map[string]builtinEntry, int, er } bare, isBareKey := bareModelID(key, e.LiteLLMProvider) - if isStale(bare, e.DeprecationDate, now) { + if isStale(bare, e.DeprecationDate, opts) { continue } byBare[bare] = append(byBare[bare], candidate{ diff --git a/tools/pricing-sync/main_test.go b/tools/pricing-sync/main_test.go index 5886f8a..9337f6e 100644 --- a/tools/pricing-sync/main_test.go +++ b/tools/pricing-sync/main_test.go @@ -9,7 +9,10 @@ import ( "time" ) -var fixtureNow = time.Date(2026, time.August, 1, 0, 0, 0, 0, time.UTC) +var ( + fixtureNow = time.Date(2026, time.August, 1, 0, 0, 0, 0, time.UTC) + fixtureOpts = pruneOptions{now: fixtureNow, maxAgeMonths: defaultMaxAgeMonths} +) // fixtureJSON is a small, hand-built stand-in for LiteLLM's real // model_prices_and_context_window.json, covering every branch of the @@ -119,7 +122,7 @@ const fixtureUpstreamKeyCount = 17 // out, covering every pruning-rule branch fixtureJSON's own doc comment // lists. func TestParseAndPrune_Golden(t *testing.T) { - got, totalUpstream, err := parseAndPrune([]byte(fixtureJSON), fixtureNow) + got, totalUpstream, err := parseAndPrune([]byte(fixtureJSON), fixtureOpts) if err != nil { t.Fatalf("parseAndPrune: %v", err) } @@ -289,7 +292,8 @@ func TestIsStale(t *testing.T) { } for _, tc := range tests { t.Run(tc.name, func(t *testing.T) { - if got := isStale(tc.id, tc.deprecated, now); got != tc.want { + opts := pruneOptions{now: now, maxAgeMonths: defaultMaxAgeMonths} + if got := isStale(tc.id, tc.deprecated, opts); got != tc.want { t.Errorf("isStale(%q, %q) = %v, want %v", tc.id, tc.deprecated, got, tc.want) } }) @@ -302,7 +306,7 @@ func TestParseAndPrune_DropsStaleEntries(t *testing.T) { "gpt-old-2024-05-13": {"litellm_provider": "openai", "mode": "chat", "input_cost_per_token": 1e-6, "output_cost_per_token": 2e-6}, "gpt-dead": {"litellm_provider": "openai", "mode": "chat", "deprecation_date": "2026-01-01", "input_cost_per_token": 1e-6, "output_cost_per_token": 2e-6} }` - got, _, err := parseAndPrune([]byte(data), fixtureNow) + got, _, err := parseAndPrune([]byte(data), fixtureOpts) if err != nil { t.Fatalf("parseAndPrune: %v", err) } @@ -313,3 +317,29 @@ func TestParseAndPrune_DropsStaleEntries(t *testing.T) { t.Errorf("gpt-fresh missing from %v", got) } } + +func TestIsStale_Options(t *testing.T) { + now := time.Date(2026, time.September, 29, 0, 0, 0, 0, time.UTC) + const old = "gpt-4o-2024-08-06" + tests := []struct { + name string + opts pruneOptions + id string + want bool + }{ + {"default drops old", pruneOptions{now: now, maxAgeMonths: 12}, old, true}, + {"zero disables rule", pruneOptions{now: now, maxAgeMonths: 0}, old, false}, + {"negative disables rule", pruneOptions{now: now, maxAgeMonths: -1}, old, false}, + {"longer window keeps", pruneOptions{now: now, maxAgeMonths: 36}, old, false}, + {"shorter window drops more", pruneOptions{now: now, maxAgeMonths: 3}, "model-2026-03-01", true}, + {"keep flag exempts", pruneOptions{now: now, maxAgeMonths: 12, keep: map[string]bool{old: true}}, old, false}, + {"alwaysKeep exempts", pruneOptions{now: now, maxAgeMonths: 12}, "claude-sonnet-4-5-20250929", false}, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + if got := isStale(tc.id, "", tc.opts); got != tc.want { + t.Errorf("isStale(%q) = %v, want %v", tc.id, got, tc.want) + } + }) + } +}