mirror of
https://github.com/coder/coder.git
synced 2026-10-16 14:48:56 +08:00
feat: distinguish default and custom AI model prices (#28204)
## Description Prices in `ai_model_prices` come from two places: the price book embedded in each Coder release, and prices set through the experimental CLI (https://github.com/coder/coder/pull/27926). Until now, they were indistinguishable, so the startup seeder re-applied the book over every row, and a price someone set would silently revert on the next restart. This adds a `source` column with values `default` and `custom`, and keeps a row for each. Readers resolve a model to its custom price when one exists and fall back to the book, including the cost path that prices an interception. Keeping both rows means the price book stays current underneath an override, so a future delete can restore it. Note: the endpoint still only accepts models the price book does not cover. Overriding a price the book already carries is handled by the follow-up issue https://linear.app/codercom/issue/AIGOV-594. ## Changes - Add a migration for the `ai_model_price_source` enum and a `source` column, backfilled to `default` and then `NOT NULL`. `source` joins the primary key, so a model can hold one row per source. - Take a `source` parameter on `UpsertAIModelPrices`. Each source keeps its own row, so the price book and a custom price never overwrite each other. - Resolve to the custom price in `GetAIModelPriceByProviderModel` and `GetAIModelPrices`, so both the interception cost snapshot and the list endpoint use it. - Seed the price book as `default`, and write prices set through the endpoint as `custom`. Closes https://linear.app/codercom/issue/AIGOV-592 > [!NOTE] > Initially generated by Claude Opus 5, modified and reviewed by @ssncferreira
This commit is contained in:
@@ -51,7 +51,10 @@ func SeedFromBytes(ctx context.Context, db database.Store, data []byte) error {
|
||||
if len(rows) == 0 {
|
||||
return xerrors.New("price seed is empty")
|
||||
}
|
||||
return db.UpsertAIModelPrices(ctx, data)
|
||||
return db.UpsertAIModelPrices(ctx, database.UpsertAIModelPricesParams{
|
||||
Seed: data,
|
||||
Source: database.AIModelPriceSourceDefault,
|
||||
})
|
||||
}
|
||||
|
||||
func parseSeed(data []byte) ([]seedRow, error) {
|
||||
|
||||
@@ -75,6 +75,24 @@ func TestSeedFromBytes(t *testing.T) {
|
||||
require.Zero(t, gpt.CacheWritePrice.Int64)
|
||||
})
|
||||
|
||||
t.Run("SeededPricesAreDefault", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := testutil.Context(t, testutil.WaitShort)
|
||||
db, _ := dbtestutil.NewDB(t)
|
||||
|
||||
require.NoError(t, prices.SeedFromBytes(ctx, db, []byte(testSeedJSON)))
|
||||
|
||||
got, err := db.GetAIModelPriceByProviderModel(ctx, database.GetAIModelPriceByProviderModelParams{
|
||||
Provider: "openai", Model: "gpt-4o",
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, database.AIModelPriceSourceDefault, got.Source)
|
||||
require.Equal(t, int64(2_500_000), got.InputPrice.Int64)
|
||||
require.Equal(t, int64(10_000_000), got.OutputPrice.Int64)
|
||||
require.Equal(t, int64(1_250_000), got.CacheReadPrice.Int64)
|
||||
require.False(t, got.CacheWritePrice.Valid)
|
||||
})
|
||||
|
||||
t.Run("Idempotent", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := testutil.Context(t, testutil.WaitShort)
|
||||
@@ -108,14 +126,17 @@ func TestSeedFromBytes(t *testing.T) {
|
||||
// cache_write_price is set to a non-NULL value here even though the
|
||||
// embedded seed leaves it NULL for OpenAI; Seed must replace it with
|
||||
// NULL to keep the table in sync with the seed.
|
||||
require.NoError(t, db.UpsertAIModelPrices(ctx, []byte(`[{
|
||||
"provider": "openai",
|
||||
"model": "gpt-4o",
|
||||
"input_price": 1,
|
||||
"output_price": 2,
|
||||
"cache_read_price": 3,
|
||||
"cache_write_price": 4
|
||||
}]`)))
|
||||
require.NoError(t, db.UpsertAIModelPrices(ctx, database.UpsertAIModelPricesParams{
|
||||
Seed: []byte(`[{
|
||||
"provider": "openai",
|
||||
"model": "gpt-4o",
|
||||
"input_price": 1,
|
||||
"output_price": 2,
|
||||
"cache_read_price": 3,
|
||||
"cache_write_price": 4
|
||||
}]`),
|
||||
Source: database.AIModelPriceSourceDefault,
|
||||
}))
|
||||
before, err := db.GetAIModelPriceByProviderModel(ctx, database.GetAIModelPriceByProviderModelParams{
|
||||
Provider: "openai", Model: "gpt-4o",
|
||||
})
|
||||
@@ -143,14 +164,17 @@ func TestSeedFromBytes(t *testing.T) {
|
||||
|
||||
// Insert a row for a (provider, model) the seed doesn't cover. After
|
||||
// Seed it should still be there with its values intact.
|
||||
require.NoError(t, db.UpsertAIModelPrices(ctx, []byte(`[{
|
||||
"provider": "test-provider",
|
||||
"model": "test-model-not-in-seed",
|
||||
"input_price": 12345,
|
||||
"output_price": 67890,
|
||||
"cache_read_price": null,
|
||||
"cache_write_price": null
|
||||
}]`)))
|
||||
require.NoError(t, db.UpsertAIModelPrices(ctx, database.UpsertAIModelPricesParams{
|
||||
Seed: []byte(`[{
|
||||
"provider": "test-provider",
|
||||
"model": "test-model-not-in-seed",
|
||||
"input_price": 12345,
|
||||
"output_price": 67890,
|
||||
"cache_read_price": null,
|
||||
"cache_write_price": null
|
||||
}]`),
|
||||
Source: database.AIModelPriceSourceDefault,
|
||||
}))
|
||||
|
||||
require.NoError(t, prices.SeedFromBytes(ctx, db, []byte(testSeedJSON)))
|
||||
|
||||
@@ -162,6 +186,48 @@ func TestSeedFromBytes(t *testing.T) {
|
||||
require.Equal(t, int64(67890), got.OutputPrice.Int64)
|
||||
})
|
||||
|
||||
t.Run("LeavesCustomPricesUntouched", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := testutil.Context(t, testutil.WaitShort)
|
||||
db, _ := dbtestutil.NewDB(t)
|
||||
|
||||
// Price a model the seed also covers.
|
||||
require.NoError(t, db.UpsertAIModelPrices(ctx, database.UpsertAIModelPricesParams{
|
||||
Seed: []byte(`[{
|
||||
"provider": "openai",
|
||||
"model": "gpt-4o",
|
||||
"input_price": 1,
|
||||
"output_price": 2,
|
||||
"cache_read_price": 3,
|
||||
"cache_write_price": 4
|
||||
}]`),
|
||||
Source: database.AIModelPriceSourceCustom,
|
||||
}))
|
||||
before, err := db.GetAIModelPriceByProviderModel(ctx, database.GetAIModelPriceByProviderModelParams{
|
||||
Provider: "openai", Model: "gpt-4o",
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, database.AIModelPriceSourceCustom, before.Source)
|
||||
require.Equal(t, int64(1), before.InputPrice.Int64)
|
||||
require.Equal(t, int64(2), before.OutputPrice.Int64)
|
||||
require.Equal(t, int64(3), before.CacheReadPrice.Int64)
|
||||
require.Equal(t, int64(4), before.CacheWritePrice.Int64)
|
||||
|
||||
// Re-applying the price book writes its own row and leaves this one be.
|
||||
require.NoError(t, prices.SeedFromBytes(ctx, db, []byte(testSeedJSON)))
|
||||
|
||||
got, err := db.GetAIModelPriceByProviderModel(ctx, database.GetAIModelPriceByProviderModelParams{
|
||||
Provider: "openai", Model: "gpt-4o",
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, int64(1), got.InputPrice.Int64)
|
||||
require.Equal(t, int64(2), got.OutputPrice.Int64)
|
||||
require.Equal(t, int64(3), got.CacheReadPrice.Int64)
|
||||
require.Equal(t, int64(4), got.CacheWritePrice.Int64)
|
||||
require.Equal(t, database.AIModelPriceSourceCustom, got.Source)
|
||||
require.Equal(t, before.UpdatedAt, got.UpdatedAt)
|
||||
})
|
||||
|
||||
// Verifies the chain: AsAIBridged context -> dbauthz wrapper auth check
|
||||
// -> subjectAibridged's permission grant. A missing or wrong action on
|
||||
// the subject would surface as "unauthorized: rbac: forbidden" here, even
|
||||
|
||||
@@ -1731,7 +1731,6 @@ func TestRecordTokenUsage(t *testing.T) {
|
||||
// No override
|
||||
expectTokenUsageCostLookups(db, intc, nil, group, nil, price)
|
||||
|
||||
// input 300 + output 1200 + cache read 15 + cache write 40.
|
||||
const wantCost int64 = 1555
|
||||
|
||||
db.EXPECT().InTx(gomock.Any(), nil).DoAndReturn(
|
||||
@@ -2197,7 +2196,6 @@ func TestRecordTokenUsage(t *testing.T) {
|
||||
func(fn func(database.Store) error, _ *database.TxOptions) error { return fn(db) },
|
||||
)
|
||||
|
||||
// input 300 + output 1200 + cache read 15 + cache write 40.
|
||||
const wantCost int64 = 1555
|
||||
|
||||
db.EXPECT().InsertAIBridgeTokenUsage(gomock.Any(), gomock.Cond(func(p database.InsertAIBridgeTokenUsageParams) bool {
|
||||
@@ -2506,7 +2504,7 @@ func TestRecordTokenUsageAuthorized(t *testing.T) {
|
||||
"cache_write_price": 4_000_000,
|
||||
}})
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, rawDB.UpsertAIModelPrices(ctx, priceSeed), "seed model prices")
|
||||
require.NoError(t, rawDB.UpsertAIModelPrices(ctx, database.UpsertAIModelPricesParams{Seed: priceSeed, Source: database.AIModelPriceSourceDefault}), "seed model prices")
|
||||
|
||||
// The interception's provider name resolves to this provider, whose type keys
|
||||
// the price lookup.
|
||||
@@ -2559,7 +2557,6 @@ func TestRecordTokenUsageAuthorized(t *testing.T) {
|
||||
require.Equal(t, sql.NullInt64{Int64: 6_000_000, Valid: true}, tokenUsage.OutputPriceMicros, "output price")
|
||||
require.Equal(t, sql.NullInt64{Int64: 300_000, Valid: true}, tokenUsage.CacheReadPriceMicros, "cache read price")
|
||||
require.Equal(t, sql.NullInt64{Int64: 4_000_000, Valid: true}, tokenUsage.CacheWritePriceMicros, "cache write price")
|
||||
// input 300 + output 1200 + cache read 15 + cache write 40.
|
||||
const wantCost int64 = 1555
|
||||
require.Equal(t, sql.NullInt64{Int64: wantCost, Valid: true}, tokenUsage.CostMicros, "cost")
|
||||
|
||||
@@ -2577,6 +2574,149 @@ func TestRecordTokenUsageAuthorized(t *testing.T) {
|
||||
require.Equal(t, wantCost, spend.SpendMicros, "spend micros")
|
||||
}
|
||||
|
||||
// TestRecordTokenUsageModelPriceResolution covers which price an interception
|
||||
// snapshots when a model carries a price from the embedded book, a price set
|
||||
// through the API, or both.
|
||||
func TestRecordTokenUsageModelPriceResolution(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
const provider, model = "anthropic", "claude-sonnet-4-6"
|
||||
|
||||
priceSeed := func(input, output, cacheRead, cacheWrite int64) json.RawMessage {
|
||||
seed, err := json.Marshal([]map[string]any{{
|
||||
"provider": provider,
|
||||
"model": model,
|
||||
"input_price": input,
|
||||
"output_price": output,
|
||||
"cache_read_price": cacheRead,
|
||||
"cache_write_price": cacheWrite,
|
||||
}})
|
||||
require.NoError(t, err)
|
||||
return seed
|
||||
}
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
defaultSeed json.RawMessage
|
||||
customSeed json.RawMessage
|
||||
want database.AIBridgeTokenUsage
|
||||
}{
|
||||
{
|
||||
name: "DefaultOnly",
|
||||
defaultSeed: priceSeed(3_000_000, 6_000_000, 300_000, 4_000_000),
|
||||
want: database.AIBridgeTokenUsage{
|
||||
InputPriceMicros: sql.NullInt64{Int64: 3_000_000, Valid: true},
|
||||
OutputPriceMicros: sql.NullInt64{Int64: 6_000_000, Valid: true},
|
||||
CacheReadPriceMicros: sql.NullInt64{Int64: 300_000, Valid: true},
|
||||
CacheWritePriceMicros: sql.NullInt64{Int64: 4_000_000, Valid: true},
|
||||
// 100 input, 200 output, 50 cache read, and 10 cache write
|
||||
// tokens, priced per million: 300 + 1200 + 15 + 40.
|
||||
CostMicros: sql.NullInt64{Int64: 1555, Valid: true},
|
||||
},
|
||||
},
|
||||
{
|
||||
// A model the price book does not cover, priced through the API.
|
||||
name: "CustomOnly",
|
||||
customSeed: priceSeed(2_000_000, 4_000_000, 200_000, 1_000_000),
|
||||
want: database.AIBridgeTokenUsage{
|
||||
InputPriceMicros: sql.NullInt64{Int64: 2_000_000, Valid: true},
|
||||
OutputPriceMicros: sql.NullInt64{Int64: 4_000_000, Valid: true},
|
||||
CacheReadPriceMicros: sql.NullInt64{Int64: 200_000, Valid: true},
|
||||
CacheWritePriceMicros: sql.NullInt64{Int64: 1_000_000, Valid: true},
|
||||
// 100 input, 200 output, 50 cache read, and 10 cache write
|
||||
// tokens, priced per million: 200 + 800 + 10 + 10.
|
||||
CostMicros: sql.NullInt64{Int64: 1020, Valid: true},
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "CustomWinsOverDefault",
|
||||
defaultSeed: priceSeed(3_000_000, 6_000_000, 300_000, 4_000_000),
|
||||
customSeed: priceSeed(9_000_000, 12_000_000, 900_000, 8_000_000),
|
||||
want: database.AIBridgeTokenUsage{
|
||||
InputPriceMicros: sql.NullInt64{Int64: 9_000_000, Valid: true},
|
||||
OutputPriceMicros: sql.NullInt64{Int64: 12_000_000, Valid: true},
|
||||
CacheReadPriceMicros: sql.NullInt64{Int64: 900_000, Valid: true},
|
||||
CacheWritePriceMicros: sql.NullInt64{Int64: 8_000_000, Valid: true},
|
||||
// 100 input, 200 output, 50 cache read, and 10 cache write
|
||||
// tokens, priced per million: 900 + 2400 + 45 + 80.
|
||||
CostMicros: sql.NullInt64{Int64: 3425, Valid: true},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
logger := testutil.Logger(t)
|
||||
|
||||
rawDB, _ := dbtestutil.NewDB(t)
|
||||
authzDB := dbauthz.New(rawDB, rbac.NewStrictAuthorizer(prometheus.NewRegistry()), logger, coderdtest.AccessControlStorePointer())
|
||||
|
||||
org := dbgen.Organization(t, rawDB, database.Organization{})
|
||||
user := dbgen.User(t, rawDB, database.User{})
|
||||
dbgen.OrganizationMember(t, rawDB, database.OrganizationMember{OrganizationID: org.ID, UserID: user.ID})
|
||||
|
||||
if tt.defaultSeed != nil {
|
||||
require.NoError(t, rawDB.UpsertAIModelPrices(ctx, database.UpsertAIModelPricesParams{
|
||||
Seed: tt.defaultSeed,
|
||||
Source: database.AIModelPriceSourceDefault,
|
||||
}), "seed model prices")
|
||||
}
|
||||
if tt.customSeed != nil {
|
||||
require.NoError(t, rawDB.UpsertAIModelPrices(ctx, database.UpsertAIModelPricesParams{
|
||||
Seed: tt.customSeed,
|
||||
Source: database.AIModelPriceSourceCustom,
|
||||
}), "set custom model price")
|
||||
}
|
||||
|
||||
aiProvider := dbgen.AIProvider(t, rawDB, database.AIProvider{
|
||||
Name: "anthropic-eu",
|
||||
Type: database.AIProviderTypeAnthropic,
|
||||
})
|
||||
intc := dbgen.AIBridgeInterception(t, rawDB, database.InsertAIBridgeInterceptionParams{
|
||||
InitiatorID: user.ID,
|
||||
Provider: provider,
|
||||
ProviderName: aiProvider.Name,
|
||||
Model: model,
|
||||
}, nil)
|
||||
|
||||
srv, err := aibridgedserver.NewServer(ctx, aibridgedserver.Options{
|
||||
Store: authzDB,
|
||||
AISeatTracker: agplaiseats.Noop{},
|
||||
AccessURL: "/",
|
||||
GatewayCfg: codersdk.AIBridgeConfig{},
|
||||
Experiments: requiredExperiments,
|
||||
Logger: logger,
|
||||
Clock: quartz.NewReal(),
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
_, err = srv.RecordTokenUsage(ctx, &proto.RecordTokenUsageRequest{
|
||||
InterceptionId: intc.ID.String(),
|
||||
MsgId: "msg_price_resolution",
|
||||
InputTokens: 100,
|
||||
OutputTokens: 200,
|
||||
CacheReadInputTokens: 50,
|
||||
CacheWriteInputTokens: 10,
|
||||
CreatedAt: timestamppb.New(time.Date(2026, 6, 25, 14, 30, 0, 0, time.UTC)),
|
||||
})
|
||||
require.NoError(t, err, "record token usage")
|
||||
|
||||
tokenUsages, err := rawDB.GetAIBridgeTokenUsagesByInterceptionID(ctx, intc.ID)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, tokenUsages, 1)
|
||||
|
||||
tokenUsage := tokenUsages[0]
|
||||
require.Equal(t, tt.want.InputPriceMicros, tokenUsage.InputPriceMicros, "input price")
|
||||
require.Equal(t, tt.want.OutputPriceMicros, tokenUsage.OutputPriceMicros, "output price")
|
||||
require.Equal(t, tt.want.CacheReadPriceMicros, tokenUsage.CacheReadPriceMicros, "cache read price")
|
||||
require.Equal(t, tt.want.CacheWritePriceMicros, tokenUsage.CacheWritePriceMicros, "cache write price")
|
||||
require.Equal(t, tt.want.CostMicros, tokenUsage.CostMicros, "cost")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestRecordTokenUsageProviderResolution covers provider resolution against a real
|
||||
// database through dbauthz, where the live-row filter and name reuse apply.
|
||||
func TestRecordTokenUsageProviderResolution(t *testing.T) {
|
||||
@@ -2601,7 +2741,7 @@ func TestRecordTokenUsageProviderResolution(t *testing.T) {
|
||||
{"provider": string(database.AIProviderTypeAzure), "model": gptModel, "input_price": azureInputPrice},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, rawDB.UpsertAIModelPrices(setupCtx, priceSeed), "seed model prices")
|
||||
require.NoError(t, rawDB.UpsertAIModelPrices(setupCtx, database.UpsertAIModelPricesParams{Seed: priceSeed, Source: database.AIModelPriceSourceDefault}), "seed model prices")
|
||||
|
||||
srv, err := aibridgedserver.NewServer(setupCtx, aibridgedserver.Options{
|
||||
Store: authzDB,
|
||||
|
||||
@@ -8851,11 +8851,11 @@ func (q *querier) UpdateWorkspacesTTLByTemplateID(ctx context.Context, arg datab
|
||||
return q.db.UpdateWorkspacesTTLByTemplateID(ctx, arg)
|
||||
}
|
||||
|
||||
func (q *querier) UpsertAIModelPrices(ctx context.Context, seed json.RawMessage) error {
|
||||
func (q *querier) UpsertAIModelPrices(ctx context.Context, arg database.UpsertAIModelPricesParams) error {
|
||||
if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceAiModelPrice); err != nil {
|
||||
return err
|
||||
}
|
||||
return q.db.UpsertAIModelPrices(ctx, seed)
|
||||
return q.db.UpsertAIModelPrices(ctx, arg)
|
||||
}
|
||||
|
||||
func (q *querier) UpsertAISeatState(ctx context.Context, arg database.UpsertAISeatStateParams) (bool, error) {
|
||||
|
||||
@@ -6927,7 +6927,10 @@ func (s *MethodTestSuite) TestAIBridge() {
|
||||
|
||||
s.Run("UpsertAIModelPrices", s.Mocked(func(db *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
|
||||
db.EXPECT().UpsertAIModelPrices(gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
|
||||
check.Args(json.RawMessage(`[]`)).Asserts(rbac.ResourceAiModelPrice, policy.ActionUpdate)
|
||||
check.Args(database.UpsertAIModelPricesParams{
|
||||
Seed: json.RawMessage(`[]`),
|
||||
Source: database.AIModelPriceSourceDefault,
|
||||
}).Asserts(rbac.ResourceAiModelPrice, policy.ActionUpdate)
|
||||
}))
|
||||
|
||||
s.Run("GetAIModelPriceByProviderModel", s.Mocked(func(db *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
|
||||
|
||||
+2
-3
@@ -5,7 +5,6 @@ package dbmetrics
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"slices"
|
||||
"time"
|
||||
|
||||
@@ -6225,9 +6224,9 @@ func (m queryMetricsStore) UpdateWorkspacesTTLByTemplateID(ctx context.Context,
|
||||
return r0
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) UpsertAIModelPrices(ctx context.Context, seed json.RawMessage) error {
|
||||
func (m queryMetricsStore) UpsertAIModelPrices(ctx context.Context, arg database.UpsertAIModelPricesParams) error {
|
||||
start := time.Now()
|
||||
r0 := m.s.UpsertAIModelPrices(ctx, seed)
|
||||
r0 := m.s.UpsertAIModelPrices(ctx, arg)
|
||||
m.queryLatencies.WithLabelValues("UpsertAIModelPrices").Observe(time.Since(start).Seconds())
|
||||
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "UpsertAIModelPrices").Inc()
|
||||
return r0
|
||||
|
||||
Generated
+4
-5
@@ -11,7 +11,6 @@ package dbmock
|
||||
|
||||
import (
|
||||
context "context"
|
||||
json "encoding/json"
|
||||
reflect "reflect"
|
||||
time "time"
|
||||
|
||||
@@ -11716,17 +11715,17 @@ func (mr *MockStoreMockRecorder) UpdateWorkspacesTTLByTemplateID(ctx, arg any) *
|
||||
}
|
||||
|
||||
// UpsertAIModelPrices mocks base method.
|
||||
func (m *MockStore) UpsertAIModelPrices(ctx context.Context, seed json.RawMessage) error {
|
||||
func (m *MockStore) UpsertAIModelPrices(ctx context.Context, arg database.UpsertAIModelPricesParams) error {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "UpsertAIModelPrices", ctx, seed)
|
||||
ret := m.ctrl.Call(m, "UpsertAIModelPrices", ctx, arg)
|
||||
ret0, _ := ret[0].(error)
|
||||
return ret0
|
||||
}
|
||||
|
||||
// UpsertAIModelPrices indicates an expected call of UpsertAIModelPrices.
|
||||
func (mr *MockStoreMockRecorder) UpsertAIModelPrices(ctx, seed any) *gomock.Call {
|
||||
func (mr *MockStoreMockRecorder) UpsertAIModelPrices(ctx, arg any) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpsertAIModelPrices", reflect.TypeOf((*MockStore)(nil).UpsertAIModelPrices), ctx, seed)
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpsertAIModelPrices", reflect.TypeOf((*MockStore)(nil).UpsertAIModelPrices), ctx, arg)
|
||||
}
|
||||
|
||||
// UpsertAISeatState mocks base method.
|
||||
|
||||
Generated
+9
-1
@@ -10,6 +10,11 @@ CREATE TYPE agent_key_scope_enum AS ENUM (
|
||||
'no_user_data'
|
||||
);
|
||||
|
||||
CREATE TYPE ai_model_price_source AS ENUM (
|
||||
'default',
|
||||
'custom'
|
||||
);
|
||||
|
||||
CREATE TYPE ai_provider_type AS ENUM (
|
||||
'openai',
|
||||
'anthropic',
|
||||
@@ -1573,6 +1578,7 @@ CREATE TABLE ai_model_prices (
|
||||
cache_write_price bigint,
|
||||
created_at timestamp with time zone DEFAULT now() NOT NULL,
|
||||
updated_at timestamp with time zone DEFAULT now() NOT NULL,
|
||||
source ai_model_price_source NOT NULL,
|
||||
CONSTRAINT ai_model_prices_cache_read_price_check CHECK ((cache_read_price >= 0)),
|
||||
CONSTRAINT ai_model_prices_cache_write_price_check CHECK ((cache_write_price >= 0)),
|
||||
CONSTRAINT ai_model_prices_input_price_check CHECK ((input_price >= 0)),
|
||||
@@ -1581,6 +1587,8 @@ CREATE TABLE ai_model_prices (
|
||||
|
||||
COMMENT ON TABLE ai_model_prices IS 'Per-model token prices used by AI Bridge to compute interception cost.';
|
||||
|
||||
COMMENT ON COLUMN ai_model_prices.source IS 'Where the price came from: default for the embedded price book, custom for a price set through the API. Both can exist for the same model.';
|
||||
|
||||
CREATE TABLE ai_provider_keys (
|
||||
id uuid DEFAULT gen_random_uuid() NOT NULL,
|
||||
provider_id uuid NOT NULL,
|
||||
@@ -4315,7 +4323,7 @@ ALTER TABLE ONLY ai_gateway_keys
|
||||
ADD CONSTRAINT ai_gateway_keys_pkey PRIMARY KEY (id);
|
||||
|
||||
ALTER TABLE ONLY ai_model_prices
|
||||
ADD CONSTRAINT ai_model_prices_pkey PRIMARY KEY (provider, model);
|
||||
ADD CONSTRAINT ai_model_prices_pkey PRIMARY KEY (provider, model, source);
|
||||
|
||||
ALTER TABLE ONLY ai_provider_keys
|
||||
ADD CONSTRAINT ai_provider_keys_pkey PRIMARY KEY (id);
|
||||
|
||||
@@ -0,0 +1,11 @@
|
||||
-- Custom prices have no place in the single-row model, and a model may hold
|
||||
-- both rows, so they are dropped to restore uniqueness on (provider, model).
|
||||
DELETE FROM ai_model_prices WHERE source = 'custom';
|
||||
|
||||
ALTER TABLE ai_model_prices DROP CONSTRAINT ai_model_prices_pkey;
|
||||
|
||||
ALTER TABLE ai_model_prices ADD PRIMARY KEY (provider, model);
|
||||
|
||||
ALTER TABLE ai_model_prices DROP COLUMN source;
|
||||
|
||||
DROP TYPE ai_model_price_source;
|
||||
@@ -0,0 +1,17 @@
|
||||
-- source records where a price came from, either the price book embedded in
|
||||
-- each Coder release or a price set through the API. Both can exist for the
|
||||
-- same model, so source joins the primary key.
|
||||
|
||||
CREATE TYPE ai_model_price_source AS ENUM ('default', 'custom');
|
||||
|
||||
ALTER TABLE ai_model_prices ADD COLUMN source ai_model_price_source;
|
||||
|
||||
UPDATE ai_model_prices SET source = 'default' WHERE source IS NULL;
|
||||
|
||||
ALTER TABLE ai_model_prices ALTER COLUMN source SET NOT NULL;
|
||||
|
||||
ALTER TABLE ai_model_prices DROP CONSTRAINT ai_model_prices_pkey;
|
||||
|
||||
ALTER TABLE ai_model_prices ADD PRIMARY KEY (provider, model, source);
|
||||
|
||||
COMMENT ON COLUMN ai_model_prices.source IS 'Where the price came from: default for the embedded price book, custom for a price set through the API. Both can exist for the same model.';
|
||||
Generated
+60
@@ -89,6 +89,64 @@ func AllAIBridgeInterceptionErrorTypeValues() []AIBridgeInterceptionErrorType {
|
||||
}
|
||||
}
|
||||
|
||||
type AIModelPriceSource string
|
||||
|
||||
const (
|
||||
AIModelPriceSourceDefault AIModelPriceSource = "default"
|
||||
AIModelPriceSourceCustom AIModelPriceSource = "custom"
|
||||
)
|
||||
|
||||
func (e *AIModelPriceSource) Scan(src interface{}) error {
|
||||
switch s := src.(type) {
|
||||
case []byte:
|
||||
*e = AIModelPriceSource(s)
|
||||
case string:
|
||||
*e = AIModelPriceSource(s)
|
||||
default:
|
||||
return fmt.Errorf("unsupported scan type for AIModelPriceSource: %T", src)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
type NullAIModelPriceSource struct {
|
||||
AIModelPriceSource AIModelPriceSource `json:"ai_model_price_source"`
|
||||
Valid bool `json:"valid"` // Valid is true if AIModelPriceSource is not NULL
|
||||
}
|
||||
|
||||
// Scan implements the Scanner interface.
|
||||
func (ns *NullAIModelPriceSource) Scan(value interface{}) error {
|
||||
if value == nil {
|
||||
ns.AIModelPriceSource, ns.Valid = "", false
|
||||
return nil
|
||||
}
|
||||
ns.Valid = true
|
||||
return ns.AIModelPriceSource.Scan(value)
|
||||
}
|
||||
|
||||
// Value implements the driver Valuer interface.
|
||||
func (ns NullAIModelPriceSource) Value() (driver.Value, error) {
|
||||
if !ns.Valid {
|
||||
return nil, nil
|
||||
}
|
||||
return string(ns.AIModelPriceSource), nil
|
||||
}
|
||||
|
||||
func (e AIModelPriceSource) Valid() bool {
|
||||
switch e {
|
||||
case AIModelPriceSourceDefault,
|
||||
AIModelPriceSourceCustom:
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func AllAIModelPriceSourceValues() []AIModelPriceSource {
|
||||
return []AIModelPriceSource{
|
||||
AIModelPriceSourceDefault,
|
||||
AIModelPriceSourceCustom,
|
||||
}
|
||||
}
|
||||
|
||||
type AIProviderType string
|
||||
|
||||
const (
|
||||
@@ -4827,6 +4885,8 @@ type AIModelPrice struct {
|
||||
CacheWritePrice sql.NullInt64 `db:"cache_write_price" json:"cache_write_price"`
|
||||
CreatedAt time.Time `db:"created_at" json:"created_at"`
|
||||
UpdatedAt time.Time `db:"updated_at" json:"updated_at"`
|
||||
// Where the price came from: default for the embedded price book, custom for a price set through the API. Both can exist for the same model.
|
||||
Source AIModelPriceSource `db:"source" json:"source"`
|
||||
}
|
||||
|
||||
// Runtime configuration for AI providers. Authoritative source for the provider set served by aibridged. Replaces deployment-time CODER_AIBRIDGE_* environment variables.
|
||||
|
||||
Generated
+12
-8
@@ -6,7 +6,6 @@ package database
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
@@ -324,7 +323,11 @@ type sqlcQuerier interface {
|
||||
// returning the matched key. The lookup is an exact match on a unique index,
|
||||
// so a returned row is itself proof the secret is valid.
|
||||
GetAIGatewayKeyByHashedSecret(ctx context.Context, hashedSecret []byte) (AIGatewayKey, error)
|
||||
// Returns the price in effect for the model, preferring a custom price over
|
||||
// the price book.
|
||||
GetAIModelPriceByProviderModel(ctx context.Context, arg GetAIModelPriceByProviderModelParams) (AIModelPrice, error)
|
||||
// Returns the price in effect for each model, preferring a custom price over
|
||||
// the price book.
|
||||
GetAIModelPrices(ctx context.Context, arg GetAIModelPricesParams) ([]AIModelPrice, error)
|
||||
GetAIProviderByID(ctx context.Context, id uuid.UUID) (AIProvider, error)
|
||||
// Lock the provider row until the model-config write completes. The
|
||||
@@ -1616,13 +1619,14 @@ type sqlcQuerier interface {
|
||||
UpdateWorkspaceTTL(ctx context.Context, arg UpdateWorkspaceTTLParams) error
|
||||
UpdateWorkspacesDormantDeletingAtByTemplateID(ctx context.Context, arg UpdateWorkspacesDormantDeletingAtByTemplateIDParams) ([]WorkspaceTable, error)
|
||||
UpdateWorkspacesTTLByTemplateID(ctx context.Context, arg UpdateWorkspacesTTLByTemplateIDParams) error
|
||||
// Upsert a batch of (provider, model) rows from a JSON array. Each element
|
||||
// must have provider, model, and the four price fields; null prices are
|
||||
// written as SQL NULL.
|
||||
// A conflicting row is only rewritten when a price differs, so updated_at
|
||||
// records when a price last changed. Prices are nullable and a NULL on
|
||||
// either side counts as a difference.
|
||||
UpsertAIModelPrices(ctx context.Context, seed json.RawMessage) error
|
||||
// Upsert a batch of model prices from a JSON array, all recorded under the
|
||||
// given source. Each element must have provider, model, and the four price
|
||||
// fields, and null prices are written as SQL NULL.
|
||||
// Each source keeps its own row, so the price book and a custom price never
|
||||
// overwrite each other. A conflicting row is only rewritten when a price
|
||||
// differs, so updated_at records when a price last changed. Prices are
|
||||
// nullable and a NULL on either side counts as a difference.
|
||||
UpsertAIModelPrices(ctx context.Context, arg UpsertAIModelPricesParams) error
|
||||
// Returns true if a new rows was inserted, false otherwise.
|
||||
UpsertAISeatState(ctx context.Context, arg UpsertAISeatStateParams) (bool, error)
|
||||
UpsertAnnouncementBanners(ctx context.Context, value string) error
|
||||
|
||||
@@ -19059,6 +19059,85 @@ func TestOAuth2ProviderScopeNotEmpty(t *testing.T) {
|
||||
})
|
||||
}
|
||||
|
||||
func TestGetAIModelPriceByProviderModel(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
const defaultSeed = `[{"provider":"anthropic","model":"model-a","input_price":1,"output_price":2,"cache_read_price":3,"cache_write_price":4}]`
|
||||
const customSeed = `[{"provider":"anthropic","model":"model-a","input_price":5,"output_price":6,"cache_read_price":0,"cache_write_price":null}]`
|
||||
|
||||
defaultPrices := database.AIModelPrice{
|
||||
InputPrice: sql.NullInt64{Int64: 1, Valid: true},
|
||||
OutputPrice: sql.NullInt64{Int64: 2, Valid: true},
|
||||
CacheReadPrice: sql.NullInt64{Int64: 3, Valid: true},
|
||||
CacheWritePrice: sql.NullInt64{Int64: 4, Valid: true},
|
||||
}
|
||||
customPrices := database.AIModelPrice{
|
||||
InputPrice: sql.NullInt64{Int64: 5, Valid: true},
|
||||
OutputPrice: sql.NullInt64{Int64: 6, Valid: true},
|
||||
CacheReadPrice: sql.NullInt64{Int64: 0, Valid: true},
|
||||
CacheWritePrice: sql.NullInt64{},
|
||||
}
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
seeds []database.UpsertAIModelPricesParams
|
||||
want database.AIModelPrice
|
||||
wantSource database.AIModelPriceSource
|
||||
}{
|
||||
{
|
||||
name: "DefaultOnly",
|
||||
seeds: []database.UpsertAIModelPricesParams{{Seed: []byte(defaultSeed), Source: database.AIModelPriceSourceDefault}},
|
||||
want: defaultPrices,
|
||||
wantSource: database.AIModelPriceSourceDefault,
|
||||
},
|
||||
{
|
||||
name: "CustomOnly",
|
||||
seeds: []database.UpsertAIModelPricesParams{{Seed: []byte(customSeed), Source: database.AIModelPriceSourceCustom}},
|
||||
want: customPrices,
|
||||
wantSource: database.AIModelPriceSourceCustom,
|
||||
},
|
||||
{
|
||||
name: "CustomWinsOverDefault",
|
||||
seeds: []database.UpsertAIModelPricesParams{
|
||||
{Seed: []byte(defaultSeed), Source: database.AIModelPriceSourceDefault},
|
||||
{Seed: []byte(customSeed), Source: database.AIModelPriceSourceCustom},
|
||||
},
|
||||
want: customPrices,
|
||||
wantSource: database.AIModelPriceSourceCustom,
|
||||
},
|
||||
{
|
||||
name: "CustomWinsWhateverTheWriteOrder",
|
||||
seeds: []database.UpsertAIModelPricesParams{
|
||||
{Seed: []byte(customSeed), Source: database.AIModelPriceSourceCustom},
|
||||
{Seed: []byte(defaultSeed), Source: database.AIModelPriceSourceDefault},
|
||||
},
|
||||
want: customPrices,
|
||||
wantSource: database.AIModelPriceSourceCustom,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := testutil.Context(t, testutil.WaitShort)
|
||||
db, _ := dbtestutil.NewDB(t)
|
||||
for _, seed := range tt.seeds {
|
||||
require.NoError(t, db.UpsertAIModelPrices(ctx, seed))
|
||||
}
|
||||
|
||||
got, err := db.GetAIModelPriceByProviderModel(ctx, database.GetAIModelPriceByProviderModelParams{
|
||||
Provider: "anthropic", Model: "model-a",
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, tt.wantSource, got.Source)
|
||||
require.Equal(t, tt.want.InputPrice, got.InputPrice)
|
||||
require.Equal(t, tt.want.OutputPrice, got.OutputPrice)
|
||||
require.Equal(t, tt.want.CacheReadPrice, got.CacheReadPrice)
|
||||
require.Equal(t, tt.want.CacheWritePrice, got.CacheWritePrice)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetAIModelPrices(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
@@ -19070,15 +19149,21 @@ func TestGetAIModelPrices(t *testing.T) {
|
||||
{"provider":"openai","model":"model-a","input_price":3,"output_price":null,"cache_read_price":null,"cache_write_price":null}
|
||||
]`
|
||||
|
||||
// A custom price for anthropic/model-a, which the seed above also covers.
|
||||
const customSeed = `[{"provider":"anthropic","model":"model-a","input_price":9,"output_price":null,"cache_read_price":null,"cache_write_price":null}]`
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
params database.GetAIModelPricesParams
|
||||
want []string
|
||||
name string
|
||||
customSeed string
|
||||
params database.GetAIModelPricesParams
|
||||
want []string
|
||||
wantPrices []int64
|
||||
}{
|
||||
{
|
||||
name: "NoFilterReturnsEveryPrice",
|
||||
params: database.GetAIModelPricesParams{},
|
||||
want: []string{"anthropic/model-a", "anthropic/model-b", "openai/model-a"},
|
||||
name: "NoFilterReturnsEveryPrice",
|
||||
params: database.GetAIModelPricesParams{},
|
||||
want: []string{"anthropic/model-a", "anthropic/model-b", "openai/model-a"},
|
||||
wantPrices: []int64{1, 2, 3},
|
||||
},
|
||||
{
|
||||
name: "ByProvider",
|
||||
@@ -19105,6 +19190,14 @@ func TestGetAIModelPrices(t *testing.T) {
|
||||
params: database.GetAIModelPricesParams{Provider: "openai", Model: "model-b"},
|
||||
want: nil,
|
||||
},
|
||||
{
|
||||
// The anthropic/model-a is reported once, at the custom one.
|
||||
name: "ResolvesToTheCustomPrice",
|
||||
customSeed: customSeed,
|
||||
params: database.GetAIModelPricesParams{},
|
||||
want: []string{"anthropic/model-a", "anthropic/model-b", "openai/model-a"},
|
||||
wantPrices: []int64{9, 2, 3},
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
@@ -19112,7 +19205,10 @@ func TestGetAIModelPrices(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := testutil.Context(t, testutil.WaitShort)
|
||||
db, _ := dbtestutil.NewDB(t)
|
||||
require.NoError(t, db.UpsertAIModelPrices(ctx, []byte(seed)))
|
||||
require.NoError(t, db.UpsertAIModelPrices(ctx, database.UpsertAIModelPricesParams{Seed: []byte(seed), Source: database.AIModelPriceSourceDefault}))
|
||||
if tt.customSeed != "" {
|
||||
require.NoError(t, db.UpsertAIModelPrices(ctx, database.UpsertAIModelPricesParams{Seed: []byte(tt.customSeed), Source: database.AIModelPriceSourceCustom}))
|
||||
}
|
||||
|
||||
prices, err := db.GetAIModelPrices(ctx, tt.params)
|
||||
require.NoError(t, err)
|
||||
@@ -19126,6 +19222,14 @@ func TestGetAIModelPrices(t *testing.T) {
|
||||
return
|
||||
}
|
||||
require.Equal(t, tt.want, got)
|
||||
|
||||
if tt.wantPrices != nil {
|
||||
gotPrices := make([]int64, 0, len(prices))
|
||||
for _, price := range prices {
|
||||
gotPrices = append(gotPrices, price.InputPrice.Int64)
|
||||
}
|
||||
require.Equal(t, tt.wantPrices, gotPrices)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
Generated
+30
-15
@@ -2862,9 +2862,11 @@ func (q *sqlQuerier) ExportOrganizationAISpend(ctx context.Context, arg ExportOr
|
||||
}
|
||||
|
||||
const getAIModelPriceByProviderModel = `-- name: GetAIModelPriceByProviderModel :one
|
||||
SELECT provider, model, input_price, output_price, cache_read_price, cache_write_price, created_at, updated_at
|
||||
SELECT provider, model, input_price, output_price, cache_read_price, cache_write_price, created_at, updated_at, source
|
||||
FROM ai_model_prices
|
||||
WHERE provider = $1 AND model = $2
|
||||
ORDER BY CASE WHEN source = 'custom' THEN 0 ELSE 1 END ASC
|
||||
LIMIT 1
|
||||
`
|
||||
|
||||
type GetAIModelPriceByProviderModelParams struct {
|
||||
@@ -2872,6 +2874,8 @@ type GetAIModelPriceByProviderModelParams struct {
|
||||
Model string `db:"model" json:"model"`
|
||||
}
|
||||
|
||||
// Returns the price in effect for the model, preferring a custom price over
|
||||
// the price book.
|
||||
func (q *sqlQuerier) GetAIModelPriceByProviderModel(ctx context.Context, arg GetAIModelPriceByProviderModelParams) (AIModelPrice, error) {
|
||||
row := q.db.QueryRowContext(ctx, getAIModelPriceByProviderModel, arg.Provider, arg.Model)
|
||||
var i AIModelPrice
|
||||
@@ -2884,12 +2888,13 @@ func (q *sqlQuerier) GetAIModelPriceByProviderModel(ctx context.Context, arg Get
|
||||
&i.CacheWritePrice,
|
||||
&i.CreatedAt,
|
||||
&i.UpdatedAt,
|
||||
&i.Source,
|
||||
)
|
||||
return i, err
|
||||
}
|
||||
|
||||
const getAIModelPrices = `-- name: GetAIModelPrices :many
|
||||
SELECT provider, model, input_price, output_price, cache_read_price, cache_write_price, created_at, updated_at
|
||||
SELECT DISTINCT ON (provider, model) provider, model, input_price, output_price, cache_read_price, cache_write_price, created_at, updated_at, source
|
||||
FROM ai_model_prices
|
||||
-- Filter by provider
|
||||
WHERE CASE
|
||||
@@ -2903,7 +2908,7 @@ WHERE CASE
|
||||
model = $2
|
||||
ELSE true
|
||||
END
|
||||
ORDER BY provider, model
|
||||
ORDER BY provider ASC, model ASC, CASE WHEN source = 'custom' THEN 0 ELSE 1 END ASC
|
||||
`
|
||||
|
||||
type GetAIModelPricesParams struct {
|
||||
@@ -2911,6 +2916,8 @@ type GetAIModelPricesParams struct {
|
||||
Model string `db:"model" json:"model"`
|
||||
}
|
||||
|
||||
// Returns the price in effect for each model, preferring a custom price over
|
||||
// the price book.
|
||||
func (q *sqlQuerier) GetAIModelPrices(ctx context.Context, arg GetAIModelPricesParams) ([]AIModelPrice, error) {
|
||||
rows, err := q.db.QueryContext(ctx, getAIModelPrices, arg.Provider, arg.Model)
|
||||
if err != nil {
|
||||
@@ -2929,6 +2936,7 @@ func (q *sqlQuerier) GetAIModelPrices(ctx context.Context, arg GetAIModelPricesP
|
||||
&i.CacheWritePrice,
|
||||
&i.CreatedAt,
|
||||
&i.UpdatedAt,
|
||||
&i.Source,
|
||||
); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -3498,7 +3506,7 @@ func (q *sqlQuerier) IncrementUserAIDailySpend(ctx context.Context, arg Incremen
|
||||
|
||||
const upsertAIModelPrices = `-- name: UpsertAIModelPrices :exec
|
||||
INSERT INTO ai_model_prices (
|
||||
provider, model, input_price, output_price, cache_read_price, cache_write_price
|
||||
provider, model, input_price, output_price, cache_read_price, cache_write_price, source
|
||||
)
|
||||
SELECT
|
||||
elem->>'provider',
|
||||
@@ -3506,9 +3514,10 @@ SELECT
|
||||
(elem->>'input_price')::bigint,
|
||||
(elem->>'output_price')::bigint,
|
||||
(elem->>'cache_read_price')::bigint,
|
||||
(elem->>'cache_write_price')::bigint
|
||||
FROM jsonb_array_elements($1::jsonb) AS elem
|
||||
ON CONFLICT (provider, model) DO UPDATE SET
|
||||
(elem->>'cache_write_price')::bigint,
|
||||
$1::ai_model_price_source
|
||||
FROM jsonb_array_elements($2::jsonb) AS elem
|
||||
ON CONFLICT (provider, model, source) DO UPDATE SET
|
||||
input_price = EXCLUDED.input_price,
|
||||
output_price = EXCLUDED.output_price,
|
||||
cache_read_price = EXCLUDED.cache_read_price,
|
||||
@@ -3527,14 +3536,20 @@ WHERE (
|
||||
)
|
||||
`
|
||||
|
||||
// Upsert a batch of (provider, model) rows from a JSON array. Each element
|
||||
// must have provider, model, and the four price fields; null prices are
|
||||
// written as SQL NULL.
|
||||
// A conflicting row is only rewritten when a price differs, so updated_at
|
||||
// records when a price last changed. Prices are nullable and a NULL on
|
||||
// either side counts as a difference.
|
||||
func (q *sqlQuerier) UpsertAIModelPrices(ctx context.Context, seed json.RawMessage) error {
|
||||
_, err := q.db.ExecContext(ctx, upsertAIModelPrices, seed)
|
||||
type UpsertAIModelPricesParams struct {
|
||||
Source AIModelPriceSource `db:"source" json:"source"`
|
||||
Seed json.RawMessage `db:"seed" json:"seed"`
|
||||
}
|
||||
|
||||
// Upsert a batch of model prices from a JSON array, all recorded under the
|
||||
// given source. Each element must have provider, model, and the four price
|
||||
// fields, and null prices are written as SQL NULL.
|
||||
// Each source keeps its own row, so the price book and a custom price never
|
||||
// overwrite each other. A conflicting row is only rewritten when a price
|
||||
// differs, so updated_at records when a price last changed. Prices are
|
||||
// nullable and a NULL on either side counts as a difference.
|
||||
func (q *sqlQuerier) UpsertAIModelPrices(ctx context.Context, arg UpsertAIModelPricesParams) error {
|
||||
_, err := q.db.ExecContext(ctx, upsertAIModelPrices, arg.Source, arg.Seed)
|
||||
return err
|
||||
}
|
||||
|
||||
|
||||
@@ -1,12 +1,13 @@
|
||||
-- name: UpsertAIModelPrices :exec
|
||||
-- Upsert a batch of (provider, model) rows from a JSON array. Each element
|
||||
-- must have provider, model, and the four price fields; null prices are
|
||||
-- written as SQL NULL.
|
||||
-- A conflicting row is only rewritten when a price differs, so updated_at
|
||||
-- records when a price last changed. Prices are nullable and a NULL on
|
||||
-- either side counts as a difference.
|
||||
-- Upsert a batch of model prices from a JSON array, all recorded under the
|
||||
-- given source. Each element must have provider, model, and the four price
|
||||
-- fields, and null prices are written as SQL NULL.
|
||||
-- Each source keeps its own row, so the price book and a custom price never
|
||||
-- overwrite each other. A conflicting row is only rewritten when a price
|
||||
-- differs, so updated_at records when a price last changed. Prices are
|
||||
-- nullable and a NULL on either side counts as a difference.
|
||||
INSERT INTO ai_model_prices (
|
||||
provider, model, input_price, output_price, cache_read_price, cache_write_price
|
||||
provider, model, input_price, output_price, cache_read_price, cache_write_price, source
|
||||
)
|
||||
SELECT
|
||||
elem->>'provider',
|
||||
@@ -14,9 +15,10 @@ SELECT
|
||||
(elem->>'input_price')::bigint,
|
||||
(elem->>'output_price')::bigint,
|
||||
(elem->>'cache_read_price')::bigint,
|
||||
(elem->>'cache_write_price')::bigint
|
||||
(elem->>'cache_write_price')::bigint,
|
||||
@source::ai_model_price_source
|
||||
FROM jsonb_array_elements(@seed::jsonb) AS elem
|
||||
ON CONFLICT (provider, model) DO UPDATE SET
|
||||
ON CONFLICT (provider, model, source) DO UPDATE SET
|
||||
input_price = EXCLUDED.input_price,
|
||||
output_price = EXCLUDED.output_price,
|
||||
cache_read_price = EXCLUDED.cache_read_price,
|
||||
@@ -35,12 +37,18 @@ WHERE (
|
||||
);
|
||||
|
||||
-- name: GetAIModelPriceByProviderModel :one
|
||||
-- Returns the price in effect for the model, preferring a custom price over
|
||||
-- the price book.
|
||||
SELECT *
|
||||
FROM ai_model_prices
|
||||
WHERE provider = @provider AND model = @model;
|
||||
WHERE provider = @provider AND model = @model
|
||||
ORDER BY CASE WHEN source = 'custom' THEN 0 ELSE 1 END ASC
|
||||
LIMIT 1;
|
||||
|
||||
-- name: GetAIModelPrices :many
|
||||
SELECT *
|
||||
-- Returns the price in effect for each model, preferring a custom price over
|
||||
-- the price book.
|
||||
SELECT DISTINCT ON (provider, model) *
|
||||
FROM ai_model_prices
|
||||
-- Filter by provider
|
||||
WHERE CASE
|
||||
@@ -54,7 +62,7 @@ WHERE CASE
|
||||
model = @model
|
||||
ELSE true
|
||||
END
|
||||
ORDER BY provider, model;
|
||||
ORDER BY provider ASC, model ASC, CASE WHEN source = 'custom' THEN 0 ELSE 1 END ASC;
|
||||
|
||||
-- name: GetGroupAIBudget :one
|
||||
SELECT *
|
||||
|
||||
Generated
+1
-1
@@ -8,7 +8,7 @@ type UniqueConstraint string
|
||||
const (
|
||||
UniqueAgentStatsPkey UniqueConstraint = "agent_stats_pkey" // ALTER TABLE ONLY workspace_agent_stats ADD CONSTRAINT agent_stats_pkey PRIMARY KEY (id);
|
||||
UniqueAIGatewayKeysPkey UniqueConstraint = "ai_gateway_keys_pkey" // ALTER TABLE ONLY ai_gateway_keys ADD CONSTRAINT ai_gateway_keys_pkey PRIMARY KEY (id);
|
||||
UniqueAIModelPricesPkey UniqueConstraint = "ai_model_prices_pkey" // ALTER TABLE ONLY ai_model_prices ADD CONSTRAINT ai_model_prices_pkey PRIMARY KEY (provider, model);
|
||||
UniqueAIModelPricesPkey UniqueConstraint = "ai_model_prices_pkey" // ALTER TABLE ONLY ai_model_prices ADD CONSTRAINT ai_model_prices_pkey PRIMARY KEY (provider, model, source);
|
||||
UniqueAIProviderKeysPkey UniqueConstraint = "ai_provider_keys_pkey" // ALTER TABLE ONLY ai_provider_keys ADD CONSTRAINT ai_provider_keys_pkey PRIMARY KEY (id);
|
||||
UniqueAIProvidersPkey UniqueConstraint = "ai_providers_pkey" // ALTER TABLE ONLY ai_providers ADD CONSTRAINT ai_providers_pkey PRIMARY KEY (id);
|
||||
UniqueAISeatStatePkey UniqueConstraint = "ai_seat_state_pkey" // ALTER TABLE ONLY ai_seat_state ADD CONSTRAINT ai_seat_state_pkey PRIMARY KEY (user_id);
|
||||
|
||||
@@ -123,7 +123,10 @@ func (api *API) upsertAIModelPrices(rw http.ResponseWriter, r *http.Request) {
|
||||
httpapi.InternalServerError(rw, err)
|
||||
return
|
||||
}
|
||||
err = api.Database.UpsertAIModelPrices(ctx, seed)
|
||||
err = api.Database.UpsertAIModelPrices(ctx, database.UpsertAIModelPricesParams{
|
||||
Seed: seed,
|
||||
Source: database.AIModelPriceSourceCustom,
|
||||
})
|
||||
if dbauthz.IsNotAuthorizedError(err) {
|
||||
httpapi.Forbidden(rw)
|
||||
return
|
||||
|
||||
Reference in New Issue
Block a user