From 211f49091058f7dfa0f9350dc5b6f9949d443753 Mon Sep 17 00:00:00 2001 From: Yevhenii Shcherbina Date: Thu, 20 Aug 2026 14:04:23 -0400 Subject: [PATCH] refactor: share the AI Gateway price book contract (#28310) Follow-up to the review discussion on the price book refresh automation: https://github.com/coder/coder/pull/28146#discussion_r3806439966 ## Problem The price seed row struct was duplicated three times, so a JSON tag change had to be repeated in each copy: - `scripts/aibridgepricesgen.priceRow` - `scripts/aibridgepricesdiff.priceRow` - `coderd/aibridge/prices.seedRow` The `(provider, model)` key and its parsing were duplicated too. ## Change Add `coderd/aibridge/prices/pricebook` holding the shared contract: `Row`, `Key`, `Row.Key()`, and `Parse`. The package depends only on the standard library, so the scripts can import it without pulling in the database layer that `coderd/aibridge/prices` brings. Consumer-specific behaviour stays where it is: upstream conversion and validation in the generator, comparison and Markdown rendering in the diff tool, DB seeding and runtime lookup in the seeder. The batch SQL upsert still extracts JSON fields by name, so it remains a separate contract boundary. The note about keeping it in sync now lives on the shared `Row`. No generated artifacts change; the serialized schema is identical. --- Opened by Coder Agents on behalf of @evgeniy-scherbina. --- coderd/aibridge/prices/pricebook/pricebook.go | 71 ++++++ .../prices/pricebook/pricebook_test.go | 235 ++++++++++++++++++ coderd/aibridge/prices/prices.go | 28 +-- coderd/aibridge/prices/prices_test.go | 35 +++ scripts/aibridgepricesdiff/main.go | 91 ++----- scripts/aibridgepricesdiff/main_test.go | 66 ++--- scripts/aibridgepricesgen/main.go | 38 +-- scripts/aibridgepricesgen/main_test.go | 6 +- 8 files changed, 413 insertions(+), 157 deletions(-) create mode 100644 coderd/aibridge/prices/pricebook/pricebook.go create mode 100644 coderd/aibridge/prices/pricebook/pricebook_test.go diff --git a/coderd/aibridge/prices/pricebook/pricebook.go b/coderd/aibridge/prices/pricebook/pricebook.go new file mode 100644 index 00000000000..c4b4283f988 --- /dev/null +++ b/coderd/aibridge/prices/pricebook/pricebook.go @@ -0,0 +1,71 @@ +// Package pricebook defines the serialized AI Gateway model price book. +package pricebook + +import ( + "encoding/json" + "io" +) + +// Row is one provider and model entry in the generated price book. Pointer +// fields preserve the distinction between a missing price and an explicit zero. +// +// NOTE: the batch SQL upsert extracts these fields from the raw JSON by name, +// so the tags here and that query must stay in sync. TestSeedFromBytes/ +// RowTagsMatchSQLColumns in the parent package guards that. +type Row struct { + Provider string `json:"provider"` + Model string `json:"model"` + InputPrice *int64 `json:"input_price"` + OutputPrice *int64 `json:"output_price"` + CacheReadPrice *int64 `json:"cache_read_price"` + CacheWritePrice *int64 `json:"cache_write_price"` +} + +// Key identifies a model within a provider. +type Key struct { + Provider string + Model string +} + +// Key returns the provider and model identity of the row. +func (r Row) Key() Key { + return Key{Provider: r.Provider, Model: r.Model} +} + +func (k Key) String() string { + return k.Provider + "/" + k.Model +} + +// Parse decodes a generated price book. +func Parse(data []byte) ([]Row, error) { + var rows []Row + if err := json.Unmarshal(data, &rows); err != nil { + return nil, err + } + return rows, nil +} + +// Write encodes a price book in the on-disk form of the generated artifact. +func Write(w io.Writer, rows []Row) error { + enc := json.NewEncoder(w) + enc.SetIndent("", " ") + return enc.Encode(rows) +} + +// SamePrices reports whether two rows carry identical prices. Provider and +// model are ignored, so it is only meaningful for rows with the same Key. A +// price moving to or from null counts as a difference, matching the null +// handling in the batch SQL upsert. +func (r Row) SamePrices(other Row) bool { + return equalPrice(r.InputPrice, other.InputPrice) && + equalPrice(r.OutputPrice, other.OutputPrice) && + equalPrice(r.CacheReadPrice, other.CacheReadPrice) && + equalPrice(r.CacheWritePrice, other.CacheWritePrice) +} + +func equalPrice(a, b *int64) bool { + if a == nil || b == nil { + return a == nil && b == nil + } + return *a == *b +} diff --git a/coderd/aibridge/prices/pricebook/pricebook_test.go b/coderd/aibridge/prices/pricebook/pricebook_test.go new file mode 100644 index 00000000000..2362846aed2 --- /dev/null +++ b/coderd/aibridge/prices/pricebook/pricebook_test.go @@ -0,0 +1,235 @@ +package pricebook_test + +import ( + "bytes" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/coder/coder/v2/coderd/aibridge/prices/pricebook" +) + +func TestParse(t *testing.T) { + t.Parallel() + + t.Run("DistinguishesNullFromZero", func(t *testing.T) { + t.Parallel() + + rows, err := pricebook.Parse([]byte(`[ + { + "provider": "openai", + "model": "gpt", + "input_price": 0, + "output_price": null, + "cache_read_price": 1, + "cache_write_price": 2 + } + ]`)) + require.NoError(t, err) + require.Len(t, rows, 1) + require.Equal(t, "openai", rows[0].Provider) + require.Equal(t, "gpt", rows[0].Model) + require.Equal(t, int64Ptr(0), rows[0].InputPrice) + require.Nil(t, rows[0].OutputPrice) + require.Equal(t, int64Ptr(1), rows[0].CacheReadPrice) + require.Equal(t, int64Ptr(2), rows[0].CacheWritePrice) + }) + + t.Run("MissingPricesAreNull", func(t *testing.T) { + t.Parallel() + + rows, err := pricebook.Parse([]byte(`[{"provider": "anthropic", "model": "claude"}]`)) + require.NoError(t, err) + require.Len(t, rows, 1) + require.Nil(t, rows[0].InputPrice) + require.Nil(t, rows[0].OutputPrice) + require.Nil(t, rows[0].CacheReadPrice) + require.Nil(t, rows[0].CacheWritePrice) + }) + + t.Run("Empty", func(t *testing.T) { + t.Parallel() + + rows, err := pricebook.Parse([]byte(`[]`)) + require.NoError(t, err) + require.Empty(t, rows) + }) + + t.Run("Malformed", func(t *testing.T) { + t.Parallel() + + _, err := pricebook.Parse([]byte(`not json`)) + require.Error(t, err) + }) +} + +func TestRowKey(t *testing.T) { + t.Parallel() + + // Provider and model stay separate fields, so a provider identifier that + // itself contains a slash cannot collide with a namespaced model. + a := pricebook.Row{ + Provider: "azure", + Model: "openai/gpt-5", + } + b := pricebook.Row{ + Provider: "azure/openai", + Model: "gpt-5", + } + require.Equal(t, pricebook.Key{Provider: "azure", Model: "openai/gpt-5"}, a.Key()) + require.NotEqual(t, a.Key(), b.Key()) +} + +func int64Ptr(v int64) *int64 { + return &v +} + +func TestWrite(t *testing.T) { + t.Parallel() + + t.Run("OnDiskForm", func(t *testing.T) { + t.Parallel() + + var buf bytes.Buffer + require.NoError(t, pricebook.Write(&buf, []pricebook.Row{{ + Provider: "openai", + Model: "gpt", + InputPrice: int64Ptr(0), + }})) + require.Equal(t, `[ + { + "provider": "openai", + "model": "gpt", + "input_price": 0, + "output_price": null, + "cache_read_price": null, + "cache_write_price": null + } +] +`, buf.String()) + }) + + t.Run("RoundTrips", func(t *testing.T) { + t.Parallel() + + // Zero and null must survive a write followed by a read, since the + // generated artifact is the input to both the diff tool and the seeder. + rows := []pricebook.Row{ + { + Provider: "anthropic", + Model: "claude", + InputPrice: int64Ptr(3_000_000), + OutputPrice: int64Ptr(0), + CacheReadPrice: nil, + CacheWritePrice: int64Ptr(1), + }, + { + Provider: "openai", + Model: "gpt", + }, + } + + var buf bytes.Buffer + require.NoError(t, pricebook.Write(&buf, rows)) + got, err := pricebook.Parse(buf.Bytes()) + require.NoError(t, err) + require.Equal(t, rows, got) + }) +} + +func TestRowSamePrices(t *testing.T) { + t.Parallel() + + cases := []struct { + name string + a pricebook.Row + b pricebook.Row + want bool + }{ + { + name: "identical", + a: pricebook.Row{ + InputPrice: int64Ptr(1), + OutputPrice: int64Ptr(2), + }, + b: pricebook.Row{ + InputPrice: int64Ptr(1), + OutputPrice: int64Ptr(2), + }, + want: true, + }, + { + name: "all null", + a: pricebook.Row{}, + b: pricebook.Row{}, + want: true, + }, + { + // Identity is not part of the comparison. + name: "different models same prices", + a: pricebook.Row{ + Provider: "openai", + Model: "gpt", + InputPrice: int64Ptr(1), + }, + b: pricebook.Row{ + Provider: "anthropic", + Model: "claude", + InputPrice: int64Ptr(1), + }, + want: true, + }, + { + name: "input differs", + a: pricebook.Row{InputPrice: int64Ptr(1)}, + b: pricebook.Row{InputPrice: int64Ptr(2)}, + want: false, + }, + { + name: "output differs", + a: pricebook.Row{OutputPrice: int64Ptr(1)}, + b: pricebook.Row{OutputPrice: int64Ptr(2)}, + want: false, + }, + { + name: "cache read differs", + a: pricebook.Row{CacheReadPrice: int64Ptr(1)}, + b: pricebook.Row{CacheReadPrice: int64Ptr(2)}, + want: false, + }, + { + name: "cache write differs", + a: pricebook.Row{CacheWritePrice: int64Ptr(1)}, + b: pricebook.Row{CacheWritePrice: int64Ptr(2)}, + want: false, + }, + { + // Zero is a populated price, distinct from an absent one. + name: "null and zero", + a: pricebook.Row{}, + b: pricebook.Row{InputPrice: int64Ptr(0)}, + want: false, + }, + { + name: "zero and null", + a: pricebook.Row{InputPrice: int64Ptr(0)}, + b: pricebook.Row{}, + want: false, + }, + { + name: "null and nonzero", + a: pricebook.Row{InputPrice: nil}, + b: pricebook.Row{InputPrice: int64Ptr(5)}, + want: false, + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + require.Equal(t, tc.want, tc.a.SamePrices(tc.b)) + // The relation is symmetric. + require.Equal(t, tc.want, tc.b.SamePrices(tc.a)) + }) + } +} diff --git a/coderd/aibridge/prices/prices.go b/coderd/aibridge/prices/prices.go index 0eed74a0926..f662b64ddc0 100644 --- a/coderd/aibridge/prices/prices.go +++ b/coderd/aibridge/prices/prices.go @@ -5,32 +5,16 @@ package prices import ( "context" _ "embed" - "encoding/json" "golang.org/x/xerrors" + "github.com/coder/coder/v2/coderd/aibridge/prices/pricebook" "github.com/coder/coder/v2/coderd/database" ) //go:embed data/prices.json var seedJSON []byte -// Pointer fields preserve the distinction between "not populated by upstream" -// (null) and "explicitly zero" (0). Used only for Go-side type validation in -// parseSeed; the upsert reads the raw JSON bytes via the batch SQL query. -// -// NOTE: the JSON contract for the price seed lives in three places that must -// stay in sync: the corresponding struct in the price generator, the column -// extraction in the batch SQL upsert, and the tags here. -type seedRow struct { - Provider string `json:"provider"` - Model string `json:"model"` - InputPrice *int64 `json:"input_price"` - OutputPrice *int64 `json:"output_price"` - CacheReadPrice *int64 `json:"cache_read_price"` - CacheWritePrice *int64 `json:"cache_write_price"` -} - // Seed applies the embedded price seed to ai_model_prices table, replacing the // price columns of any existing (provider, model) row and inserting new ones. // Rows already in the table that no longer appear in the seed are left @@ -43,7 +27,7 @@ func Seed(ctx context.Context, db database.Store) error { // which applies the seed embedded in this binary; SeedFromBytes is exposed // for tests that need to inject a deterministic seed. func SeedFromBytes(ctx context.Context, db database.Store, data []byte) error { - rows, err := parseSeed(data) + rows, err := pricebook.Parse(data) if err != nil { return xerrors.Errorf("parse price seed: %w", err) } @@ -55,11 +39,3 @@ func SeedFromBytes(ctx context.Context, db database.Store, data []byte) error { Source: database.AIModelPriceSourceDefault, }) } - -func parseSeed(data []byte) ([]seedRow, error) { - var rows []seedRow - if err := json.Unmarshal(data, &rows); err != nil { - return nil, err - } - return rows, nil -} diff --git a/coderd/aibridge/prices/prices_test.go b/coderd/aibridge/prices/prices_test.go index f97ee2e081d..4e5d6294ac1 100644 --- a/coderd/aibridge/prices/prices_test.go +++ b/coderd/aibridge/prices/prices_test.go @@ -1,6 +1,7 @@ package prices_test import ( + "bytes" "fmt" "testing" @@ -9,11 +10,13 @@ import ( "cdr.dev/slog/v3/sloggers/slogtest" "github.com/coder/coder/v2/coderd/aibridge/prices" + "github.com/coder/coder/v2/coderd/aibridge/prices/pricebook" "github.com/coder/coder/v2/coderd/coderdtest" "github.com/coder/coder/v2/coderd/database" "github.com/coder/coder/v2/coderd/database/dbauthz" "github.com/coder/coder/v2/coderd/database/dbtestutil" "github.com/coder/coder/v2/coderd/rbac" + "github.com/coder/coder/v2/coderd/util/ptr" "github.com/coder/coder/v2/testutil" ) @@ -186,6 +189,38 @@ func TestSeedFromBytes(t *testing.T) { require.Equal(t, int64(67890), got.OutputPrice.Int64) }) + t.Run("RowTagsMatchSQLColumns", func(t *testing.T) { + t.Parallel() + ctx := testutil.Context(t, testutil.WaitShort) + db, _ := dbtestutil.NewDB(t) + + // The upsert extracts each field from the raw JSON by name, so the + // struct tags on pricebook.Row and the query have to agree. Building + // the seed from the struct rather than a JSON literal is what makes a + // renamed tag fail here: the query would look for a name the seed no + // longer carries and write NULL instead. Each price gets a distinct + // value so a query that reads the wrong field also fails. + var seed bytes.Buffer + require.NoError(t, pricebook.Write(&seed, []pricebook.Row{{ + Provider: "test-provider", + Model: "tag-contract", + InputPrice: ptr.Ref(int64(11)), + OutputPrice: ptr.Ref(int64(22)), + CacheReadPrice: ptr.Ref(int64(33)), + CacheWritePrice: ptr.Ref(int64(44)), + }})) + require.NoError(t, prices.SeedFromBytes(ctx, db, seed.Bytes())) + + got, err := db.GetAIModelPriceByProviderModel(ctx, database.GetAIModelPriceByProviderModelParams{ + Provider: "test-provider", Model: "tag-contract", + }) + require.NoError(t, err) + require.Equal(t, int64(11), got.InputPrice.Int64) + require.Equal(t, int64(22), got.OutputPrice.Int64) + require.Equal(t, int64(33), got.CacheReadPrice.Int64) + require.Equal(t, int64(44), got.CacheWritePrice.Int64) + }) + t.Run("LeavesCustomPricesUntouched", func(t *testing.T) { t.Parallel() ctx := testutil.Context(t, testutil.WaitShort) diff --git a/scripts/aibridgepricesdiff/main.go b/scripts/aibridgepricesdiff/main.go index c2167fb809d..1109dbf3164 100644 --- a/scripts/aibridgepricesdiff/main.go +++ b/scripts/aibridgepricesdiff/main.go @@ -13,7 +13,6 @@ package main import ( - "encoding/json" "flag" "fmt" "io" @@ -22,46 +21,22 @@ import ( "strings" "golang.org/x/xerrors" + + "github.com/coder/coder/v2/coderd/aibridge/prices/pricebook" ) -// priceRow mirrors the seed file schema written by aibridgepricesgen. Pointer -// fields preserve the distinction between "not populated by upstream" (null) -// and "explicitly zero" (0). -type priceRow struct { - Provider string `json:"provider"` - Model string `json:"model"` - InputPrice *int64 `json:"input_price"` - OutputPrice *int64 `json:"output_price"` - CacheReadPrice *int64 `json:"cache_read_price"` - CacheWritePrice *int64 `json:"cache_write_price"` -} - -// modelKey identifies a model across the two snapshots. -type modelKey struct { - provider string - model string -} - -func (r priceRow) key() modelKey { - return modelKey{provider: r.Provider, model: r.Model} -} - -func (k modelKey) String() string { - return k.provider + "/" + k.model -} - -func less(a, b modelKey) bool { - if a.provider != b.provider { - return a.provider < b.provider +func less(a, b pricebook.Key) bool { + if a.Provider != b.Provider { + return a.Provider < b.Provider } - return a.model < b.model + return a.Model < b.Model } // diff is the full comparison between two snapshots type diff struct { - added []modelKey - removed []modelKey - changed []modelKey + added []pricebook.Key + removed []pricebook.Key + changed []pricebook.Key } func (d diff) empty() bool { @@ -94,65 +69,49 @@ func run(oldPath, newPath string, w io.Writer) error { return err } -func readRows(path string) ([]priceRow, error) { +func readRows(path string) ([]pricebook.Row, error) { data, err := os.ReadFile(path) if err != nil { return nil, err } - var rows []priceRow - if err := json.Unmarshal(data, &rows); err != nil { + rows, err := pricebook.Parse(data) + if err != nil { return nil, xerrors.Errorf("parse: %w", err) } return rows, nil } // compare classifies every model as added, removed, or changed. -func compare(oldRows, newRows []priceRow) diff { - oldByKey := make(map[modelKey]priceRow, len(oldRows)) +func compare(oldRows, newRows []pricebook.Row) diff { + oldByKey := make(map[pricebook.Key]pricebook.Row, len(oldRows)) for _, r := range oldRows { - oldByKey[r.key()] = r + oldByKey[r.Key()] = r } var d diff - seen := make(map[modelKey]struct{}, len(newRows)) + seen := make(map[pricebook.Key]struct{}, len(newRows)) for _, r := range newRows { - seen[r.key()] = struct{}{} - prev, ok := oldByKey[r.key()] + seen[r.Key()] = struct{}{} + prev, ok := oldByKey[r.Key()] switch { case !ok: - d.added = append(d.added, r.key()) - case !samePrices(prev, r): - d.changed = append(d.changed, r.key()) + d.added = append(d.added, r.Key()) + case !prev.SamePrices(r): + d.changed = append(d.changed, r.Key()) } } for _, r := range oldRows { - if _, ok := seen[r.key()]; !ok { - d.removed = append(d.removed, r.key()) + if _, ok := seen[r.Key()]; !ok { + d.removed = append(d.removed, r.Key()) } } - for _, keys := range [][]modelKey{d.added, d.removed, d.changed} { + for _, keys := range [][]pricebook.Key{d.added, d.removed, d.changed} { sort.Slice(keys, func(i, j int) bool { return less(keys[i], keys[j]) }) } return d } -// samePrices reports whether two rows for the same model carry identical -// prices. A price moving to or from null counts as a change. -func samePrices(a, b priceRow) bool { - return equalPrice(a.InputPrice, b.InputPrice) && - equalPrice(a.OutputPrice, b.OutputPrice) && - equalPrice(a.CacheReadPrice, b.CacheReadPrice) && - equalPrice(a.CacheWritePrice, b.CacheWritePrice) -} - -func equalPrice(a, b *int64) bool { - if a == nil || b == nil { - return a == nil && b == nil - } - return *a == *b -} - // render writes the Markdown summary: counts, then the models in each // category. Empty categories are omitted. func render(d diff) string { @@ -173,7 +132,7 @@ func render(d diff) string { plural(len(d.changed), "model"), ) - renderList := func(heading string, keys []modelKey) { + renderList := func(heading string, keys []pricebook.Key) { if len(keys) == 0 { return } diff --git a/scripts/aibridgepricesdiff/main_test.go b/scripts/aibridgepricesdiff/main_test.go index 0b334ce131d..985733103c1 100644 --- a/scripts/aibridgepricesdiff/main_test.go +++ b/scripts/aibridgepricesdiff/main_test.go @@ -7,6 +7,8 @@ import ( "testing" "github.com/stretchr/testify/require" + + "github.com/coder/coder/v2/coderd/aibridge/prices/pricebook" ) func TestCompare(t *testing.T) { @@ -14,23 +16,23 @@ func TestCompare(t *testing.T) { cases := []struct { name string - old []priceRow - new []priceRow + old []pricebook.Row + new []pricebook.Row wantAdded []string wantRemoved []string wantChanged []string }{ { name: "identical", - old: []priceRow{row("anthropic", "claude", 1, 2)}, - new: []priceRow{row("anthropic", "claude", 1, 2)}, + old: []pricebook.Row{row("anthropic", "claude", 1, 2)}, + new: []pricebook.Row{row("anthropic", "claude", 1, 2)}, }, { name: "added", - old: []priceRow{ + old: []pricebook.Row{ row("anthropic", "claude", 1, 2), }, - new: []priceRow{ + new: []pricebook.Row{ row("anthropic", "claude", 1, 2), row("openai", "gpt", 3, 4), }, @@ -38,30 +40,30 @@ func TestCompare(t *testing.T) { }, { name: "removed", - old: []priceRow{ + old: []pricebook.Row{ row("anthropic", "claude", 1, 2), row("openai", "gpt", 3, 4), }, - new: []priceRow{ + new: []pricebook.Row{ row("anthropic", "claude", 1, 2), }, wantRemoved: []string{"openai/gpt"}, }, { name: "changed", - old: []priceRow{row("anthropic", "claude", 1, 2)}, - new: []priceRow{row("anthropic", "claude", 1, 5)}, + old: []pricebook.Row{row("anthropic", "claude", 1, 2)}, + new: []pricebook.Row{row("anthropic", "claude", 1, 5)}, wantChanged: []string{"anthropic/claude"}, }, { // Every price field participates, not just input and output. name: "cache price changed", - old: []priceRow{{ + old: []pricebook.Row{{ Provider: "anthropic", Model: "claude", CacheReadPrice: int64Ptr(1), }}, - new: []priceRow{{ + new: []pricebook.Row{{ Provider: "anthropic", Model: "claude", CacheReadPrice: int64Ptr(2), @@ -71,11 +73,11 @@ func TestCompare(t *testing.T) { { // A previously missing price becoming populated is a change. name: "price set", - old: []priceRow{{ + old: []pricebook.Row{{ Provider: "anthropic", Model: "claude", }}, - new: []priceRow{{ + new: []pricebook.Row{{ Provider: "anthropic", Model: "claude", InputPrice: int64Ptr(1), @@ -85,10 +87,10 @@ func TestCompare(t *testing.T) { { // A model whose price becomes null is a change, not a removal. name: "price unset", - old: []priceRow{ + old: []pricebook.Row{ row("anthropic", "claude", 1, 2), }, - new: []priceRow{{ + new: []pricebook.Row{{ Provider: "anthropic", Model: "claude", InputPrice: int64Ptr(1), @@ -98,11 +100,11 @@ func TestCompare(t *testing.T) { { // Zero is a populated price, distinct from an absent one. name: "null becomes zero", - old: []priceRow{{ + old: []pricebook.Row{{ Provider: "openai", Model: "gpt", }}, - new: []priceRow{{ + new: []pricebook.Row{{ Provider: "openai", Model: "gpt", InputPrice: int64Ptr(0), @@ -112,12 +114,12 @@ func TestCompare(t *testing.T) { { // Zero is a real price, distinct from an absent one. name: "zero is not null", - old: []priceRow{{ + old: []pricebook.Row{{ Provider: "openai", Model: "gpt", InputPrice: int64Ptr(0), }}, - new: []priceRow{{ + new: []pricebook.Row{{ Provider: "openai", Model: "gpt", }}, @@ -126,11 +128,11 @@ func TestCompare(t *testing.T) { { // Same model identifier under two providers must not collide. name: "same model different providers", - old: []priceRow{ + old: []pricebook.Row{ row("anthropic", "shared", 1, 2), row("openai", "shared", 1, 2), }, - new: []priceRow{ + new: []pricebook.Row{ row("anthropic", "shared", 1, 2), row("openai", "shared", 9, 2), }, @@ -138,11 +140,11 @@ func TestCompare(t *testing.T) { }, { name: "added removed and changed together", - old: []priceRow{ + old: []pricebook.Row{ row("anthropic", "old-model", 1, 2), row("openai", "gpt", 3, 4), }, - new: []priceRow{ + new: []pricebook.Row{ row("anthropic", "new-model", 5, 6), row("openai", "gpt", 3, 7), }, @@ -167,7 +169,7 @@ func TestCompare(t *testing.T) { func TestCompareSortsDeterministically(t *testing.T) { t.Parallel() - newRows := []priceRow{ + newRows := []pricebook.Row{ row("openai", "b", 1, 1), row("anthropic", "z", 1, 1), row("anthropic", "a", 1, 1), @@ -184,8 +186,8 @@ func TestRender(t *testing.T) { t.Parallel() out := render(compare( - []priceRow{row("anthropic", "claude", 1, 2)}, - []priceRow{row("anthropic", "claude", 1, 2)}, + []pricebook.Row{row("anthropic", "claude", 1, 2)}, + []pricebook.Row{row("anthropic", "claude", 1, 2)}, )) require.Equal(t, `## Price book changes @@ -197,11 +199,11 @@ No price changes. t.Parallel() out := render(compare( - []priceRow{ + []pricebook.Row{ row("anthropic", "gone", 1, 2), row("openai", "gpt", 1, 2), }, - []priceRow{ + []pricebook.Row{ row("anthropic", "fresh", 3, 4), row("openai", "gpt", 2, 3), }, @@ -285,8 +287,8 @@ func TestRun(t *testing.T) { }) } -func row(provider, model string, input, output int64) priceRow { - return priceRow{ +func row(provider, model string, input, output int64) pricebook.Row { + return pricebook.Row{ Provider: provider, Model: model, InputPrice: int64Ptr(input), @@ -294,7 +296,7 @@ func row(provider, model string, input, output int64) priceRow { } } -func names(keys []modelKey) []string { +func names(keys []pricebook.Key) []string { if len(keys) == 0 { return nil } diff --git a/scripts/aibridgepricesgen/main.go b/scripts/aibridgepricesgen/main.go index 4d41e3f1dfa..6fcdb5c51e8 100644 --- a/scripts/aibridgepricesgen/main.go +++ b/scripts/aibridgepricesgen/main.go @@ -17,13 +17,13 @@ import ( "encoding/json" "flag" "fmt" - "io" "math" "os" "sort" "golang.org/x/xerrors" + "github.com/coder/coder/v2/coderd/aibridge/prices/pricebook" "github.com/coder/coder/v2/coderd/aibridge/prices/providers" ) @@ -63,21 +63,6 @@ func (c *upstreamCost) hasPricing() bool { c.CacheRead != nil || c.CacheWrite != nil } -// Pointer fields preserve the distinction between "not populated by upstream" -// (null) and "explicitly zero" (0). -// -// NOTE: the JSON contract for the price seed lives in three places that must -// stay in sync: the tags here, the corresponding struct in the price seeder, -// and the column extraction in the batch SQL upsert. -type priceRow struct { - Provider string `json:"provider"` - Model string `json:"model"` - InputPrice *int64 `json:"input_price"` - OutputPrice *int64 `json:"output_price"` - CacheReadPrice *int64 `json:"cache_read_price"` - CacheWritePrice *int64 `json:"cache_write_price"` -} - func main() { format := flag.String("format", "prices", `output format: "prices" (cost-control seed) or "catalog" (frontend known-models JSON)`) upstreamPath := flag.String("upstream", "", "path to a models.dev api.json snapshot (required)") @@ -131,8 +116,8 @@ func runPrices(upstream map[string]upstreamProvider) error { if err := validate(rows); err != nil { return err } - if err := write(os.Stdout, rows); err != nil { - return err + if err := pricebook.Write(os.Stdout, rows); err != nil { + return xerrors.Errorf("encode: %w", err) } _, _ = fmt.Fprintf(os.Stderr, "aibridgepricesgen: wrote %d prices for %d provider(s)\n", len(rows), len(providers.Supported)) return nil @@ -158,9 +143,9 @@ func runCatalog(upstream map[string]upstreamProvider) error { // providers. If any configured provider is absent from the upstream payload, // every missing provider is reported and the function returns an error so the // caller doesn't ship an incomplete seed. -func convert(upstream map[string]upstreamProvider, providerIDs []string) ([]priceRow, error) { +func convert(upstream map[string]upstreamProvider, providerIDs []string) ([]pricebook.Row, error) { var ( - rows []priceRow + rows []pricebook.Row missing []string ) for _, providerID := range providerIDs { @@ -173,7 +158,7 @@ func convert(upstream map[string]upstreamProvider, providerIDs []string) ([]pric if !m.Cost.hasPricing() { continue } - rows = append(rows, priceRow{ + rows = append(rows, pricebook.Row{ Provider: providerID, Model: modelID, InputPrice: toMicros(m.Cost.Input), @@ -200,7 +185,7 @@ func convert(upstream map[string]upstreamProvider, providerIDs []string) ([]pric // changes that produce structurally valid but semantically broken seed // data, e.g. a renamed `cost` key that leaves every row with all-null // prices. -func validate(rows []priceRow) error { +func validate(rows []pricebook.Row) error { for _, r := range rows { if r.InputPrice != nil || r.OutputPrice != nil { return nil @@ -223,12 +208,3 @@ func toMicros(price *float64) *int64 { micros := int64(math.Round(*price * 1_000_000)) return µs } - -func write(w io.Writer, rows []priceRow) error { - enc := json.NewEncoder(w) - enc.SetIndent("", " ") - if err := enc.Encode(rows); err != nil { - return xerrors.Errorf("encode: %w", err) - } - return nil -} diff --git a/scripts/aibridgepricesgen/main_test.go b/scripts/aibridgepricesgen/main_test.go index b21793f0d62..a3b32c2b296 100644 --- a/scripts/aibridgepricesgen/main_test.go +++ b/scripts/aibridgepricesgen/main_test.go @@ -5,6 +5,8 @@ import ( "testing" "github.com/stretchr/testify/require" + + "github.com/coder/coder/v2/coderd/aibridge/prices/pricebook" ) func TestToMicros(t *testing.T) { @@ -136,7 +138,7 @@ func TestValidate(t *testing.T) { t.Run("PassesWhenAnyRowHasPricing", func(t *testing.T) { t.Parallel() - rows := []priceRow{ + rows := []pricebook.Row{ {Provider: "openai", Model: "no-prices"}, {Provider: "anthropic", Model: "claude", InputPrice: int64Ptr(3_000_000)}, } @@ -148,7 +150,7 @@ func TestValidate(t *testing.T) { // Mirrors what would happen if upstream renamed the `cost` key: // Go's decoder silently drops it, every row gets all-null prices, // and convert returns syntactically valid rows with no pricing. - rows := []priceRow{ + rows := []pricebook.Row{ {Provider: "anthropic", Model: "claude-x"}, {Provider: "openai", Model: "gpt-x"}, }