mirror of
https://github.com/coder/coder.git
synced 2026-10-31 02:43:21 +08:00
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. --- <sub>Opened by Coder Agents on behalf of @evgeniy-scherbina.</sub>
This commit is contained in:
@@ -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
|
||||
}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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"},
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user