From 2d9b6eda8f0a2e374c4197121008c3fd381c4b86 Mon Sep 17 00:00:00 2001 From: Susana Ferreira Date: Thu, 13 Aug 2026 15:00:36 +0100 Subject: [PATCH] feat: add experimental CLI to price unpriced AI models (#27926) ## Description AI Gateway computes the cost of an interception from `ai_model_prices`, which is seeded on every server start from a price book embedded in the binary. A model the price book does not cover records a NULL cost, so its spend is invisible to cost reporting and is not enforced against budgets. The only fix was to wait for a Coder release that added the model. This adds an experimental CLI, backed by an experimental HTTP endpoint, for pricing those models. Models the price book already covers are rejected, because the seeder re-applies the book on every start and would overwrite an operator price. Support for custom pricing will be handled in https://linear.app/codercom/issue/AIGOV-589/extend-experimental-cli-command-to-set-custom-ai-model-prices. ## Commands ``` coder exp ai-model-prices list [--provider] [--model] coder exp ai-model-prices update [file|-] [--provider] [--model] [--input-price] [--output-price] [--cache-read-price] [--cache-write-price] [--yes] ``` ## Changes - Add `GET` and `POST /api/experimental/ai/model-prices`, gated behind the AI Bridge entitlement and the existing `ai_model_price` RBAC resource. - Add a `GetAIModelPrices` query with optional `provider` and `model` filters applied in SQL. - Validate the whole request before writing anything, so one bad entry cannot leave the table half updated, and report every problem at once. - Reject prices for models the embedded price book already covers, through a new `prices.IsDefaultPriced`. - Add the `coder exp ai-model-prices` command with `list` and `update`. `update` accepts a JSON document or the single-model flags and prints a plan, asking to confirm unless the document is piped in or `--yes` is passed. - Consolidate the supported provider list into `coderd/aibridge/prices/providers` so the price generator and the server share one definition. - Add `codersdk` types and client methods for both endpoints, and bound the request body at 1 MiB. - Document the command in the AI Gateway cost controls page. Closes https://linear.app/codercom/issue/AIGOV-567/experimental-cli-command-to-set-prices-for-unpriced-ai-models > [!NOTE] > Initially generated by Claude Opus 5, modified and reviewed by @ssncferreira --- coderd/aibridge/prices/prices.go | 30 + coderd/aibridge/prices/prices_test.go | 53 ++ coderd/aibridge/prices/providers/providers.go | 28 + coderd/apidoc/docs.go | 144 +++++ coderd/apidoc/swagger.json | 136 +++++ coderd/database/db2sdk/db2sdk.go | 21 + coderd/database/dbauthz/dbauthz.go | 7 + coderd/database/dbauthz/dbauthz_test.go | 5 + coderd/database/dbmetrics/querymetrics.go | 8 + coderd/database/dbmock/dbmock.go | 15 + coderd/database/querier.go | 1 + coderd/database/querier_test.go | 71 +++ coderd/database/queries.sql.go | 55 ++ coderd/database/queries/aicostcontrol.sql | 17 + codersdk/aimodelprices.go | 96 ++++ docs/ai-coder/ai-gateway/cost-controls.md | 41 +- docs/reference/api/schemas.md | 75 +++ enterprise/cli/exp_aimodelprices.go | 521 ++++++++++++++++++ .../cli/exp_aimodelprices_internal_test.go | 131 +++++ enterprise/cli/exp_aimodelprices_test.go | 418 ++++++++++++++ enterprise/cli/root.go | 6 +- enterprise/coderd/aimodelprices.go | 275 +++++++++ .../coderd/aimodelprices_internal_test.go | 194 +++++++ enterprise/coderd/aimodelprices_test.go | 416 ++++++++++++++ enterprise/coderd/coderd.go | 11 + scripts/aibridgepricesgen/main.go | 23 +- site/src/api/typesGenerated.ts | 50 ++ 27 files changed, 2828 insertions(+), 20 deletions(-) create mode 100644 coderd/aibridge/prices/providers/providers.go create mode 100644 codersdk/aimodelprices.go create mode 100644 enterprise/cli/exp_aimodelprices.go create mode 100644 enterprise/cli/exp_aimodelprices_internal_test.go create mode 100644 enterprise/cli/exp_aimodelprices_test.go create mode 100644 enterprise/coderd/aimodelprices.go create mode 100644 enterprise/coderd/aimodelprices_internal_test.go create mode 100644 enterprise/coderd/aimodelprices_test.go diff --git a/coderd/aibridge/prices/prices.go b/coderd/aibridge/prices/prices.go index bbb5689ea0..5cac067c55 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 ae30b7e449..4ecdabd2ca 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 0000000000..62d09d6a5f --- /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 f3dd16d01d..6660dcff75 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 e0aa90bda2..1b8d8b86e9 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 67bff9308a..943c2bb955 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 ad1becaf35..c661819287 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 f175c6b216..86dc818f54 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 795e777ee8..4a986ad1eb 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 418a4c205d..4f09e5a4b7 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 34e74dedc4..bbc90e0285 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 d416a3f688..084e61602c 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 5ada09d52b..bb57a04bb3 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 5552e8c7af..def9c0d316 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 0000000000..00cae36d18 --- /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 52c4fa13c4..dba6be0e63 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 1e6dbe6e0e..527d79f185 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 0000000000..c3b39734e8 --- /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 0000000000..2f8a7ec4e3 --- /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 0000000000..258c34eaf5 --- /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 f081335722..877842a122 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 0000000000..78f2d06846 --- /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 0000000000..75d09e5fb0 --- /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 0000000000..0902d2a060 --- /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 7fbdc5d776..bb58222b49 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 51ad5bb8bd..4d41e3f1df 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 0222095336..cfc4fc9465 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 { /**