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"}, }