Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit d295caf

Browse files
ssncferreirassncferreira
authored
fix: only write AI model prices that changed (#27923) (#28105)
Backport of #27923 Original PR: #27923 - only write AI model prices that changed Merge commit: 5efa7ab Requested by: @ssncferreira
1 parent 2934aa0 commit d295caf

4 files changed

Lines changed: 137 additions & 3 deletions

File tree

‎coderd/aibridge/prices/prices_test.go‎

Lines changed: 105 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package prices_test
22

33
import (
4+
"fmt"
45
"testing"
56

67
"github.com/prometheus/client_golang/prometheus"
@@ -58,6 +59,7 @@ func TestSeedFromBytes(t *testing.T) {
5859
require.Equal(t, int64(25_000_000), opus.OutputPrice.Int64)
5960
require.Equal(t, int64(500_000), opus.CacheReadPrice.Int64)
6061
require.Equal(t, int64(6_250_000), opus.CacheWritePrice.Int64)
62+
require.Equal(t, opus.CreatedAt, opus.UpdatedAt)
6163

6264
// Spot-check a row where the seed has a NULL price (OpenAI does not
6365
// publish a cache_write_price). The column should land as SQL NULL.
@@ -90,11 +92,11 @@ func TestSeedFromBytes(t *testing.T) {
9092
})
9193
require.NoError(t, err)
9294

93-
// Prices must be identical across runs and CreatedAt must be
94-
// preserved (only updated_at moves on a no-op upsert).
95+
// A re-seed that changes nothing must not touch the row at all.
9596
require.Equal(t, first.InputPrice, second.InputPrice)
9697
require.Equal(t, first.OutputPrice, second.OutputPrice)
9798
require.Equal(t, first.CreatedAt, second.CreatedAt)
99+
require.Equal(t, first.UpdatedAt, second.UpdatedAt)
98100
})
99101

100102
t.Run("OverwritesExistingPrices", func(t *testing.T) {
@@ -114,6 +116,10 @@ func TestSeedFromBytes(t *testing.T) {
114116
"cache_read_price": 3,
115117
"cache_write_price": 4
116118
}]`)))
119+
before, err := db.GetAIModelPriceByProviderModel(ctx, database.GetAIModelPriceByProviderModelParams{
120+
Provider: "openai", Model: "gpt-4o",
121+
})
122+
require.NoError(t, err)
117123

118124
require.NoError(t, prices.SeedFromBytes(ctx, db, []byte(testSeedJSON)))
119125

@@ -126,6 +132,8 @@ func TestSeedFromBytes(t *testing.T) {
126132
require.Equal(t, int64(1_250_000), got.CacheReadPrice.Int64)
127133
require.False(t, got.CacheWritePrice.Valid)
128134
require.Zero(t, got.CacheWritePrice.Int64)
135+
require.Equal(t, before.CreatedAt, got.CreatedAt)
136+
require.True(t, got.UpdatedAt.After(before.UpdatedAt))
129137
})
130138

131139
t.Run("LeavesOrphanRowsUntouched", func(t *testing.T) {
@@ -174,6 +182,101 @@ func TestSeedFromBytes(t *testing.T) {
174182
require.True(t, got.InputPrice.Valid)
175183
require.Equal(t, int64(2_500_000), got.InputPrice.Int64)
176184
})
185+
186+
// Every price column counts toward the comparison, and a NULL on either
187+
// side counts as a difference.
188+
t.Run("UpdatedAtTracksPriceChanges", func(t *testing.T) {
189+
t.Parallel()
190+
191+
key := database.GetAIModelPriceByProviderModelParams{Provider: "openai", Model: "gpt-4o"}
192+
seed := func(priceFields string) []byte {
193+
return fmt.Appendf(nil, `[{"provider": %q, "model": %q, %s}]`, key.Provider, key.Model, priceFields)
194+
}
195+
196+
tests := []struct {
197+
name string
198+
initial, updated string
199+
}{
200+
{
201+
name: "InputPriceChanged",
202+
initial: `"input_price": 100, "output_price": 200, "cache_read_price": 300, "cache_write_price": 400`,
203+
updated: `"input_price": 111, "output_price": 200, "cache_read_price": 300, "cache_write_price": 400`,
204+
},
205+
{
206+
name: "InputPriceSetFromNull",
207+
initial: `"input_price": null, "output_price": 200, "cache_read_price": 300, "cache_write_price": 400`,
208+
updated: `"input_price": 100, "output_price": 200, "cache_read_price": 300, "cache_write_price": 400`,
209+
},
210+
{
211+
name: "InputPriceClearedToNull",
212+
initial: `"input_price": 100, "output_price": 200, "cache_read_price": 300, "cache_write_price": 400`,
213+
updated: `"input_price": null, "output_price": 200, "cache_read_price": 300, "cache_write_price": 400`,
214+
},
215+
{
216+
name: "OutputPriceChanged",
217+
initial: `"input_price": 100, "output_price": 200, "cache_read_price": 300, "cache_write_price": 400`,
218+
updated: `"input_price": 100, "output_price": 222, "cache_read_price": 300, "cache_write_price": 400`,
219+
},
220+
{
221+
name: "OutputPriceSetFromNull",
222+
initial: `"input_price": 100, "output_price": null, "cache_read_price": 300, "cache_write_price": 400`,
223+
updated: `"input_price": 100, "output_price": 200, "cache_read_price": 300, "cache_write_price": 400`,
224+
},
225+
{
226+
name: "OutputPriceClearedToNull",
227+
initial: `"input_price": 100, "output_price": 200, "cache_read_price": 300, "cache_write_price": 400`,
228+
updated: `"input_price": 100, "output_price": null, "cache_read_price": 300, "cache_write_price": 400`,
229+
},
230+
{
231+
name: "CacheReadPriceChanged",
232+
initial: `"input_price": 100, "output_price": 200, "cache_read_price": 300, "cache_write_price": 400`,
233+
updated: `"input_price": 100, "output_price": 200, "cache_read_price": 333, "cache_write_price": 400`,
234+
},
235+
{
236+
name: "CacheReadPriceSetFromNull",
237+
initial: `"input_price": 100, "output_price": 200, "cache_read_price": null, "cache_write_price": 400`,
238+
updated: `"input_price": 100, "output_price": 200, "cache_read_price": 300, "cache_write_price": 400`,
239+
},
240+
{
241+
name: "CacheReadPriceClearedToNull",
242+
initial: `"input_price": 100, "output_price": 200, "cache_read_price": 300, "cache_write_price": 400`,
243+
updated: `"input_price": 100, "output_price": 200, "cache_read_price": null, "cache_write_price": 400`,
244+
},
245+
{
246+
name: "CacheWritePriceChanged",
247+
initial: `"input_price": 100, "output_price": 200, "cache_read_price": 300, "cache_write_price": 400`,
248+
updated: `"input_price": 100, "output_price": 200, "cache_read_price": 300, "cache_write_price": 444`,
249+
},
250+
{
251+
name: "CacheWritePriceSetFromNull",
252+
initial: `"input_price": 100, "output_price": 200, "cache_read_price": 300, "cache_write_price": null`,
253+
updated: `"input_price": 100, "output_price": 200, "cache_read_price": 300, "cache_write_price": 400`,
254+
},
255+
{
256+
name: "CacheWritePriceClearedToNull",
257+
initial: `"input_price": 100, "output_price": 200, "cache_read_price": 300, "cache_write_price": 400`,
258+
updated: `"input_price": 100, "output_price": 200, "cache_read_price": 300, "cache_write_price": null`,
259+
},
260+
}
261+
262+
for _, tt := range tests {
263+
t.Run(tt.name, func(t *testing.T) {
264+
t.Parallel()
265+
ctx := testutil.Context(t, testutil.WaitShort)
266+
db, _ := dbtestutil.NewDB(t)
267+
268+
require.NoError(t, prices.SeedFromBytes(ctx, db, seed(tt.initial)))
269+
before, err := db.GetAIModelPriceByProviderModel(ctx, key)
270+
require.NoError(t, err)
271+
272+
require.NoError(t, prices.SeedFromBytes(ctx, db, seed(tt.updated)))
273+
after, err := db.GetAIModelPriceByProviderModel(ctx, key)
274+
require.NoError(t, err)
275+
276+
require.True(t, after.UpdatedAt.After(before.UpdatedAt), "updated_at should advance when a price changes")
277+
})
278+
}
279+
})
177280
}
178281

179282
// TestSeed exercises the real embedded prices.json so we catch a corrupted,

‎coderd/database/querier.go‎

Lines changed: 3 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/queries.sql.go‎

Lines changed: 14 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/queries/aicostcontrol.sql‎

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,9 @@
22
-- Upsert a batch of (provider, model) rows from a JSON array. Each element
33
-- must have provider, model, and the four price fields; null prices are
44
-- written as SQL NULL.
5+
-- A conflicting row is only rewritten when a price differs, so updated_at
6+
-- records when a price last changed. Prices are nullable and a NULL on
7+
-- either side counts as a difference.
58
INSERT INTO ai_model_prices (
69
provider, model, input_price, output_price, cache_read_price, cache_write_price
710
)
@@ -18,7 +21,18 @@ ON CONFLICT (provider, model) DO UPDATE SET
1821
output_price = EXCLUDED.output_price,
1922
cache_read_price = EXCLUDED.cache_read_price,
2023
cache_write_price = EXCLUDED.cache_write_price,
21-
updated_at = NOW();
24+
updated_at = NOW()
25+
WHERE (
26+
ai_model_prices.input_price,
27+
ai_model_prices.output_price,
28+
ai_model_prices.cache_read_price,
29+
ai_model_prices.cache_write_price
30+
) IS DISTINCT FROM (
31+
EXCLUDED.input_price,
32+
EXCLUDED.output_price,
33+
EXCLUDED.cache_read_price,
34+
EXCLUDED.cache_write_price
35+
);
2236

2337
-- name: GetAIModelPriceByProviderModel :one
2438
SELECT *

0 commit comments

Comments
 (0)