diff --git a/coderd/aibridge/prices/prices.go b/coderd/aibridge/prices/prices.go index bbb5689ea02..5cac067c554 100644 --- a/coderd/aibridge/prices/prices.go +++ b/coderd/aibridge/prices/prices.go @@ -6,6 +6,7 @@ import ( "context" _ "embed" "encoding/json" + "sync" "golang.org/x/xerrors" @@ -60,3 +61,32 @@ func parseSeed(data []byte) ([]seedRow, error) { } return rows, nil } + +// modelKey identifies a priced model. +type modelKey struct { + provider string + model string +} + +// defaultPricedModels indexes the embedded price book by provider and model. +// Built on first use, since a deployment that never sets a price never needs +// it. +var defaultPricedModels = sync.OnceValue(func() map[modelKey]struct{} { + rows, err := parseSeed(seedJSON) + if err != nil { + panic(xerrors.Errorf("parse embedded price seed: %w", err)) + } + index := make(map[modelKey]struct{}, len(rows)) + for _, row := range rows { + index[modelKey{provider: row.Provider, model: row.Model}] = struct{}{} + } + return index +}) + +// IsDefaultPriced reports whether the embedded price book already carries a +// price for the model. Coder owns those prices and re-applies them on every +// startup, so an operator price set for one would not survive a restart. +func IsDefaultPriced(provider, model string) bool { + _, ok := defaultPricedModels()[modelKey{provider: provider, model: model}] + return ok +} diff --git a/coderd/aibridge/prices/prices_test.go b/coderd/aibridge/prices/prices_test.go index ae30b7e449e..4ecdabd2caa 100644 --- a/coderd/aibridge/prices/prices_test.go +++ b/coderd/aibridge/prices/prices_test.go @@ -289,3 +289,56 @@ func TestSeed(t *testing.T) { db, _ := dbtestutil.NewDB(t) require.NoError(t, prices.Seed(ctx, db)) } + +// TestIsDefaultPriced reads the real embedded price book, so it uses a model +// the generator injects rather than one that could drift out of upstream. +func TestIsDefaultPriced(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + provider string + model string + want bool + }{ + { + name: "ModelInThePriceBook", + provider: "anthropic", + model: "claude-opus-5", + want: true, + }, + { + name: "ModelNotInThePriceBook", + provider: "anthropic", + model: "not-a-real-model", + want: false, + }, + { + // The book is keyed on both columns, so the same model under + // another provider is a different entry. + name: "SameModelUnderAnotherProvider", + provider: "openai", + model: "claude-opus-5", + want: false, + }, + { + name: "UnknownProvider", + provider: "unknown-provider", + model: "claude-opus-5", + want: false, + }, + { + name: "Empty", + provider: "", + model: "", + want: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + require.Equal(t, tt.want, prices.IsDefaultPriced(tt.provider, tt.model)) + }) + } +} diff --git a/coderd/aibridge/prices/providers/providers.go b/coderd/aibridge/prices/providers/providers.go new file mode 100644 index 00000000000..62d09d6a5fc --- /dev/null +++ b/coderd/aibridge/prices/providers/providers.go @@ -0,0 +1,28 @@ +package providers + +import "github.com/coder/coder/v2/coderd/database" + +// Supported lists the providers a model price may be set for. +// +// openai-compat is excluded: it is a generic passthrough, so the upstream +// vendor is unknown and a price cannot be attributed to it. Listed explicitly +// rather than derived from ai_provider_type so a new provider is opt-in. +var Supported = []database.AIProviderType{ + database.AIProviderTypeAnthropic, + database.AIProviderTypeAzure, + database.AIProviderTypeBedrock, + database.AIProviderTypeCopilot, + database.AIProviderTypeGoogle, + database.AIProviderTypeOpenai, + database.AIProviderTypeOpenrouter, + database.AIProviderTypeVercel, +} + +// SupportedStrings returns the supported providers as plain strings. +func SupportedStrings() []string { + ids := make([]string, len(Supported)) + for i, provider := range Supported { + ids[i] = string(provider) + } + return ids +} diff --git a/coderd/apidoc/docs.go b/coderd/apidoc/docs.go index f3dd16d01d0..6660dcff756 100644 --- a/coderd/apidoc/docs.go +++ b/coderd/apidoc/docs.go @@ -64,6 +64,85 @@ const docTemplate = `{ } } }, + "/api/experimental/ai/model-prices": { + "get": { + "produces": [ + "application/json" + ], + "tags": [ + "Enterprise" + ], + "summary": "List AI model prices", + "operationId": "list-ai-model-prices", + "parameters": [ + { + "type": "string", + "description": "Only return prices for this provider", + "name": "provider", + "in": "query" + }, + { + "type": "string", + "description": "Only return prices for this model", + "name": "model", + "in": "query" + } + ], + "responses": { + "200": { + "description": "OK", + "schema": { + "type": "array", + "items": { + "$ref": "#/definitions/codersdk.AIModelPrice" + } + } + } + }, + "security": [ + { + "CoderSessionToken": [] + } + ], + "x-apidocgen": { + "skip": true + } + }, + "post": { + "consumes": [ + "application/json" + ], + "tags": [ + "Enterprise" + ], + "summary": "Upsert AI model prices", + "operationId": "upsert-ai-model-prices", + "parameters": [ + { + "description": "Prices to set", + "name": "request", + "in": "body", + "required": true, + "schema": { + "$ref": "#/definitions/codersdk.UpsertAIModelPricesRequest" + } + } + ], + "responses": { + "204": { + "description": "No Content" + } + }, + "security": [ + { + "CoderSessionToken": [] + } + ], + "x-apidocgen": { + "skip": true + } + } + }, "/api/experimental/chats": { "get": { "description": "Experimental: this endpoint is subject to change.", @@ -15939,6 +16018,60 @@ const docTemplate = `{ } } }, + "codersdk.AIModelPrice": { + "type": "object", + "properties": { + "cache_read_price": { + "type": "integer" + }, + "cache_write_price": { + "type": "integer" + }, + "created_at": { + "type": "string", + "format": "date-time" + }, + "input_price": { + "type": "integer" + }, + "model": { + "type": "string" + }, + "output_price": { + "type": "integer" + }, + "provider": { + "type": "string" + }, + "updated_at": { + "type": "string", + "format": "date-time" + } + } + }, + "codersdk.AIModelPriceUpsert": { + "type": "object", + "properties": { + "cache_read_price": { + "type": "integer" + }, + "cache_write_price": { + "type": "integer" + }, + "input_price": { + "type": "integer" + }, + "model": { + "type": "string" + }, + "output_price": { + "type": "integer" + }, + "provider": { + "type": "string" + } + } + }, "codersdk.AIProvider": { "type": "object", "properties": { @@ -26475,6 +26608,17 @@ const docTemplate = `{ } } }, + "codersdk.UpsertAIModelPricesRequest": { + "type": "object", + "properties": { + "prices": { + "type": "array", + "items": { + "$ref": "#/definitions/codersdk.AIModelPriceUpsert" + } + } + } + }, "codersdk.UpsertGroupAIBudgetRequest": { "type": "object", "properties": { diff --git a/coderd/apidoc/swagger.json b/coderd/apidoc/swagger.json index e0aa90bda21..1b8d8b86e96 100644 --- a/coderd/apidoc/swagger.json +++ b/coderd/apidoc/swagger.json @@ -49,6 +49,77 @@ } } }, + "/api/experimental/ai/model-prices": { + "get": { + "produces": ["application/json"], + "tags": ["Enterprise"], + "summary": "List AI model prices", + "operationId": "list-ai-model-prices", + "parameters": [ + { + "type": "string", + "description": "Only return prices for this provider", + "name": "provider", + "in": "query" + }, + { + "type": "string", + "description": "Only return prices for this model", + "name": "model", + "in": "query" + } + ], + "responses": { + "200": { + "description": "OK", + "schema": { + "type": "array", + "items": { + "$ref": "#/definitions/codersdk.AIModelPrice" + } + } + } + }, + "security": [ + { + "CoderSessionToken": [] + } + ], + "x-apidocgen": { + "skip": true + } + }, + "post": { + "consumes": ["application/json"], + "tags": ["Enterprise"], + "summary": "Upsert AI model prices", + "operationId": "upsert-ai-model-prices", + "parameters": [ + { + "description": "Prices to set", + "name": "request", + "in": "body", + "required": true, + "schema": { + "$ref": "#/definitions/codersdk.UpsertAIModelPricesRequest" + } + } + ], + "responses": { + "204": { + "description": "No Content" + } + }, + "security": [ + { + "CoderSessionToken": [] + } + ], + "x-apidocgen": { + "skip": true + } + } + }, "/api/experimental/chats": { "get": { "description": "Experimental: this endpoint is subject to change.", @@ -14226,6 +14297,60 @@ } } }, + "codersdk.AIModelPrice": { + "type": "object", + "properties": { + "cache_read_price": { + "type": "integer" + }, + "cache_write_price": { + "type": "integer" + }, + "created_at": { + "type": "string", + "format": "date-time" + }, + "input_price": { + "type": "integer" + }, + "model": { + "type": "string" + }, + "output_price": { + "type": "integer" + }, + "provider": { + "type": "string" + }, + "updated_at": { + "type": "string", + "format": "date-time" + } + } + }, + "codersdk.AIModelPriceUpsert": { + "type": "object", + "properties": { + "cache_read_price": { + "type": "integer" + }, + "cache_write_price": { + "type": "integer" + }, + "input_price": { + "type": "integer" + }, + "model": { + "type": "string" + }, + "output_price": { + "type": "integer" + }, + "provider": { + "type": "string" + } + } + }, "codersdk.AIProvider": { "type": "object", "properties": { @@ -24351,6 +24476,17 @@ } } }, + "codersdk.UpsertAIModelPricesRequest": { + "type": "object", + "properties": { + "prices": { + "type": "array", + "items": { + "$ref": "#/definitions/codersdk.AIModelPriceUpsert" + } + } + } + }, "codersdk.UpsertGroupAIBudgetRequest": { "type": "object", "properties": { diff --git a/coderd/database/db2sdk/db2sdk.go b/coderd/database/db2sdk/db2sdk.go index 67bff9308ad..943c2bb955a 100644 --- a/coderd/database/db2sdk/db2sdk.go +++ b/coderd/database/db2sdk/db2sdk.go @@ -1725,6 +1725,27 @@ func chatMessageParts(m database.ChatMessage) ([]codersdk.ChatMessagePart, error return filtered, nil } +func AIModelPrices(dbPrices []database.AIModelPrice) []codersdk.AIModelPrice { + out := make([]codersdk.AIModelPrice, 0, len(dbPrices)) + for _, dbPrice := range dbPrices { + out = append(out, AIModelPrice(dbPrice)) + } + return out +} + +func AIModelPrice(dbPrice database.AIModelPrice) codersdk.AIModelPrice { + return codersdk.AIModelPrice{ + Provider: dbPrice.Provider, + Model: dbPrice.Model, + InputPrice: nullInt64Ptr(dbPrice.InputPrice), + OutputPrice: nullInt64Ptr(dbPrice.OutputPrice), + CacheReadPrice: nullInt64Ptr(dbPrice.CacheReadPrice), + CacheWritePrice: nullInt64Ptr(dbPrice.CacheWritePrice), + CreatedAt: dbPrice.CreatedAt, + UpdatedAt: dbPrice.UpdatedAt, + } +} + func nullUUIDPtr(v uuid.NullUUID) *uuid.UUID { if !v.Valid { return nil diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index ad1becaf35c..c6618192877 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/coderd/database/dbauthz/dbauthz.go @@ -2841,6 +2841,13 @@ func (q *querier) GetAIModelPriceByProviderModel(ctx context.Context, arg databa return q.db.GetAIModelPriceByProviderModel(ctx, arg) } +func (q *querier) GetAIModelPrices(ctx context.Context, arg database.GetAIModelPricesParams) ([]database.AIModelPrice, error) { + if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceAiModelPrice); err != nil { + return nil, err + } + return q.db.GetAIModelPrices(ctx, arg) +} + func (q *querier) GetAIProviderByID(ctx context.Context, id uuid.UUID) (database.AIProvider, error) { if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceAIProvider); err != nil { return database.AIProvider{}, err diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go index f175c6b2166..86dc818f548 100644 --- a/coderd/database/dbauthz/dbauthz_test.go +++ b/coderd/database/dbauthz/dbauthz_test.go @@ -6849,6 +6849,11 @@ func (s *MethodTestSuite) TestAIBridge() { check.Args(database.GetAIModelPriceByProviderModelParams{}).Asserts(rbac.ResourceAiModelPrice, policy.ActionRead) })) + s.Run("GetAIModelPrices", s.Mocked(func(db *dbmock.MockStore, _ *gofakeit.Faker, check *expects) { + db.EXPECT().GetAIModelPrices(gomock.Any(), gomock.Any()).Return([]database.AIModelPrice{}, nil).AnyTimes() + check.Args(database.GetAIModelPricesParams{}).Asserts(rbac.ResourceAiModelPrice, policy.ActionRead) + })) + s.Run("GetOrganizationGroupsAISpend", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { org := testutil.Fake(s.T(), faker, database.Organization{}) row1 := testutil.Fake(s.T(), faker, database.GetOrganizationGroupsAISpendRow{OrganizationID: org.ID}) diff --git a/coderd/database/dbmetrics/querymetrics.go b/coderd/database/dbmetrics/querymetrics.go index 795e777ee83..4a986ad1eb1 100644 --- a/coderd/database/dbmetrics/querymetrics.go +++ b/coderd/database/dbmetrics/querymetrics.go @@ -1129,6 +1129,14 @@ func (m queryMetricsStore) GetAIModelPriceByProviderModel(ctx context.Context, a return r0, r1 } +func (m queryMetricsStore) GetAIModelPrices(ctx context.Context, arg database.GetAIModelPricesParams) ([]database.AIModelPrice, error) { + start := time.Now() + r0, r1 := m.s.GetAIModelPrices(ctx, arg) + m.queryLatencies.WithLabelValues("GetAIModelPrices").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetAIModelPrices").Inc() + return r0, r1 +} + func (m queryMetricsStore) GetAIProviderByID(ctx context.Context, id uuid.UUID) (database.AIProvider, error) { start := time.Now() r0, r1 := m.s.GetAIProviderByID(ctx, id) diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go index 418a4c205dd..4f09e5a4b7f 100644 --- a/coderd/database/dbmock/dbmock.go +++ b/coderd/database/dbmock/dbmock.go @@ -1965,6 +1965,21 @@ func (mr *MockStoreMockRecorder) GetAIModelPriceByProviderModel(ctx, arg any) *g return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAIModelPriceByProviderModel", reflect.TypeOf((*MockStore)(nil).GetAIModelPriceByProviderModel), ctx, arg) } +// GetAIModelPrices mocks base method. +func (m *MockStore) GetAIModelPrices(ctx context.Context, arg database.GetAIModelPricesParams) ([]database.AIModelPrice, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetAIModelPrices", ctx, arg) + ret0, _ := ret[0].([]database.AIModelPrice) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetAIModelPrices indicates an expected call of GetAIModelPrices. +func (mr *MockStoreMockRecorder) GetAIModelPrices(ctx, arg any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAIModelPrices", reflect.TypeOf((*MockStore)(nil).GetAIModelPrices), ctx, arg) +} + // GetAIProviderByID mocks base method. func (m *MockStore) GetAIProviderByID(ctx context.Context, id uuid.UUID) (database.AIProvider, error) { m.ctrl.T.Helper() diff --git a/coderd/database/querier.go b/coderd/database/querier.go index 34e74dedc4f..bbc90e02859 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -318,6 +318,7 @@ type sqlcQuerier interface { // so a returned row is itself proof the secret is valid. GetAIGatewayKeyByHashedSecret(ctx context.Context, hashedSecret []byte) (AIGatewayKey, error) GetAIModelPriceByProviderModel(ctx context.Context, arg GetAIModelPriceByProviderModelParams) (AIModelPrice, error) + 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 // transaction alone does not stop a concurrent soft-delete or disable diff --git a/coderd/database/querier_test.go b/coderd/database/querier_test.go index d416a3f6880..084e61602c8 100644 --- a/coderd/database/querier_test.go +++ b/coderd/database/querier_test.go @@ -18881,3 +18881,74 @@ func TestGetActiveUsersAuthorizationRolesParity(t *testing.T) { require.ElementsMatch(t, single.Groups, row.Groups, "groups diverged for user %s", row.ID) } } + +func TestGetAIModelPrices(t *testing.T) { + t.Parallel() + + // Two anthropic models, and an openai model sharing a name with one of + // them, so provider and model can be told apart. + const seed = `[ + {"provider":"anthropic","model":"model-a","input_price":1,"output_price":null,"cache_read_price":null,"cache_write_price":null}, + {"provider":"anthropic","model":"model-b","input_price":2,"output_price":null,"cache_read_price":null,"cache_write_price":null}, + {"provider":"openai","model":"model-a","input_price":3,"output_price":null,"cache_read_price":null,"cache_write_price":null} + ]` + + tests := []struct { + name string + params database.GetAIModelPricesParams + want []string + }{ + { + name: "NoFilterReturnsEveryPrice", + params: database.GetAIModelPricesParams{}, + want: []string{"anthropic/model-a", "anthropic/model-b", "openai/model-a"}, + }, + { + name: "ByProvider", + params: database.GetAIModelPricesParams{Provider: "anthropic"}, + want: []string{"anthropic/model-a", "anthropic/model-b"}, + }, + { + name: "ByModelSpansProviders", + params: database.GetAIModelPricesParams{Model: "model-a"}, + want: []string{"anthropic/model-a", "openai/model-a"}, + }, + { + name: "ByProviderAndModel", + params: database.GetAIModelPricesParams{Provider: "anthropic", Model: "model-a"}, + want: []string{"anthropic/model-a"}, + }, + { + name: "UnknownProviderMatchesNothing", + params: database.GetAIModelPricesParams{Provider: "unknown-provider"}, + want: nil, + }, + { + name: "MismatchedProviderAndModel", + params: database.GetAIModelPricesParams{Provider: "openai", Model: "model-b"}, + want: nil, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + ctx := testutil.Context(t, testutil.WaitShort) + db, _ := dbtestutil.NewDB(t) + require.NoError(t, db.UpsertAIModelPrices(ctx, []byte(seed))) + + prices, err := db.GetAIModelPrices(ctx, tt.params) + require.NoError(t, err) + + got := make([]string, 0, len(prices)) + for _, price := range prices { + got = append(got, price.Provider+"/"+price.Model) + } + if len(tt.want) == 0 { + require.Empty(t, got) + return + } + require.Equal(t, tt.want, got) + }) + } +} diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index 5ada09d52b0..bb57a04bb30 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -2888,6 +2888,61 @@ func (q *sqlQuerier) GetAIModelPriceByProviderModel(ctx context.Context, arg Get 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 +FROM ai_model_prices + -- Filter by provider +WHERE CASE + WHEN $1::text != '' THEN + provider = $1 + ELSE true + END + -- Filter by model + AND CASE + WHEN $2::text != '' THEN + model = $2 + ELSE true + END +ORDER BY provider, model +` + +type GetAIModelPricesParams struct { + Provider string `db:"provider" json:"provider"` + Model string `db:"model" json:"model"` +} + +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 { + return nil, err + } + defer rows.Close() + var items []AIModelPrice + for rows.Next() { + var i AIModelPrice + if err := rows.Scan( + &i.Provider, + &i.Model, + &i.InputPrice, + &i.OutputPrice, + &i.CacheReadPrice, + &i.CacheWritePrice, + &i.CreatedAt, + &i.UpdatedAt, + ); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + const getGroupAIBudget = `-- name: GetGroupAIBudget :one SELECT group_id, spend_limit_micros, created_at, updated_at FROM group_ai_budgets diff --git a/coderd/database/queries/aicostcontrol.sql b/coderd/database/queries/aicostcontrol.sql index 5552e8c7af4..def9c0d316e 100644 --- a/coderd/database/queries/aicostcontrol.sql +++ b/coderd/database/queries/aicostcontrol.sql @@ -39,6 +39,23 @@ SELECT * FROM ai_model_prices WHERE provider = @provider AND model = @model; +-- name: GetAIModelPrices :many +SELECT * +FROM ai_model_prices + -- Filter by provider +WHERE CASE + WHEN @provider::text != '' THEN + provider = @provider + ELSE true + END + -- Filter by model + AND CASE + WHEN @model::text != '' THEN + model = @model + ELSE true + END +ORDER BY provider, model; + -- name: GetGroupAIBudget :one SELECT * FROM group_ai_budgets diff --git a/codersdk/aimodelprices.go b/codersdk/aimodelprices.go new file mode 100644 index 00000000000..00cae36d186 --- /dev/null +++ b/codersdk/aimodelprices.go @@ -0,0 +1,96 @@ +package codersdk + +import ( + "context" + "net/http" + "time" +) + +// AIModelPrice is a per-model token price used by AI Gateway to compute the +// cost of an interception. +// +// Prices are integer micro-units per million tokens, so 1000000 is $1.00 per +// million tokens. A nil price means the price is not known, which the cost +// calculation treats the same as zero. Distinguish that from an explicit 0, +// which declares the model free of charge. +type AIModelPrice 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"` + CreatedAt time.Time `json:"created_at" format:"date-time"` + UpdatedAt time.Time `json:"updated_at" format:"date-time"` +} + +// MaxAIModelPricesBytes bounds an upsert request body. +const MaxAIModelPricesBytes = 1 << 20 // 1 MiB + +// UpsertAIModelPricesRequest sets prices for the listed models. Models absent +// from the request are left untouched. +type UpsertAIModelPricesRequest struct { + Prices []AIModelPriceUpsert `json:"prices"` +} + +// AIModelPriceUpsert is one model's prices in an upsert request. It carries +// only the writable fields of AIModelPrice. +type AIModelPriceUpsert 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"` +} + +// AIModelPricesFilter narrows the listed model prices. An empty field does not +// filter on that attribute. +// +// @typescript-ignore AIModelPricesFilter +type AIModelPricesFilter struct { + Provider string `json:"provider,omitempty"` + Model string `json:"model,omitempty"` +} + +func (f AIModelPricesFilter) asRequestOption() RequestOption { + return func(r *http.Request) { + query := r.URL.Query() + if f.Provider != "" { + query.Set("provider", f.Provider) + } + if f.Model != "" { + query.Set("model", f.Model) + } + r.URL.RawQuery = query.Encode() + } +} + +// ListAIModelPrices returns the AI model prices matching the filter. +func (c *ExperimentalClient) ListAIModelPrices(ctx context.Context, filter AIModelPricesFilter) ([]AIModelPrice, error) { + res, err := c.Request(ctx, http.MethodGet, "/api/experimental/ai/model-prices", nil, filter.asRequestOption()) + if err != nil { + return nil, err + } + defer res.Body.Close() + if res.StatusCode != http.StatusOK { + return nil, ReadBodyAsError(res) + } + + var prices []AIModelPrice + return prices, ReadBodyAsJSON(res, &prices) +} + +// UpsertAIModelPrices sets prices for the models in req. The request is +// rejected in full if any model fails validation. +func (c *ExperimentalClient) UpsertAIModelPrices(ctx context.Context, req UpsertAIModelPricesRequest) error { + res, err := c.Request(ctx, http.MethodPost, "/api/experimental/ai/model-prices", req) + if err != nil { + return err + } + defer res.Body.Close() + if res.StatusCode != http.StatusNoContent { + return ReadBodyAsError(res) + } + return nil +} diff --git a/docs/ai-coder/ai-gateway/cost-controls.md b/docs/ai-coder/ai-gateway/cost-controls.md index 52c4fa13c44..dba6be0e631 100644 --- a/docs/ai-coder/ai-gateway/cost-controls.md +++ b/docs/ai-coder/ai-gateway/cost-controls.md @@ -212,7 +212,46 @@ Replace `` with your Coder minor version, for example `2.36`. Monitor `coder_ai_gateway_cost_control_unpriced_token_usage_records_total`, labeled by `provider` and `model`, to detect unpriced usage. Any non-zero value means spend is under-counted. Because the price book ships with the release, a -newly launched model can remain unpriced until you upgrade Coder. +newly launched model is unpriced until you upgrade Coder or set a price for it +yourself. + +### Set model prices + +Use the experimental `coder exp ai-model-prices` command to set prices for +models the price book does not cover. It requires the AI Governance add-on and +the `ai_model_price:update` permission. Run +`coder exp ai-model-prices --help` for the full reference. + +List the prices this deployment holds, optionally narrowed to one provider or +model: + +```sh +coder exp ai-model-prices list --provider anthropic +``` + +Price a single model. Prices are micro-units per million tokens, so `3000000` +is $3.00 per million tokens. Use `null` for a price you do not have, and `0` to +declare a model free: + +```sh +coder exp ai-model-prices update --provider anthropic --model my-model \ + --input-price 3000000 --output-price 15000000 \ + --cache-read-price null --cache-write-price null +``` + +Price several models at once from a JSON document in the same shape as the +price book: + +```sh +coder exp ai-model-prices update prices.json +``` + +> [!IMPORTANT] +> +> - Prices are not retroactive. Usage recorded before you set a price stays +> unpriced, so past spend does not change. +> - You can only set prices for models the price book does not cover. +> - This command is experimental and can change without notice. ## Monitor spend diff --git a/docs/reference/api/schemas.md b/docs/reference/api/schemas.md index 1e6dbe6e0ef..527d79f185e 100644 --- a/docs/reference/api/schemas.md +++ b/docs/reference/api/schemas.md @@ -1138,6 +1138,58 @@ title: Schemas | `last_heartbeat_at` | string | false | | | | `name` | string | false | | | +## codersdk.AIModelPrice + +```json +{ + "cache_read_price": 0, + "cache_write_price": 0, + "created_at": "2019-08-24T14:15:22Z", + "input_price": 0, + "model": "string", + "output_price": 0, + "provider": "string", + "updated_at": "2019-08-24T14:15:22Z" +} +``` + +### Properties + +| Name | Type | Required | Restrictions | Description | +|---------------------|---------|----------|--------------|-------------| +| `cache_read_price` | integer | false | | | +| `cache_write_price` | integer | false | | | +| `created_at` | string | false | | | +| `input_price` | integer | false | | | +| `model` | string | false | | | +| `output_price` | integer | false | | | +| `provider` | string | false | | | +| `updated_at` | string | false | | | + +## codersdk.AIModelPriceUpsert + +```json +{ + "cache_read_price": 0, + "cache_write_price": 0, + "input_price": 0, + "model": "string", + "output_price": 0, + "provider": "string" +} +``` + +### Properties + +| Name | Type | Required | Restrictions | Description | +|---------------------|---------|----------|--------------|-------------| +| `cache_read_price` | integer | false | | | +| `cache_write_price` | integer | false | | | +| `input_price` | integer | false | | | +| `model` | string | false | | | +| `output_price` | integer | false | | | +| `provider` | string | false | | | + ## codersdk.AIProvider ```json @@ -14464,6 +14516,29 @@ If the schedule is empty, the user will be updated to use the default schedule.| |--------|--------|----------|--------------|-------------| | `hash` | string | false | | | +## codersdk.UpsertAIModelPricesRequest + +```json +{ + "prices": [ + { + "cache_read_price": 0, + "cache_write_price": 0, + "input_price": 0, + "model": "string", + "output_price": 0, + "provider": "string" + } + ] +} +``` + +### Properties + +| Name | Type | Required | Restrictions | Description | +|----------|---------------------------------------------------------------------|----------|--------------|-------------| +| `prices` | array of [codersdk.AIModelPriceUpsert](#codersdkaimodelpriceupsert) | false | | | + ## codersdk.UpsertGroupAIBudgetRequest ```json diff --git a/enterprise/cli/exp_aimodelprices.go b/enterprise/cli/exp_aimodelprices.go new file mode 100644 index 00000000000..c3b39734e87 --- /dev/null +++ b/enterprise/cli/exp_aimodelprices.go @@ -0,0 +1,521 @@ +package cli + +import ( + "encoding/json" + "fmt" + "io" + "os" + "strconv" + "strings" + + "github.com/dustin/go-humanize" + "github.com/mattn/go-isatty" + "golang.org/x/xerrors" + + agplcli "github.com/coder/coder/v2/cli" + "github.com/coder/coder/v2/cli/cliui" + "github.com/coder/coder/v2/codersdk" + "github.com/coder/serpent" +) + +func (r *RootCmd) aiModelPricesCommand() *serpent.Command { + return &serpent.Command{ + Use: "ai-model-prices", + Short: "Manage AI Governance model prices", + Handler: func(inv *serpent.Invocation) error { + return inv.Command.HelpHandler(inv) + }, + Children: []*serpent.Command{ + r.aiModelPricesList(), + r.aiModelPricesUpdate(), + }, + } +} + +const modelPricesUpdateDescriptionLong = `Sets prices for models that Coder's price book does not cover. Models the +price book covers cannot be changed. + +The JSON document is an array of model prices, in the same shape as +Coder's price book: + [ + { + "provider": "anthropic", + "model": "my-model", + "input_price": 3000000, + "output_price": 15000000, + "cache_read_price": 300000, + "cache_write_price": null + } + ] + * Prices are micro-units per million tokens, so 3000000 is $3.00 per + million tokens. + * A 'null' price is unknown and adds no cost. An explicit 0 declares the + model free. + * Every entry sets all four prices, so all four are required. Use 'null' + for a price you do not have. +` + +// aiModelPriceRow renders prices as dollars per million tokens for the table +// output. The embedded price carries the JSON output, keeping the raw +// micro-units for scripting. +type aiModelPriceRow struct { + // For JSON format: + codersdk.AIModelPrice `table:"-"` + + // For table format: + Provider string `json:"-" table:"provider,default_sort"` + Model string `json:"-" table:"model"` + InputPrice string `json:"-" table:"input price"` + OutputPrice string `json:"-" table:"output price"` + CacheReadPrice string `json:"-" table:"cache read price"` + CacheWritePrice string `json:"-" table:"cache write price"` + CreatedAt string `json:"-" table:"created at"` + UpdatedAt string `json:"-" table:"updated at"` +} + +func (r *RootCmd) aiModelPricesList() *serpent.Command { + var ( + provider string + model string + formatter = cliui.NewOutputFormatter( + cliui.TableFormat([]aiModelPriceRow{}, []string{ + "provider", "model", "input price", "output price", "cache read price", "cache write price", + }), + cliui.JSONFormat(), + ) + ) + + cmd := &serpent.Command{ + Use: "list", + Short: "List AI Governance model prices", + Long: "Lists every model priced for this deployment. Prices are shown in " + + "dollars per million tokens. Narrow the output with --provider or --model.", + Middleware: serpent.Chain(serpent.RequireNArgs(0)), + Options: serpent.OptionSet{ + { + Flag: "provider", + Description: "Only show models for this provider.", + Value: serpent.StringOf(&provider), + }, + { + Flag: "model", + Description: "Only show this model.", + Value: serpent.StringOf(&model), + }, + }, + Handler: func(inv *serpent.Invocation) error { + ctx := inv.Context() + client, err := r.InitClient(inv) + if err != nil { + return err + } + + prices, err := codersdk.NewExperimentalClient(client).ListAIModelPrices(ctx, codersdk.AIModelPricesFilter{ + Provider: provider, + Model: model, + }) + if err != nil { + return xerrors.Errorf("list model prices: %w", err) + } + + rows := make([]aiModelPriceRow, 0, len(prices)) + for _, price := range prices { + rows = append(rows, aiModelPriceRow{ + AIModelPrice: price, + Provider: price.Provider, + Model: price.Model, + InputPrice: formatMicros(price.InputPrice), + OutputPrice: formatMicros(price.OutputPrice), + CacheReadPrice: formatMicros(price.CacheReadPrice), + CacheWritePrice: formatMicros(price.CacheWritePrice), + CreatedAt: humanize.Time(price.CreatedAt), + UpdatedAt: humanize.Time(price.UpdatedAt), + }) + } + + // JSON output keeps the empty array so scripts can parse it. + if len(rows) == 0 && formatter.FormatID() == "table" { + cliui.Infof(inv.Stderr, "No model prices found.") + return nil + } + + out, err := formatter.Format(ctx, rows) + if err != nil { + return xerrors.Errorf("format output: %w", err) + } + _, err = fmt.Fprintln(inv.Stdout, out) + return err + }, + } + formatter.AttachOptions(&cmd.Options) + return cmd +} + +func (r *RootCmd) aiModelPricesUpdate() *serpent.Command { + var ( + provider string + model string + input string + output string + cacheRead string + cacheWrite string + ) + + cmd := &serpent.Command{ + Use: "update [file|-]", + Short: "Set AI Governance model prices", + Long: modelPricesUpdateDescriptionLong + "\n" + agplcli.FormatExamples( + agplcli.Example{ + Description: "Set prices for several models from a JSON document", + Command: "coder exp ai-model-prices update prices.json", + }, + agplcli.Example{ + Description: "Read the document from stdin, applied without confirmation", + Command: "coder exp ai-model-prices update < prices.json", + }, + agplcli.Example{ + Description: "Set prices for a single model", + Command: "coder exp ai-model-prices update --provider anthropic --model my-model " + + "--input-price 3000000 --output-price 15000000 --cache-read-price 300000 --cache-write-price null", + }, + agplcli.Example{ + Description: "Set prices without confirmation", + Command: "coder exp ai-model-prices update prices.json --yes", + }, + ), + Middleware: serpent.Chain(serpent.RequireRangeArgs(0, 1)), + Options: serpent.OptionSet{ + { + Flag: "provider", + Description: "Provider of the model to price.", + Value: serpent.StringOf(&provider), + }, + { + Flag: "model", + Description: "Model to price. Requires --provider.", + Value: serpent.StringOf(&model), + }, + { + Flag: "input-price", + Description: "Input price in micro-units per million tokens, or 'null' if unknown.", + Value: serpent.StringOf(&input), + }, + { + Flag: "output-price", + Description: "Output price in micro-units per million tokens, or 'null' if unknown.", + Value: serpent.StringOf(&output), + }, + { + Flag: "cache-read-price", + Description: "Cache read price in micro-units per million tokens, or 'null' if unknown.", + Value: serpent.StringOf(&cacheRead), + }, + { + Flag: "cache-write-price", + Description: "Cache write price in micro-units per million tokens, or 'null' if unknown.", + Value: serpent.StringOf(&cacheWrite), + }, + cliui.SkipPromptOption(), + }, + Handler: func(inv *serpent.Invocation) error { + ctx := inv.Context() + client, err := r.InitClient(inv) + if err != nil { + return err + } + exp := codersdk.NewExperimentalClient(client) + + requested, fromStdin, err := readAIModelPrices(inv, provider, model) + if err != nil { + return err + } + + current, err := exp.ListAIModelPrices(ctx, codersdk.AIModelPricesFilter{}) + if err != nil { + return xerrors.Errorf("list model prices: %w", err) + } + additions, changes := diffAIModelPrices(requested, current) + if len(additions) == 0 && len(changes) == 0 { + _, err = fmt.Fprintln(inv.Stdout, "No changes to apply.") + return err + } + printAIModelPriceChanges(inv, additions, changes) + + // Avoid prompting when the document came from stdin (already drained). + if !fromStdin { + if _, err := cliui.Prompt(inv, cliui.PromptOptions{ + Text: "Apply?", + IsConfirm: true, + Default: cliui.ConfirmNo, + }); err != nil { + return err + } + } + + if err := exp.UpsertAIModelPrices(ctx, codersdk.UpsertAIModelPricesRequest{Prices: requested}); err != nil { + return xerrors.Errorf("update model prices: %w", err) + } + _, err = fmt.Fprintf(inv.Stdout, "Updated prices for %d model(s).\n", len(additions)+len(changes)) + return err + }, + } + return cmd +} + +var modelPriceFlags = []string{"input-price", "output-price", "cache-read-price", "cache-write-price"} + +// readAIModelPrices picks the input source and returns the requested prices, +// reporting whether they were read from stdin. The single-model flags win when +// any of them is set, a JSON document is read otherwise, and supplying both is +// rejected rather than guessing which one the caller meant. Field-level rules +// are enforced by the server. +func readAIModelPrices(inv *serpent.Invocation, provider, model string) ([]codersdk.AIModelPriceUpsert, bool, error) { + providedFlags := userSetPriceFlags(inv) + if provider == "" && model == "" && len(providedFlags) == 0 { + return readAIModelPricesDocument(inv) + } + if len(inv.Args) > 0 { + return nil, false, xerrors.New("pass either a JSON document or the single-model flags, not both") + } + if err := validateModelPriceFlags(provider, model, providedFlags); err != nil { + return nil, false, err + } + price, err := modelPriceFromFlags(inv, provider, model) + if err != nil { + return nil, false, err + } + return []codersdk.AIModelPriceUpsert{price}, false, nil +} + +// userSetPriceFlags lists the price flags the caller supplied. +func userSetPriceFlags(inv *serpent.Invocation) []string { + var provided []string + for _, name := range modelPriceFlags { + if opt := inv.Command.Options.ByFlag(name); opt != nil && opt.ValueSource != serpent.ValueSourceNone { + provided = append(provided, name) + } + } + return provided +} + +// validateModelPriceFlags checks the flag combination names one complete model. +func validateModelPriceFlags(provider, model string, providedFlags []string) error { + if provider == "" || model == "" { + return xerrors.New("--provider and --model are both required to price a single model") + } + // An entry sets all four prices, so every flag is required. Leaving one out + // would clear that price rather than preserve it. + if len(providedFlags) != len(modelPriceFlags) { + return xerrors.Errorf("all price flags are required: --%s. Pass 'null' for a price you do not have", + strings.Join(modelPriceFlags, ", --")) + } + return nil +} + +// modelPriceFromFlags builds the entry named by the single-model flags. +func modelPriceFromFlags(inv *serpent.Invocation, provider, model string) (codersdk.AIModelPriceUpsert, error) { + price := codersdk.AIModelPriceUpsert{Provider: provider, Model: model} + for _, name := range modelPriceFlags { + // A nil value is an unknown price, spelled "null" to match the JSON + // document and the stored column. + raw := inv.Command.Options.ByFlag(name).Value.String() + var value *int64 + if !strings.EqualFold(raw, "null") { + parsed, err := strconv.ParseInt(raw, 10, 64) + if err != nil { + return codersdk.AIModelPriceUpsert{}, xerrors.Errorf("--%s: want a whole number or 'null', got %q", name, raw) + } + value = &parsed + } + switch name { + case "input-price": + price.InputPrice = value + case "output-price": + price.OutputPrice = value + case "cache-read-price": + price.CacheReadPrice = value + case "cache-write-price": + price.CacheWritePrice = value + } + } + return price, nil +} + +// isTerminalStdin reports whether stdin is an interactive terminal, meaning no +// document was piped or redirected in. +func isTerminalStdin(inv *serpent.Invocation) bool { + file, ok := inv.Stdin.(*os.File) + if !ok { + return false + } + return isatty.IsTerminal(file.Fd()) +} + +// readAIModelPricesDocument reads the JSON document from the named file, or +// from stdin when the argument is absent or "-", and reports which of the two +// it read. Draining stdin leaves nothing behind for a prompt to read. +func readAIModelPricesDocument(inv *serpent.Invocation) ([]codersdk.AIModelPriceUpsert, bool, error) { + var ( + data []byte + err error + fromStdin bool + ) + if len(inv.Args) == 0 || inv.Args[0] == "-" { + // Reading a terminal would block until the caller sends EOF, so leave + // the document empty rather than appearing to hang. + if !isTerminalStdin(inv) { + fromStdin = true + data, err = io.ReadAll(inv.Stdin) + if err != nil { + return nil, false, xerrors.Errorf("read stdin: %w", err) + } + } + } else { + data, err = os.ReadFile(inv.Args[0]) + if err != nil { + return nil, false, xerrors.Errorf("read %s: %w", inv.Args[0], err) + } + } + if len(data) == 0 { + return nil, false, xerrors.New("no prices given, pass a JSON document or set --provider, --model and the four price flags") + } + var requested []codersdk.AIModelPriceUpsert + if err := json.Unmarshal(data, &requested); err != nil { + return nil, false, xerrors.Errorf("parse prices: %w", err) + } + return requested, fromStdin, nil +} + +// aiModelPriceChange is a requested price alongside the one it replaces. +type aiModelPriceChange struct { + price codersdk.AIModelPriceUpsert + old codersdk.AIModelPrice +} + +// diffAIModelPrices splits the requested prices into models the deployment has +// no price for and models whose prices would move. Requests that match the +// stored prices exactly are dropped. +func diffAIModelPrices(requested []codersdk.AIModelPriceUpsert, current []codersdk.AIModelPrice) ([]codersdk.AIModelPriceUpsert, []aiModelPriceChange) { + stored := make(map[string]codersdk.AIModelPrice, len(current)) + for _, row := range current { + stored[row.Provider+"/"+row.Model] = row + } + + var ( + additions []codersdk.AIModelPriceUpsert + changes []aiModelPriceChange + ) + for _, entry := range requested { + old, ok := stored[entry.Provider+"/"+entry.Model] + if !ok { + additions = append(additions, entry) + continue + } + if pricesEqual(entry, old) { + continue + } + changes = append(changes, aiModelPriceChange{price: entry, old: old}) + } + return additions, changes +} + +// pricesEqual reports whether all four prices already hold the requested values. +func pricesEqual(requested codersdk.AIModelPriceUpsert, stored codersdk.AIModelPrice) bool { + return priceEqual(requested.InputPrice, stored.InputPrice) && + priceEqual(requested.OutputPrice, stored.OutputPrice) && + priceEqual(requested.CacheReadPrice, stored.CacheReadPrice) && + priceEqual(requested.CacheWritePrice, stored.CacheWritePrice) +} + +// priceEqual compares two prices, treating unknown as equal only to unknown. +func priceEqual(a, b *int64) bool { + if a == nil || b == nil { + return a == nil && b == nil + } + return *a == *b +} + +// printAIModelPriceChanges previews the requested changes, marking a model the +// deployment does not price yet with "+" and one whose prices change with "~". +// Nothing is written until the plan is applied. +func printAIModelPriceChanges(inv *serpent.Invocation, additions []codersdk.AIModelPriceUpsert, changes []aiModelPriceChange) { + var summary []string + if len(additions) > 0 { + summary = append(summary, fmt.Sprintf("%d to add", len(additions))) + } + if len(changes) > 0 { + summary = append(summary, fmt.Sprintf("%d to change", len(changes))) + } + cliui.Infof(inv.Stdout, "Plan: %s.", strings.Join(summary, ", ")) + + for _, price := range additions { + _, _ = fmt.Fprintf(inv.Stdout, " + %s/%s %s\n", price.Provider, price.Model, describePrices(price)) + } + for _, change := range changes { + _, _ = fmt.Fprintf(inv.Stdout, " ~ %s/%s\n", change.price.Provider, change.price.Model) + for _, line := range describePriceChanges(change) { + _, _ = fmt.Fprintf(inv.Stdout, " %s\n", line) + } + } +} + +// describePrices renders an entry as a one-line list of its set prices. +func describePrices(price codersdk.AIModelPriceUpsert) string { + var parts []string + for _, named := range namedPrices(price) { + if named.value == nil { + continue + } + parts = append(parts, fmt.Sprintf("%s %s", named.name, formatMicros(named.value))) + } + return strings.Join(parts, " ") +} + +// describePriceChanges renders one "old -> new" line per price that moves. +func describePriceChanges(change aiModelPriceChange) []string { + old := map[string]*int64{ + "input_price": change.old.InputPrice, + "output_price": change.old.OutputPrice, + "cache_read_price": change.old.CacheReadPrice, + "cache_write_price": change.old.CacheWritePrice, + } + var lines []string + for _, named := range namedPrices(change.price) { + if priceEqual(named.value, old[named.name]) { + continue + } + lines = append(lines, fmt.Sprintf("%-18s %s -> %s", named.name, formatMicros(old[named.name]), formatMicros(named.value))) + } + return lines +} + +// namedPrice pairs a price with the field name it is reported under. +type namedPrice struct { + name string + value *int64 +} + +// namedPrices lists an entry's four prices in field order. +func namedPrices(price codersdk.AIModelPriceUpsert) []namedPrice { + return []namedPrice{ + {"input_price", price.InputPrice}, + {"output_price", price.OutputPrice}, + {"cache_read_price", price.CacheReadPrice}, + {"cache_write_price", price.CacheWritePrice}, + } +} + +// formatMicros renders a micro-unit price as dollars per million tokens. An +// unknown price shows as "-", while a zero price shows as $0.00. A price under +// a cent carries enough decimals to stay distinct from another price that also +// rounds to $0.00. +func formatMicros(price *int64) string { + if price == nil { + return "-" + } + dollars := float64(*price) / 1_000_000 + if *price > 0 && *price < 10_000 { + return strings.TrimRight(fmt.Sprintf("$%.6f", dollars), "0") + } + return fmt.Sprintf("$%.2f", dollars) +} diff --git a/enterprise/cli/exp_aimodelprices_internal_test.go b/enterprise/cli/exp_aimodelprices_internal_test.go new file mode 100644 index 00000000000..2f8a7ec4e33 --- /dev/null +++ b/enterprise/cli/exp_aimodelprices_internal_test.go @@ -0,0 +1,131 @@ +package cli + +import ( + "testing" + + "github.com/stretchr/testify/require" + + "github.com/coder/coder/v2/codersdk" +) + +func TestFormatMicros(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + price *int64 + want string + }{ + {name: "Unknown", price: nil, want: "-"}, + {name: "Zero", price: new(int64(0)), want: "$0.00"}, + {name: "WholeDollars", price: new(int64(3_000_000)), want: "$3.00"}, + {name: "Fractional", price: new(int64(2_500_000)), want: "$2.50"}, + {name: "OneCent", price: new(int64(10_000)), want: "$0.01"}, + {name: "UnderACent", price: new(int64(3_600)), want: "$0.0036"}, + {name: "UnderACentTrailingZeros", price: new(int64(1_000)), want: "$0.001"}, + {name: "UnderACentManyDecimals", price: new(int64(3_625)), want: "$0.003625"}, + {name: "SmallestUnit", price: new(int64(1)), want: "$0.000001"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + require.Equal(t, tt.want, formatMicros(tt.price)) + }) + } +} + +func TestDiffAIModelPrices(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + current []codersdk.AIModelPrice + requested []codersdk.AIModelPriceUpsert + wantAdded int + wantChanged int + }{ + { + name: "UnknownModelIsAnAddition", + current: nil, + requested: []codersdk.AIModelPriceUpsert{{ + Provider: "anthropic", Model: "my-model", InputPrice: new(int64(100)), + }}, + wantAdded: 1, + wantChanged: 0, + }, + { + name: "ChangedPriceIsAChange", + current: []codersdk.AIModelPrice{{ + Provider: "anthropic", Model: "my-model", InputPrice: new(int64(100)), + }}, + requested: []codersdk.AIModelPriceUpsert{{ + Provider: "anthropic", Model: "my-model", InputPrice: new(int64(200)), + }}, + wantAdded: 0, + wantChanged: 1, + }, + { + name: "IdenticalPriceIsDropped", + current: []codersdk.AIModelPrice{{ + Provider: "anthropic", Model: "my-model", InputPrice: new(int64(100)), + }}, + requested: []codersdk.AIModelPriceUpsert{{ + Provider: "anthropic", Model: "my-model", InputPrice: new(int64(100)), + }}, + wantAdded: 0, + wantChanged: 0, + }, + { + name: "UnknownToValueIsAChange", + current: []codersdk.AIModelPrice{{ + Provider: "anthropic", Model: "my-model", InputPrice: new(int64(100)), + }}, + requested: []codersdk.AIModelPriceUpsert{{ + Provider: "anthropic", Model: "my-model", InputPrice: new(int64(100)), OutputPrice: new(int64(200)), + }}, + wantAdded: 0, + wantChanged: 1, + }, + { + name: "ValueToUnknownIsAChange", + current: []codersdk.AIModelPrice{{ + Provider: "anthropic", Model: "my-model", InputPrice: new(int64(100)), OutputPrice: new(int64(200)), + }}, + requested: []codersdk.AIModelPriceUpsert{{ + Provider: "anthropic", Model: "my-model", InputPrice: new(int64(100)), + }}, + wantAdded: 0, + wantChanged: 1, + }, + { + // A model of the same name under another provider is a different + // row, so it does not match. + name: "SameModelDifferentProvider", + current: []codersdk.AIModelPrice{{ + Provider: "openai", Model: "my-model", InputPrice: new(int64(100)), + }}, + requested: []codersdk.AIModelPriceUpsert{{ + Provider: "anthropic", Model: "my-model", InputPrice: new(int64(100)), + }}, + wantAdded: 1, + wantChanged: 0, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + additions, changes := diffAIModelPrices(tt.requested, tt.current) + require.Len(t, additions, tt.wantAdded) + require.Len(t, changes, tt.wantChanged) + + for i, change := range changes { + require.Equal(t, tt.requested[i], change.price) + require.Equal(t, tt.current[i].InputPrice, change.old.InputPrice) + require.Equal(t, tt.current[i].OutputPrice, change.old.OutputPrice) + } + }) + } +} diff --git a/enterprise/cli/exp_aimodelprices_test.go b/enterprise/cli/exp_aimodelprices_test.go new file mode 100644 index 00000000000..258c34eaf57 --- /dev/null +++ b/enterprise/cli/exp_aimodelprices_test.go @@ -0,0 +1,418 @@ +package cli_test + +import ( + "bytes" + "encoding/json" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/coder/coder/v2/cli/clitest" + "github.com/coder/coder/v2/codersdk" + "github.com/coder/coder/v2/enterprise/coderd/coderdenttest" + "github.com/coder/coder/v2/enterprise/coderd/license" + "github.com/coder/coder/v2/testutil" +) + +const aiModelPricesDocument = `[{ + "provider": "anthropic", + "model": "my-model", + "input_price": 100, + "output_price": 200, + "cache_read_price": null, + "cache_write_price": null +}]` + +// setupAIModelPricesCLI returns a client entitled to manage model prices. +func setupAIModelPricesCLI(t *testing.T) *codersdk.Client { + t.Helper() + + client, _ := coderdenttest.New(t, &coderdenttest.Options{ + LicenseOptions: &coderdenttest.LicenseOptions{ + Features: license.Features{ + codersdk.FeatureTemplateRBAC: 1, + codersdk.FeatureAIBridge: 1, + }, + }, + }) + return client +} + +func TestAIModelPricesUpdate(t *testing.T) { + t.Parallel() + + t.Run("RejectsInvalidInput", func(t *testing.T) { + t.Parallel() + + client := setupAIModelPricesCLI(t) + + tests := []struct { + name string + args []string + stdin string + wantErr string + }{ + { + name: "NoDocumentAndNoFlags", + args: []string{"exp", "ai-model-prices", "update"}, + wantErr: "no prices given, pass a JSON document or set --provider, --model and the four price flags", + }, + { + name: "ProviderWithoutModel", + args: []string{"exp", "ai-model-prices", "update", "--provider", "anthropic"}, + wantErr: "--provider and --model are both required", + }, + { + name: "ModelWithoutProvider", + args: []string{"exp", "ai-model-prices", "update", "--model", "my-model"}, + wantErr: "--provider and --model are both required", + }, + { + name: "PriceFlagWithoutProviderAndModel", + args: []string{"exp", "ai-model-prices", "update", "--input-price", "100"}, + wantErr: "--provider and --model are both required", + }, + { + name: "SomePriceFlags", + args: []string{ + "exp", "ai-model-prices", "update", + "--provider", "anthropic", "--model", "my-model", + "--input-price", "100", "--output-price", "200", + }, + wantErr: "all price flags are required", + }, + { + name: "NoPriceFlags", + args: []string{ + "exp", "ai-model-prices", "update", + "--provider", "anthropic", "--model", "my-model", + }, + wantErr: "all price flags are required", + }, + { + name: "NonNumericPrice", + args: []string{ + "exp", "ai-model-prices", "update", + "--provider", "anthropic", "--model", "my-model", + "--input-price", "abc", "--output-price", "null", + "--cache-read-price", "null", "--cache-write-price", "null", + }, + wantErr: `want a whole number or 'null', got "abc"`, + }, + { + name: "DocumentAndFlags", + args: []string{ + "exp", "ai-model-prices", "update", "prices.json", + "--provider", "anthropic", "--model", "my-model", + }, + wantErr: "not both", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + inv, conf := newCLI(t, tt.args...) + clitest.SetupConfig(t, client, conf) //nolint:gocritic // requires owner + inv.Stdin = strings.NewReader(tt.stdin) + + err := inv.Run() + require.Error(t, err) + require.Contains(t, err.Error(), tt.wantErr) + }) + } + }) + + t.Run("AppliesADocumentFromStdin", func(t *testing.T) { + t.Parallel() + + // Given: a licensed deployment and a document on stdin, with no --yes. + // Draining stdin leaves nothing for a prompt to read, so piping a + // document applies it without confirmation. + client := setupAIModelPricesCLI(t) + inv, conf := newCLI(t, "exp", "ai-model-prices", "update") + clitest.SetupConfig(t, client, conf) //nolint:gocritic // requires owner + + var stdout bytes.Buffer + inv.Stdin = strings.NewReader(aiModelPricesDocument) + inv.Stdout = &stdout + + // When: the command runs. + require.NoError(t, inv.Run()) + + // Then: the addition is planned and applied. + require.Contains(t, stdout.String(), "Plan: 1 to add.") + require.Contains(t, stdout.String(), "+ anthropic/my-model") + require.Contains(t, stdout.String(), "Updated prices for 1 model(s).") + + // Then: every column in the document is stored, nulls included. + ctx := testutil.Context(t, testutil.WaitLong) + prices, err := codersdk.NewExperimentalClient(client).ListAIModelPrices(ctx, + codersdk.AIModelPricesFilter{Provider: "anthropic", Model: "my-model"}) + require.NoError(t, err) + require.Len(t, prices, 1) + require.Equal(t, int64(100), *prices[0].InputPrice) + require.Equal(t, int64(200), *prices[0].OutputPrice) + require.Nil(t, prices[0].CacheReadPrice) + require.Nil(t, prices[0].CacheWritePrice) + }) + + t.Run("AppliesADocumentFromAFile", func(t *testing.T) { + t.Parallel() + + // Given: the document written to disk. + client := setupAIModelPricesCLI(t) + path := filepath.Join(t.TempDir(), "prices.json") + require.NoError(t, os.WriteFile(path, []byte(aiModelPricesDocument), 0o600)) + + inv, conf := newCLI(t, "exp", "ai-model-prices", "update", path, "--yes") + clitest.SetupConfig(t, client, conf) //nolint:gocritic // requires owner + + var stdout bytes.Buffer + inv.Stdout = &stdout + + // When: the command runs against the file. + require.NoError(t, inv.Run()) + + // Then: the price is applied. + require.Contains(t, stdout.String(), "Updated prices for 1 model(s).") + }) + + t.Run("AppliesASingleModelFromFlags", func(t *testing.T) { + t.Parallel() + + // Given: a licensed deployment. + client := setupAIModelPricesCLI(t) + inv, conf := newCLI(t, + "exp", "ai-model-prices", "update", + "--provider", "anthropic", "--model", "flag-model", + "--input-price", "100", "--output-price", "200", + "--cache-read-price", "300", "--cache-write-price", "null", + "--yes", + ) + clitest.SetupConfig(t, client, conf) //nolint:gocritic // requires owner + + var stdout bytes.Buffer + inv.Stdout = &stdout + + // When: the model is priced through the flags. + require.NoError(t, inv.Run()) + require.Contains(t, stdout.String(), "Updated prices for 1 model(s).") + + // Then: each flag lands in its own column, and the null flag is stored + // as unknown rather than zero. + ctx := testutil.Context(t, testutil.WaitLong) + prices, err := codersdk.NewExperimentalClient(client).ListAIModelPrices(ctx, + codersdk.AIModelPricesFilter{Provider: "anthropic", Model: "flag-model"}) + require.NoError(t, err) + require.Len(t, prices, 1) + require.Equal(t, int64(100), *prices[0].InputPrice) + require.Equal(t, int64(200), *prices[0].OutputPrice) + require.Equal(t, int64(300), *prices[0].CacheReadPrice) + require.Nil(t, prices[0].CacheWritePrice) + }) + + t.Run("PreviewsAChangedPrice", func(t *testing.T) { + t.Parallel() + + // Given: anthropic/change-model priced on all four columns. + client := setupAIModelPricesCLI(t) + exp := codersdk.NewExperimentalClient(client) + ctx := testutil.Context(t, testutil.WaitLong) + + //nolint:gocritic // Managing AI model prices is owner-only. + require.NoError(t, exp.UpsertAIModelPrices(ctx, + codersdk.UpsertAIModelPricesRequest{ + Prices: []codersdk.AIModelPriceUpsert{{ + Provider: "anthropic", Model: "change-model", + InputPrice: new(int64(3_000_000)), + OutputPrice: new(int64(15_000_000)), + CacheReadPrice: new(int64(300_000)), + CacheWritePrice: new(int64(1_000_000)), + }}, + })) + + inv, conf := newCLI(t, + "exp", "ai-model-prices", "update", + "--provider", "anthropic", "--model", "change-model", + "--input-price", "5000000", "--output-price", "16000000", + "--cache-read-price", "400000", "--cache-write-price", "null", + "--yes", + ) + clitest.SetupConfig(t, client, conf) //nolint:gocritic // requires owner + + var stdout bytes.Buffer + inv.Stdout = &stdout + + // When: every price is changed, including one cleared to unknown. + require.NoError(t, inv.Run()) + + // Then: the plan marks it as a change and shows each transition. + require.Contains(t, stdout.String(), "Plan: 1 to change.") + require.Contains(t, stdout.String(), "~ anthropic/change-model") + require.Contains(t, stdout.String(), "$3.00 -> $5.00") + require.Contains(t, stdout.String(), "$15.00 -> $16.00") + require.Contains(t, stdout.String(), "$0.30 -> $0.40") + require.Contains(t, stdout.String(), "$1.00 -> -") + + // Then: every column holds the new price, and the cleared one is unknown. + prices, err := exp.ListAIModelPrices(ctx, + codersdk.AIModelPricesFilter{Provider: "anthropic", Model: "change-model"}) + require.NoError(t, err) + require.Len(t, prices, 1) + require.Equal(t, int64(5_000_000), *prices[0].InputPrice) + require.Equal(t, int64(16_000_000), *prices[0].OutputPrice) + require.Equal(t, int64(400_000), *prices[0].CacheReadPrice) + require.Nil(t, prices[0].CacheWritePrice) + }) + + t.Run("ReportsNoChangesOnAReapply", func(t *testing.T) { + t.Parallel() + + client := setupAIModelPricesCLI(t) + apply := func() string { + inv, conf := newCLI(t, "exp", "ai-model-prices", "update") + clitest.SetupConfig(t, client, conf) //nolint:gocritic // requires owner + + var stdout bytes.Buffer + inv.Stdin = strings.NewReader(aiModelPricesDocument) + inv.Stdout = &stdout + require.NoError(t, inv.Run()) + return stdout.String() + } + + // Given: a document that has already been applied. + require.Contains(t, apply(), "Updated prices for 1 model(s).") + + // When: the same document is applied again. + // Then: the diff is empty, so nothing is written. + require.Contains(t, apply(), "No changes to apply.") + }) + + t.Run("RejectsAModelInThePriceBook", func(t *testing.T) { + t.Parallel() + + // Given: a model Coder already prices. + client := setupAIModelPricesCLI(t) + inv, conf := newCLI(t, + "exp", "ai-model-prices", "update", + "--provider", "anthropic", "--model", "claude-opus-5", + "--input-price", "100", "--output-price", "null", + "--cache-read-price", "null", "--cache-write-price", "null", + "--yes", + ) + clitest.SetupConfig(t, client, conf) //nolint:gocritic // requires owner + inv.Stdout = &bytes.Buffer{} + + // When: it is priced. Then: the server rejects it. + err := inv.Run() + require.Error(t, err) + require.Contains(t, err.Error(), "price book") + }) +} + +func TestAIModelPricesList(t *testing.T) { + t.Parallel() + + t.Run("JSONCarriesRawMicros", func(t *testing.T) { + t.Parallel() + + // Given: a model priced on three columns, with one left unknown. + client := setupAIModelPricesCLI(t) + ctx := testutil.Context(t, testutil.WaitLong) + + //nolint:gocritic // Managing AI model prices is owner-only. + require.NoError(t, codersdk.NewExperimentalClient(client).UpsertAIModelPrices(ctx, + codersdk.UpsertAIModelPricesRequest{ + Prices: []codersdk.AIModelPriceUpsert{{ + Provider: "anthropic", Model: "json-model", + InputPrice: new(int64(3_000_000)), + OutputPrice: new(int64(15_000_000)), + CacheReadPrice: new(int64(300_000)), + }}, + })) + + inv, conf := newCLI(t, "exp", "ai-model-prices", "list", + "--provider", "anthropic", "--model", "json-model", "--output", "json") + clitest.SetupConfig(t, client, conf) //nolint:gocritic // requires owner + + var stdout bytes.Buffer + inv.Stdout = &stdout + + // When: the prices are listed as JSON. + require.NoError(t, inv.Run()) + + // Then: every column comes back as raw micro-units, not the table's + // dollar strings, and the unknown one stays null. + var prices []codersdk.AIModelPrice + require.NoError(t, json.Unmarshal(stdout.Bytes(), &prices)) + require.Len(t, prices, 1) + require.Equal(t, "anthropic", prices[0].Provider) + require.Equal(t, "json-model", prices[0].Model) + require.Equal(t, int64(3_000_000), *prices[0].InputPrice) + require.Equal(t, int64(15_000_000), *prices[0].OutputPrice) + require.Equal(t, int64(300_000), *prices[0].CacheReadPrice) + require.Nil(t, prices[0].CacheWritePrice) + }) + + t.Run("TableRendersDollarsPerMillionTokens", func(t *testing.T) { + t.Parallel() + + // Given: a model priced on three columns, one of them under a cent, + // with the fourth left unknown. + client := setupAIModelPricesCLI(t) + ctx := testutil.Context(t, testutil.WaitLong) + + //nolint:gocritic // Managing AI model prices is owner-only. + require.NoError(t, codersdk.NewExperimentalClient(client).UpsertAIModelPrices(ctx, + codersdk.UpsertAIModelPricesRequest{ + Prices: []codersdk.AIModelPriceUpsert{{ + Provider: "anthropic", Model: "mymodel", + InputPrice: new(int64(3_000_000)), + OutputPrice: new(int64(15_000_000)), + CacheReadPrice: new(int64(3_600)), + }}, + })) + + inv, conf := newCLI(t, "exp", "ai-model-prices", "list", + "--provider", "anthropic", "--model", "mymodel") + clitest.SetupConfig(t, client, conf) //nolint:gocritic // requires owner + + var stdout bytes.Buffer + inv.Stdout = &stdout + + // When: the prices are listed as a table. + require.NoError(t, inv.Run()) + + // Then: each column is shown in dollars, a sub-cent price keeps enough + // decimals to stay distinct, and the unknown one shows as a dash. + require.Contains(t, stdout.String(), "mymodel") + require.Contains(t, stdout.String(), "$3.00") + require.Contains(t, stdout.String(), "$15.00") + require.Contains(t, stdout.String(), "$0.0036") + require.Contains(t, stdout.String(), "-") + }) + + t.Run("SaysWhenNothingMatches", func(t *testing.T) { + t.Parallel() + + // Given: a filter matching no model. + client := setupAIModelPricesCLI(t) + inv, conf := newCLI(t, "exp", "ai-model-prices", "list", "--model", "no-such-model") + clitest.SetupConfig(t, client, conf) //nolint:gocritic // requires owner + + var stderr bytes.Buffer + inv.Stdout = &bytes.Buffer{} + inv.Stderr = &stderr + + // When: the prices are listed. + require.NoError(t, inv.Run()) + + // Then: an empty table says so rather than printing nothing. + require.Contains(t, stderr.String(), "No model prices found.") + }) +} diff --git a/enterprise/cli/root.go b/enterprise/cli/root.go index f0813357226..877842a1222 100644 --- a/enterprise/cli/root.go +++ b/enterprise/cli/root.go @@ -31,8 +31,10 @@ func (r *RootCmd) enterpriseOnly() []*serpent.Command { } } -func (*RootCmd) enterpriseExperimental() []*serpent.Command { - return []*serpent.Command{} +func (r *RootCmd) enterpriseExperimental() []*serpent.Command { + return []*serpent.Command{ + r.aiModelPricesCommand(), + } } func (r *RootCmd) EnterpriseSubcommands() []*serpent.Command { diff --git a/enterprise/coderd/aimodelprices.go b/enterprise/coderd/aimodelprices.go new file mode 100644 index 00000000000..78f2d068469 --- /dev/null +++ b/enterprise/coderd/aimodelprices.go @@ -0,0 +1,275 @@ +package coderd + +import ( + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "slices" + "strings" + + "cdr.dev/slog/v3" + "github.com/coder/coder/v2/coderd/aibridge/prices" + "github.com/coder/coder/v2/coderd/aibridge/prices/providers" + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/database/db2sdk" + "github.com/coder/coder/v2/coderd/database/dbauthz" + "github.com/coder/coder/v2/coderd/httpapi" + "github.com/coder/coder/v2/coderd/httpmw" + "github.com/coder/coder/v2/codersdk" +) + +// EXPERIMENTAL: this endpoint is experimental and is subject to change. +// +// @Summary List AI model prices +// @ID list-ai-model-prices +// @Security CoderSessionToken +// @Produce json +// @Tags Enterprise +// @Param provider query string false "Only return prices for this provider" +// @Param model query string false "Only return prices for this model" +// @Success 200 {array} codersdk.AIModelPrice +// @Router /api/experimental/ai/model-prices [get] +// @x-apidocgen {"skip": true} +func (api *API) listAIModelPrices(rw http.ResponseWriter, r *http.Request) { + ctx := r.Context() + + dbPrices, err := api.Database.GetAIModelPrices(ctx, database.GetAIModelPricesParams{ + Provider: r.URL.Query().Get("provider"), + Model: r.URL.Query().Get("model"), + }) + if dbauthz.IsNotAuthorizedError(err) { + httpapi.Forbidden(rw) + return + } + if err != nil { + api.Logger.Error(ctx, "list ai model prices", slog.Error(err)) + httpapi.InternalServerError(rw, err) + return + } + httpapi.Write(ctx, rw, http.StatusOK, db2sdk.AIModelPrices(dbPrices)) +} + +// EXPERIMENTAL: this endpoint is experimental and is subject to change. +// +// @Summary Upsert AI model prices +// @ID upsert-ai-model-prices +// @Security CoderSessionToken +// @Accept json +// @Tags Enterprise +// @Param request body codersdk.UpsertAIModelPricesRequest true "Prices to set" +// @Success 204 +// @Router /api/experimental/ai/model-prices [post] +// @x-apidocgen {"skip": true} +func (api *API) upsertAIModelPrices(rw http.ResponseWriter, r *http.Request) { + ctx := r.Context() + + r.Body = http.MaxBytesReader(rw, r.Body, codersdk.MaxAIModelPricesBytes) + body, err := io.ReadAll(r.Body) + if err != nil { + if _, ok := errors.AsType[*http.MaxBytesError](err); ok { + httpapi.Write(ctx, rw, http.StatusRequestEntityTooLarge, codersdk.Response{ + Message: "Request body too large.", + Detail: err.Error(), + }) + return + } + httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{ + Message: "Failed to read request body.", + Detail: err.Error(), + }) + return + } + + // Decoding a price into *int64 loses whether its key was there at all, so + // decode the entries as raw keys too. An absent key is not the same as an + // explicit null. + var rawReq struct { + Prices []map[string]json.RawMessage `json:"prices"` + } + if err := json.Unmarshal(body, &rawReq); err != nil { + httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{ + Message: "Request body must be valid JSON.", + Detail: err.Error(), + }) + return + } + var req codersdk.UpsertAIModelPricesRequest + if err := json.Unmarshal(body, &req); err != nil { + httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{ + Message: "Request body must be valid JSON.", + Detail: err.Error(), + }) + return + } + + // Validate the whole request before writing anything, so a single bad + // entry cannot leave the table half-updated. + validations := validateAIModelPrices(req.Prices, rawReq.Prices) + if len(validations) > 0 { + httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{ + Message: "Invalid AI model prices.", + Validations: validations, + }) + return + } + + // The batch upsert reads the rows as a JSON array, matching the embedded + // price seed's wire format. + seed, err := json.Marshal(req.Prices) + if err != nil { + api.Logger.Error(ctx, "marshal ai model prices", slog.Error(err)) + httpapi.InternalServerError(rw, err) + return + } + err = api.Database.UpsertAIModelPrices(ctx, seed) + if dbauthz.IsNotAuthorizedError(err) { + httpapi.Forbidden(rw) + return + } + if err != nil { + api.Logger.Error(ctx, "upsert ai model prices", slog.Error(err), + slog.F("user_id", httpmw.APIKey(r).UserID), + slog.F("count", len(req.Prices)), + ) + httpapi.InternalServerError(rw, err) + return + } + + // Model prices feed cost reporting and budget enforcement, so record who + // changed how many. + // TODO(ssncferreira): replace with audit logging once ai_model_price is an + // auditable resource (AIGOV-590). + api.Logger.Info(ctx, "ai model prices updated", + slog.F("user_id", httpmw.APIKey(r).UserID), + slog.F("count", len(req.Prices)), + ) + + rw.WriteHeader(http.StatusNoContent) +} + +// modelKey identifies a priced model. +type modelKey struct { + provider string + model string +} + +// validateAIModelPrices reports every problem with the requested prices: a +// supported provider, a model Coder's price book does not already cover, all +// four price keys, non-negative prices with at least one set, and no repeated +// model. +func validateAIModelPrices(requested []codersdk.AIModelPriceUpsert, raw []map[string]json.RawMessage) []codersdk.ValidationError { + if len(requested) == 0 { + return []codersdk.ValidationError{{ + Field: "prices", + Detail: "At least one model price is required.", + }} + } + + supportedProviders := strings.Join(providers.SupportedStrings(), ", ") + seen := make(map[modelKey]struct{}, len(requested)) + var validations []codersdk.ValidationError + + for i, price := range requested { + field := fmt.Sprintf("prices[%d]", i) + + // Provider and model identify the row, so report them first. + switch { + case price.Provider == "": + validations = append(validations, codersdk.ValidationError{ + Field: field + ".provider", + Detail: fmt.Sprintf("Provider is required. Supported providers: %s.", supportedProviders), + }) + case !slices.Contains(providers.Supported, database.AIProviderType(price.Provider)): + validations = append(validations, codersdk.ValidationError{ + Field: field + ".provider", + Detail: fmt.Sprintf("Provider %q is not supported. Supported providers: %s.", price.Provider, supportedProviders), + }) + } + if price.Model == "" { + validations = append(validations, codersdk.ValidationError{ + Field: field + ".model", + Detail: "Model is required.", + }) + } + // The price book is re-applied on every server start, so a price set for + // a model it covers would not survive a restart. + // TODO(ssncferreira): drop this once custom pricing is supported + // (AIGOV-589). + if prices.IsDefaultPriced(price.Provider, price.Model) { + validations = append(validations, codersdk.ValidationError{ + Field: field, + Detail: fmt.Sprintf("%s/%s is priced by Coder's default price book. Overriding a default price is not supported.", price.Provider, price.Model), + }) + } + + named := []struct { + name string + value *int64 + }{ + {"input_price", price.InputPrice}, + {"output_price", price.OutputPrice}, + {"cache_read_price", price.CacheReadPrice}, + {"cache_write_price", price.CacheWritePrice}, + } + // raw and requested are decoded from the same body, so this index + // always exists. Bounds-check it anyway rather than risk a panic. + var rawEntry map[string]json.RawMessage + if i < len(raw) { + rawEntry = raw[i] + } + + var present, populated int + for _, p := range named { + if _, ok := rawEntry[p.name]; ok { + present++ + } + if p.value == nil { + continue + } + populated++ + if *p.value < 0 { + validations = append(validations, codersdk.ValidationError{ + Field: field + "." + p.name, + Detail: "Price must not be negative.", + }) + } + } + + // An entry sets all four columns, so an absent key would clear that + // price rather than leave it alone. + if present > 0 && present < len(named) { + for _, p := range named { + if _, ok := rawEntry[p.name]; ok { + continue + } + validations = append(validations, codersdk.ValidationError{ + Field: field + "." + p.name, + Detail: "Price is required. Use 'null' for a price that is not known.", + }) + } + } + + // An all-null entry creates a row that computes zero cost, so the model + // stops counting as unpriced without being priced. The price book skips + // such models rather than seeding a row for them. + if populated == 0 { + validations = append(validations, codersdk.ValidationError{ + Field: field, + Detail: "At least one price must be set. Use 0 to declare a model free of charge.", + }) + } + + key := modelKey{provider: price.Provider, model: price.Model} + if _, duplicate := seen[key]; duplicate { + validations = append(validations, codersdk.ValidationError{ + Field: field, + Detail: fmt.Sprintf("%s/%s appears more than once.", price.Provider, price.Model), + }) + } + seen[key] = struct{}{} + } + + return validations +} diff --git a/enterprise/coderd/aimodelprices_internal_test.go b/enterprise/coderd/aimodelprices_internal_test.go new file mode 100644 index 00000000000..75d09e5fb09 --- /dev/null +++ b/enterprise/coderd/aimodelprices_internal_test.go @@ -0,0 +1,194 @@ +package coderd + +import ( + "encoding/json" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/coder/coder/v2/codersdk" +) + +// decodeAIModelPrices decodes a prices document the same way the handler does, +// so a case can be written as the JSON an operator would send. +func decodeAIModelPrices(t *testing.T, body string) ([]codersdk.AIModelPriceUpsert, []map[string]json.RawMessage) { + t.Helper() + + var typed codersdk.UpsertAIModelPricesRequest + require.NoError(t, json.Unmarshal([]byte(body), &typed)) + + var raw struct { + Prices []map[string]json.RawMessage `json:"prices"` + } + require.NoError(t, json.Unmarshal([]byte(body), &raw)) + + return typed.Prices, raw.Prices +} + +func TestValidateAIModelPrices(t *testing.T) { + t.Parallel() + + const allPrices = `"input_price": 100, "output_price": 200, "cache_read_price": 300, "cache_write_price": 400` + + tests := []struct { + name string + body string + // want is every validation error, in order. Nil means the document is + // accepted. + want []codersdk.ValidationError + }{ + { + name: "EmptyRequest", + body: `{"prices":[]}`, + want: []codersdk.ValidationError{ + {Field: "prices", Detail: "At least one model price is required."}, + }, + }, + { + name: "MissingProvider", + body: `{"prices":[{"model":"my-model",` + allPrices + `}]}`, + want: []codersdk.ValidationError{ + {Field: "prices[0].provider", Detail: "Provider is required. Supported providers: anthropic, azure, bedrock, copilot, google, openai, openrouter, vercel."}, + }, + }, + { + name: "UnsupportedProvider", + body: `{"prices":[{"provider":"unknown-provider","model":"my-model",` + allPrices + `}]}`, + want: []codersdk.ValidationError{ + {Field: "prices[0].provider", Detail: `Provider "unknown-provider" is not supported. Supported providers: anthropic, azure, bedrock, copilot, google, openai, openrouter, vercel.`}, + }, + }, + { + // openai-compat is a generic passthrough, so a price cannot be + // attributed to the model behind it. + name: "OpenAICompatRejected", + body: `{"prices":[{"provider":"openai-compat","model":"my-model",` + allPrices + `}]}`, + want: []codersdk.ValidationError{ + {Field: "prices[0].provider", Detail: `Provider "openai-compat" is not supported. Supported providers: anthropic, azure, bedrock, copilot, google, openai, openrouter, vercel.`}, + }, + }, + { + name: "MissingModel", + body: `{"prices":[{"provider":"anthropic","model":"",` + allPrices + `}]}`, + want: []codersdk.ValidationError{ + {Field: "prices[0].model", Detail: "Model is required."}, + }, + }, + { + // The price book is re-applied on every restart, so this price + // would not survive one. + name: "ModelInPriceBook", + body: `{"prices":[{"provider":"anthropic","model":"claude-opus-5",` + allPrices + `}]}`, + want: []codersdk.ValidationError{ + {Field: "prices[0]", Detail: "anthropic/claude-opus-5 is priced by Coder's default price book. Overriding a default price is not supported."}, + }, + }, + { + name: "NegativePrice", + body: `{"prices":[{"provider":"anthropic","model":"my-model","input_price":-1,"output_price":200,"cache_read_price":null,"cache_write_price":null}]}`, + want: []codersdk.ValidationError{ + {Field: "prices[0].input_price", Detail: "Price must not be negative."}, + }, + }, + { + // Some keys present and some absent would clear the absent ones. + name: "MissingSomePriceKeys", + body: `{"prices":[{"provider":"anthropic","model":"my-model","input_price":100}]}`, + want: []codersdk.ValidationError{ + {Field: "prices[0].output_price", Detail: "Price is required. Use 'null' for a price that is not known."}, + {Field: "prices[0].cache_read_price", Detail: "Price is required. Use 'null' for a price that is not known."}, + {Field: "prices[0].cache_write_price", Detail: "Price is required. Use 'null' for a price that is not known."}, + }, + }, + { + // Reported once rather than four times, since no price was given at + // all. + name: "MissingAllPriceKeys", + body: `{"prices":[{"provider":"anthropic","model":"my-model"}]}`, + want: []codersdk.ValidationError{ + {Field: "prices[0]", Detail: "At least one price must be set. Use 0 to declare a model free of charge."}, + }, + }, + { + name: "AllPricesNull", + body: `{"prices":[{"provider":"anthropic","model":"my-model","input_price":null,"output_price":null,"cache_read_price":null,"cache_write_price":null}]}`, + want: []codersdk.ValidationError{ + {Field: "prices[0]", Detail: "At least one price must be set. Use 0 to declare a model free of charge."}, + }, + }, + { + name: "DuplicateEntry", + body: `{"prices":[{"provider":"anthropic","model":"my-model",` + allPrices + `},` + + `{"provider":"anthropic","model":"my-model",` + allPrices + `}]}`, + want: []codersdk.ValidationError{ + {Field: "prices[1]", Detail: "anthropic/my-model appears more than once."}, + }, + }, + { + // Model names may carry a "/", as openrouter IDs do. + name: "SeparatorInAModelNameIsAccepted", + body: `{"prices":[{"provider":"openrouter","model":"anthropic/my-model",` + allPrices + `}]}`, + want: nil, + }, + { + // The two entries share a provider/model concatenation, so keying + // on the pair is what keeps them apart. + name: "SeparatorDoesNotCollideWithAProviderName", + body: `{"prices":[{"provider":"openrouter","model":"anthropic/my-model",` + allPrices + `},` + + `{"provider":"openrouter/anthropic","model":"my-model",` + allPrices + `}]}`, + want: []codersdk.ValidationError{ + {Field: "prices[1].provider", Detail: `Provider "openrouter/anthropic" is not supported. Supported providers: anthropic, azure, bedrock, copilot, google, openai, openrouter, vercel.`}, + }, + }, + { + // A model repeated under a different provider is a different row. + name: "SameModelDifferentProvider", + body: `{"prices":[{"provider":"anthropic","model":"my-model",` + allPrices + `},` + + `{"provider":"openai","model":"my-model",` + allPrices + `}]}`, + want: nil, + }, + { + // Provider comes first, and the absent price keys are reported once. + name: "ReportsProviderBeforePrices", + body: `{"prices":[{"provider":"unknown-provider","model":"my-model"}]}`, + want: []codersdk.ValidationError{ + {Field: "prices[0].provider", Detail: `Provider "unknown-provider" is not supported. Supported providers: anthropic, azure, bedrock, copilot, google, openai, openrouter, vercel.`}, + {Field: "prices[0]", Detail: "At least one price must be set. Use 0 to declare a model free of charge."}, + }, + }, + { + // Every entry is reported, not just the first bad one. + name: "ReportsEveryEntry", + body: `{"prices":[{"provider":"unknown-provider","model":"a",` + allPrices + `},` + + `{"provider":"anthropic","model":"",` + allPrices + `}]}`, + want: []codersdk.ValidationError{ + {Field: "prices[0].provider", Detail: `Provider "unknown-provider" is not supported. Supported providers: anthropic, azure, bedrock, copilot, google, openai, openrouter, vercel.`}, + {Field: "prices[1].model", Detail: "Model is required."}, + }, + }, + { + name: "Valid", + body: `{"prices":[{"provider":"anthropic","model":"my-model",` + allPrices + `}]}`, + want: nil, + }, + { + name: "ValidWithNullPrices", + body: `{"prices":[{"provider":"anthropic","model":"my-model","input_price":100,"output_price":null,"cache_read_price":null,"cache_write_price":null}]}`, + want: nil, + }, + { + name: "ValidWithZeroPrice", + body: `{"prices":[{"provider":"anthropic","model":"my-model","input_price":0,"output_price":0,"cache_read_price":null,"cache_write_price":null}]}`, + want: nil, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + requested, raw := decodeAIModelPrices(t, tt.body) + require.Equal(t, tt.want, validateAIModelPrices(requested, raw)) + }) + } +} diff --git a/enterprise/coderd/aimodelprices_test.go b/enterprise/coderd/aimodelprices_test.go new file mode 100644 index 00000000000..0902d2a0601 --- /dev/null +++ b/enterprise/coderd/aimodelprices_test.go @@ -0,0 +1,416 @@ +package coderd_test + +import ( + "encoding/json" + "fmt" + "net/http" + "strings" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/coder/coder/v2/coderd/coderdtest" + "github.com/coder/coder/v2/coderd/rbac" + "github.com/coder/coder/v2/coderd/util/ptr" + "github.com/coder/coder/v2/codersdk" + "github.com/coder/coder/v2/enterprise/coderd/coderdenttest" + "github.com/coder/coder/v2/enterprise/coderd/license" + "github.com/coder/coder/v2/testutil" +) + +// setupAIModelPricesTest returns a client entitled to manage model prices. +func setupAIModelPricesTest(t *testing.T) (*codersdk.Client, codersdk.CreateFirstUserResponse) { + t.Helper() + + return coderdenttest.New(t, &coderdenttest.Options{ + LicenseOptions: &coderdenttest.LicenseOptions{ + Features: license.Features{ + codersdk.FeatureTemplateRBAC: 1, + codersdk.FeatureAIBridge: 1, + }, + }, + }) +} + +func newAIModelPrice(provider, model string, input int64) codersdk.AIModelPriceUpsert { + return codersdk.AIModelPriceUpsert{ + Provider: provider, + Model: model, + InputPrice: &input, + } +} + +func TestUpsertAIModelPrices(t *testing.T) { + t.Parallel() + + t.Run("LicenseEntitlement", func(t *testing.T) { + t.Parallel() + + // Given: a deployment without the AI Bridge feature. + ownerClient, _ := coderdenttest.New(t, &coderdenttest.Options{ + LicenseOptions: &coderdenttest.LicenseOptions{ + Features: license.Features{}, + }, + }) + ctx := testutil.Context(t, testutil.WaitLong) + + // When: an owner sets a price for anthropic/my-model. + //nolint:gocritic // Managing AI model prices is owner-only. + err := codersdk.NewExperimentalClient(ownerClient).UpsertAIModelPrices(ctx, codersdk.UpsertAIModelPricesRequest{ + Prices: []codersdk.AIModelPriceUpsert{newAIModelPrice("anthropic", "my-model", 100)}, + }) + + // Then: RequireFeatureMW rejects it as a Premium feature. + var sdkErr *codersdk.Error + require.ErrorAs(t, err, &sdkErr) + require.Equal(t, http.StatusForbidden, sdkErr.StatusCode()) + require.Contains(t, sdkErr.Message, "Premium feature") + }) + t.Run("Forbidden", func(t *testing.T) { + t.Parallel() + + // Given: a member without ai_model_price:update. + ownerClient, owner := setupAIModelPricesTest(t) + memberClient, _ := coderdtest.CreateAnotherUser(t, ownerClient, owner.OrganizationID) + ctx := testutil.Context(t, testutil.WaitLong) + + // When: they set a price for anthropic/my-model. + err := codersdk.NewExperimentalClient(memberClient).UpsertAIModelPrices(ctx, codersdk.UpsertAIModelPricesRequest{ + Prices: []codersdk.AIModelPriceUpsert{newAIModelPrice("anthropic", "my-model", 100)}, + }) + + // Then: the request is forbidden by ai_model_price:update. + var sdkErr *codersdk.Error + require.ErrorAs(t, err, &sdkErr) + require.Equal(t, http.StatusForbidden, sdkErr.StatusCode()) + }) + + t.Run("RejectsAnOversizedBody", func(t *testing.T) { + t.Parallel() + + // Given: an entitled deployment and a body over the size cap. + ownerClient, _ := setupAIModelPricesTest(t) + ctx := testutil.Context(t, testutil.WaitLong) + body := fmt.Sprintf(`{"prices":[{"provider":"anthropic","model":%q}]}`, + strings.Repeat("a", codersdk.MaxAIModelPricesBytes)) + + // When: the body is sent. + //nolint:gocritic // Managing AI model prices is owner-only. + res, err := ownerClient.Request(ctx, http.MethodPost, + "/api/experimental/ai/model-prices", json.RawMessage(body)) + require.NoError(t, err) + defer res.Body.Close() + + // Then: it is rejected before the body is decoded. + require.Equal(t, http.StatusRequestEntityTooLarge, res.StatusCode) + }) + + t.Run("RejectsMalformedBody", func(t *testing.T) { + t.Parallel() + + // Given: an entitled deployment. + ownerClient, _ := setupAIModelPricesTest(t) + + tests := []struct { + name string + body string + wantDetail string + }{ + { + name: "NotAnArray", + body: `{"prices": "not-an-array"}`, + wantDetail: "cannot unmarshal string", + }, + { + name: "WrongFieldType", + body: `{"prices":[{"provider":"anthropic","model":"my-model","input_price":"abc","output_price":null,"cache_read_price":null,"cache_write_price":null}]}`, + wantDetail: "input_price", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + ctx := testutil.Context(t, testutil.WaitLong) + + // When: the body is sent. json.RawMessage marshals verbatim, so + // it reaches the handler unchanged. + //nolint:gocritic // Managing AI model prices is owner-only. + res, err := ownerClient.Request(ctx, http.MethodPost, + "/api/experimental/ai/model-prices", json.RawMessage(tt.body)) + require.NoError(t, err) + defer res.Body.Close() + + // Then: the detail names the problem, since both decodes report + // the same message. + require.Equal(t, http.StatusBadRequest, res.StatusCode) + var sdkErr *codersdk.Error + require.ErrorAs(t, codersdk.ReadBodyAsError(res), &sdkErr) + require.Equal(t, "Request body must be valid JSON.", sdkErr.Message) + require.Contains(t, sdkErr.Detail, tt.wantDetail) + }) + } + }) + + t.Run("RejectsInvalidPrices", func(t *testing.T) { + t.Parallel() + + // Given: an entitled deployment. + ownerClient, _ := setupAIModelPricesTest(t) + exp := codersdk.NewExperimentalClient(ownerClient) + ctx := testutil.Context(t, testutil.WaitLong) + + // When: an entry carries an empty model name. + //nolint:gocritic // Managing AI model prices is owner-only. + err := exp.UpsertAIModelPrices(ctx, codersdk.UpsertAIModelPricesRequest{ + Prices: []codersdk.AIModelPriceUpsert{newAIModelPrice("anthropic", "", 100)}, + }) + + // Then: a 400 comes back naming the fields that failed. + var sdkErr *codersdk.Error + require.ErrorAs(t, err, &sdkErr) + require.Equal(t, http.StatusBadRequest, sdkErr.StatusCode()) + require.Equal(t, "Invalid AI model prices.", sdkErr.Message) + require.NotEmpty(t, sdkErr.Validations) + }) + + // A rejected request must leave the table untouched, so a large document + // cannot be applied halfway. + t.Run("WritesNothingWhenAnyEntryIsInvalid", func(t *testing.T) { + t.Parallel() + + // Given: the prices already stored. + ownerClient, _ := setupAIModelPricesTest(t) + exp := codersdk.NewExperimentalClient(ownerClient) + ctx := testutil.Context(t, testutil.WaitLong) + + //nolint:gocritic // Managing AI model prices is owner-only. + before, err := exp.ListAIModelPrices(ctx, codersdk.AIModelPricesFilter{}) + require.NoError(t, err) + + // When: a document holds a valid good-model ahead of an entry with an + // empty model name. + err = exp.UpsertAIModelPrices(ctx, codersdk.UpsertAIModelPricesRequest{ + Prices: []codersdk.AIModelPriceUpsert{ + newAIModelPrice("anthropic", "good-model", 100), + newAIModelPrice("anthropic", "", 100), + }, + }) + require.Error(t, err) + + // Then: good-model is not written either, so the row count is unchanged. + after, err := exp.ListAIModelPrices(ctx, codersdk.AIModelPricesFilter{}) + require.NoError(t, err) + require.Len(t, after, len(before), "no price should have been written") + }) + + t.Run("SetsPrices", func(t *testing.T) { + t.Parallel() + + // Given: anthropic/my-model, which the price book does not cover. + ownerClient, _ := setupAIModelPricesTest(t) + exp := codersdk.NewExperimentalClient(ownerClient) + ctx := testutil.Context(t, testutil.WaitLong) + + // When: it is priced with an input price only. + //nolint:gocritic // Managing AI model prices is owner-only. + require.NoError(t, exp.UpsertAIModelPrices(ctx, codersdk.UpsertAIModelPricesRequest{ + Prices: []codersdk.AIModelPriceUpsert{newAIModelPrice("anthropic", "my-model", 3_000_000)}, + })) + + // Then: the input price is stored and the other three are null. + prices, err := exp.ListAIModelPrices(ctx, codersdk.AIModelPricesFilter{ + Provider: "anthropic", + Model: "my-model", + }) + require.NoError(t, err) + require.Len(t, prices, 1) + require.Equal(t, int64(3_000_000), *prices[0].InputPrice) + require.Nil(t, prices[0].OutputPrice) + require.Nil(t, prices[0].CacheReadPrice) + require.Nil(t, prices[0].CacheWritePrice) + }) + + t.Run("UpdatesAPriceItSet", func(t *testing.T) { + t.Parallel() + + // Given: anthropic/my-model priced at 100 through this endpoint. + ownerClient, _ := setupAIModelPricesTest(t) + exp := codersdk.NewExperimentalClient(ownerClient) + ctx := testutil.Context(t, testutil.WaitLong) + + //nolint:gocritic // Managing AI model prices is owner-only. + require.NoError(t, exp.UpsertAIModelPrices(ctx, codersdk.UpsertAIModelPricesRequest{ + Prices: []codersdk.AIModelPriceUpsert{newAIModelPrice("anthropic", "my-model", 100)}, + })) + + // When: the same model is priced again at 200. + require.NoError(t, exp.UpsertAIModelPrices(ctx, codersdk.UpsertAIModelPricesRequest{ + Prices: []codersdk.AIModelPriceUpsert{newAIModelPrice("anthropic", "my-model", 200)}, + })) + + // Then: the row is updated to 200. + prices, err := exp.ListAIModelPrices(ctx, codersdk.AIModelPricesFilter{ + Provider: "anthropic", + Model: "my-model", + }) + require.NoError(t, err) + require.Len(t, prices, 1) + require.Equal(t, int64(200), *prices[0].InputPrice) + }) + + t.Run("StoresAModelNameWithASeparator", func(t *testing.T) { + t.Parallel() + + // Given: an openrouter model whose ID carries a "/". + ownerClient, _ := setupAIModelPricesTest(t) + exp := codersdk.NewExperimentalClient(ownerClient) + ctx := testutil.Context(t, testutil.WaitLong) + + // When: it is priced. + //nolint:gocritic // Managing AI model prices is owner-only. + require.NoError(t, exp.UpsertAIModelPrices(ctx, codersdk.UpsertAIModelPricesRequest{ + Prices: []codersdk.AIModelPriceUpsert{ + newAIModelPrice("openrouter", "anthropic/my-model", 100), + }, + })) + + // Then: the row is stored under the full model name, and filtering on + // it round-trips through the query parameter. + prices, err := exp.ListAIModelPrices(ctx, codersdk.AIModelPricesFilter{ + Provider: "openrouter", Model: "anthropic/my-model", + }) + require.NoError(t, err) + require.Len(t, prices, 1) + require.Equal(t, "openrouter", prices[0].Provider) + require.Equal(t, "anthropic/my-model", prices[0].Model) + require.Equal(t, int64(100), *prices[0].InputPrice) + }) +} + +func TestListAIModelPrices(t *testing.T) { + t.Parallel() + + t.Run("ReturnsPriceBook", func(t *testing.T) { + t.Parallel() + + // Given: a deployment seeded with the embedded price book at startup. + ownerClient, _ := setupAIModelPricesTest(t) + exp := codersdk.NewExperimentalClient(ownerClient) + ctx := testutil.Context(t, testutil.WaitLong) + + // When: the prices are listed. + //nolint:gocritic // Reading AI model prices is owner-only. + prices, err := exp.ListAIModelPrices(ctx, codersdk.AIModelPricesFilter{}) + require.NoError(t, err) + require.NotEmpty(t, prices, "the embedded price book is seeded at startup") + + // Then: a model the book covers comes back with all four prices. + seeded, err := exp.ListAIModelPrices(ctx, codersdk.AIModelPricesFilter{ + Provider: "anthropic", Model: "claude-opus-5", + }) + require.NoError(t, err) + require.Len(t, seeded, 1) + require.Equal(t, int64(5_000_000), *seeded[0].InputPrice) + require.Equal(t, int64(25_000_000), *seeded[0].OutputPrice) + require.Equal(t, int64(500_000), *seeded[0].CacheReadPrice) + require.Equal(t, int64(6_250_000), *seeded[0].CacheWritePrice) + }) + + t.Run("Filters", func(t *testing.T) { + t.Parallel() + + // Given: two anthropic models, and an openai model sharing a name with + // one of them. + ownerClient, _ := setupAIModelPricesTest(t) + exp := codersdk.NewExperimentalClient(ownerClient) + setupCtx := testutil.Context(t, testutil.WaitLong) + + //nolint:gocritic // Managing AI model prices is owner-only. + require.NoError(t, exp.UpsertAIModelPrices(setupCtx, codersdk.UpsertAIModelPricesRequest{ + Prices: []codersdk.AIModelPriceUpsert{ + newAIModelPrice("anthropic", "model-a", 1), + newAIModelPrice("anthropic", "model-b", 2), + {Provider: "openai", Model: "model-a", InputPrice: ptr.Ref(int64(3))}, + }, + })) + + tests := []struct { + name string + filter codersdk.AIModelPricesFilter + want []string + }{ + { + name: "NoFilter", + filter: codersdk.AIModelPricesFilter{}, + want: []string{"anthropic/model-a", "anthropic/model-b", "openai/model-a"}, + }, + { + name: "ByProvider", + filter: codersdk.AIModelPricesFilter{Provider: "anthropic"}, + want: []string{"anthropic/model-a", "anthropic/model-b"}, + }, + { + name: "ByModelSpansProviders", + filter: codersdk.AIModelPricesFilter{Model: "model-a"}, + want: []string{"anthropic/model-a", "openai/model-a"}, + }, + { + name: "ByProviderAndModel", + filter: codersdk.AIModelPricesFilter{Provider: "anthropic", Model: "model-a"}, + want: []string{"anthropic/model-a"}, + }, + { + name: "UnknownProviderMatchesNothing", + filter: codersdk.AIModelPricesFilter{Provider: "unknown-provider"}, + want: nil, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + ctx := testutil.Context(t, testutil.WaitLong) + + // When: the prices are listed with the filter. + prices, err := exp.ListAIModelPrices(ctx, tt.filter) + require.NoError(t, err) + + // Then: the expected models come back, and nothing outside the + // filter does. + got := make([]string, 0, len(prices)) + for _, price := range prices { + got = append(got, price.Provider+"/"+price.Model) + if tt.filter.Provider != "" { + require.Equal(t, tt.filter.Provider, price.Provider) + } + if tt.filter.Model != "" { + require.Equal(t, tt.filter.Model, price.Model) + } + } + if len(tt.want) == 0 { + require.Empty(t, got) + return + } + require.Subset(t, got, tt.want) + }) + } + }) + + t.Run("Forbidden", func(t *testing.T) { + t.Parallel() + + // Given: a member without ai_model_price:read. + ownerClient, owner := setupAIModelPricesTest(t) + memberClient, _ := coderdtest.CreateAnotherUser(t, ownerClient, owner.OrganizationID, rbac.RoleMember()) + ctx := testutil.Context(t, testutil.WaitLong) + + // When: they list the prices. + _, err := codersdk.NewExperimentalClient(memberClient).ListAIModelPrices(ctx, codersdk.AIModelPricesFilter{}) + + // Then: the request is forbidden by ai_model_price:read. + var sdkErr *codersdk.Error + require.ErrorAs(t, err, &sdkErr) + require.Equal(t, http.StatusForbidden, sdkErr.StatusCode()) + }) +} diff --git a/enterprise/coderd/coderd.go b/enterprise/coderd/coderd.go index 7fbdc5d7767..bb58222b49f 100644 --- a/enterprise/coderd/coderd.go +++ b/enterprise/coderd/coderd.go @@ -334,6 +334,17 @@ func New(ctx context.Context, options *Options) (_ *API, err error) { }) }) + api.AGPL.ExperimentalHandler.Group(func(r chi.Router) { + r.Route("/ai/model-prices", func(r chi.Router) { + r.Use( + apiKeyMiddleware, + api.RequireFeatureMW(codersdk.FeatureAIBridge), + ) + r.Get("/", api.listAIModelPrices) + r.Post("/", api.upsertAIModelPrices) + }) + }) + api.AGPL.APIHandler.Group(func(r chi.Router) { r.Get("/entitlements", api.serveEntitlements) // /regions overrides the AGPL /regions endpoint diff --git a/scripts/aibridgepricesgen/main.go b/scripts/aibridgepricesgen/main.go index 51ad5bb8bda..4d41e3f1dfa 100644 --- a/scripts/aibridgepricesgen/main.go +++ b/scripts/aibridgepricesgen/main.go @@ -23,20 +23,9 @@ import ( "sort" "golang.org/x/xerrors" -) -// supportedProviders lists the providers we ship prices for. Adding a -// provider here is enough to include it on the next regeneration. -var supportedProviders = []string{ - "anthropic", - "azure", - "bedrock", - "copilot", - "google", - "openai", - "openrouter", - "vercel", -} + "github.com/coder/coder/v2/coderd/aibridge/prices/providers" +) // upstreamProvider is the subset of a models.dev per-provider entry we read. type upstreamProvider struct { @@ -135,7 +124,7 @@ func readUpstream(path string) (map[string]upstreamProvider, error) { } func runPrices(upstream map[string]upstreamProvider) error { - rows, err := convert(upstream, supportedProviders) + rows, err := convert(upstream, providers.SupportedStrings()) if err != nil { return err } @@ -145,7 +134,7 @@ func runPrices(upstream map[string]upstreamProvider) error { if err := write(os.Stdout, rows); err != nil { return err } - _, _ = fmt.Fprintf(os.Stderr, "aibridgepricesgen: wrote %d prices for %d provider(s)\n", len(rows), len(supportedProviders)) + _, _ = fmt.Fprintf(os.Stderr, "aibridgepricesgen: wrote %d prices for %d provider(s)\n", len(rows), len(providers.Supported)) return nil } @@ -169,12 +158,12 @@ 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, providers []string) ([]priceRow, error) { +func convert(upstream map[string]upstreamProvider, providerIDs []string) ([]priceRow, error) { var ( rows []priceRow missing []string ) - for _, providerID := range providers { + for _, providerID := range providerIDs { provider, ok := upstream[providerID] if !ok || len(provider.Models) == 0 { missing = append(missing, providerID) diff --git a/site/src/api/typesGenerated.ts b/site/src/api/typesGenerated.ts index 0222095336d..cfc4fc94650 100644 --- a/site/src/api/typesGenerated.ts +++ b/site/src/api/typesGenerated.ts @@ -352,6 +352,41 @@ export interface AIGatewayKey { */ export const AIGatewayKeyHeader = "X-Coder-AI-Governance-Gateway-Key"; +// From codersdk/aimodelprices.go +/** + * AIModelPrice is a per-model token price used by AI Gateway to compute the + * cost of an interception. + * + * Prices are integer micro-units per million tokens, so 1000000 is $1.00 per + * million tokens. A nil price means the price is not known, which the cost + * calculation treats the same as zero. Distinguish that from an explicit 0, + * which declares the model free of charge. + */ +export interface AIModelPrice { + readonly provider: string; + readonly model: string; + readonly input_price: number | null; + readonly output_price: number | null; + readonly cache_read_price: number | null; + readonly cache_write_price: number | null; + readonly created_at: string; + readonly updated_at: string; +} + +// From codersdk/aimodelprices.go +/** + * AIModelPriceUpsert is one model's prices in an upsert request. It carries + * only the writable fields of AIModelPrice. + */ +export interface AIModelPriceUpsert { + readonly provider: string; + readonly model: string; + readonly input_price: number | null; + readonly output_price: number | null; + readonly cache_read_price: number | null; + readonly cache_write_price: number | null; +} + // From codersdk/aiproviders.go /** * AIProvider represents an AI provider configuration row as returned @@ -5929,6 +5964,12 @@ export interface MatchedProvisioners { readonly most_recently_seen?: string; } +// From codersdk/aimodelprices.go +/** + * MaxAIModelPricesBytes bounds an upsert request body. + */ +export const MaxAIModelPricesBytes = 1048576; // 1 MiB + // From codersdk/aibridge.go /** * MaxAISpendLimitMicros is the highest AI spend limit that can be configured, @@ -10068,6 +10109,15 @@ export interface UploadResponse { readonly hash: string; } +// From codersdk/aimodelprices.go +/** + * UpsertAIModelPricesRequest sets prices for the listed models. Models absent + * from the request are left untouched. + */ +export interface UpsertAIModelPricesRequest { + readonly prices: readonly AIModelPriceUpsert[]; +} + // From codersdk/aibridge.go export interface UpsertGroupAIBudgetRequest { /**