mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
feat: add ai_model_prices table (#24932)
# Summary Implements https://linear.app/codercom/issue/AIGOV-282/add-ai-model-price-table-and-seed-generator This PR lays the groundwork for AI Bridge cost controls (per the AI Governance RFC). It adds the foundation needed for future cost tracking: a place to store per-model token prices, a way to keep those prices in sync with upstream pricing data, and a startup mechanism that ensures every deployment has prices loaded before AI Bridge starts processing requests. The price data comes from [models.dev](https://models.dev/), a community-maintained catalogue of AI provider pricing. A generator script fetches the latest prices, filters to Anthropic and OpenAI for now, and produces a seed file checked into the repository. On every server startup the seed is applied to the database, so new releases automatically pick up any price corrections that landed since the previous one. Existing rows are overwritten with the latest prices; rows for models no longer in the seed are left untouched. # Batching the AI model price seed: three approaches Context: at server startup we seed the `ai_model_prices` table from an embedded JSON price book (~70 rows today, will grow as we add providers, potentially 4000+). Each row is: ```text (provider, model, input_price, output_price, cache_read_price, cache_write_price) ``` Any of the four price columns can be: - `NULL` → “price unknown for this dimension” - explicit `0` → “free” The batch must be an UPSERT so re-running is idempotent and existing rows pick up new prices. We considered three implementations. --- ## Approach 1 — Per-row UPSERT in a Go loop ```go for _, row := range rows { if err := db.UpsertAIModelPrice(ctx, database.UpsertAIModelPriceParams{ Provider: row.Provider, Model: row.Model, InputPrice: nullInt64(row.InputPrice), // ... }); err != nil { return err } } ``` ### Pros - Trivial. - NULL handling falls out naturally from `sql.NullInt64`. ### Cons - `N` round-trips per seed. - With ~70 rows that means ~70 statement executions on every startup, even inside a transaction. - Doesn't scale gracefully as the price book grows, potentially 4000+. --- ## Approach 2 — `UNNEST` with parallel arrays Pass each column as a separate Go slice. Postgres unnests them in parallel into a virtual table, then `INSERT ... SELECT`. ```sql INSERT INTO ai_model_prices ( provider, model, input_price, output_price, cache_read_price, cache_write_price ) SELECT UNNEST(@providers::text[]), UNNEST(@models::text[]), NULLIF(UNNEST(@input_prices::bigint[]), -1), NULLIF(UNNEST(@output_prices::bigint[]), -1), NULLIF(UNNEST(@cache_read_prices::bigint[]), -1), NULLIF(UNNEST(@cache_write_prices::bigint[]), -1) ON CONFLICT (provider, model) DO UPDATE SET input_price = EXCLUDED.input_price, output_price = EXCLUDED.output_price, cache_read_price = EXCLUDED.cache_read_price, cache_write_price = EXCLUDED.cache_write_price, updated_at = NOW(); ``` Go side: flatten rows into six parallel slices. Use a sentinel (`-1`) for “missing”, since `lib/pq` can't encode `NULL` into a `bigint[]` element. ```go providers := make([]string, len(rows)) models := make([]string, len(rows)) inputs := make([]int64, len(rows)) outputs := make([]int64, len(rows)) cacheR := make([]int64, len(rows)) cacheW := make([]int64, len(rows)) for i, r := range rows { providers[i] = r.Provider models[i] = r.Model inputs[i] = -1 if r.InputPrice != nil { inputs[i] = *r.InputPrice } outputs[i] = -1 if r.OutputPrice != nil { outputs[i] = *r.OutputPrice } cacheR[i] = -1 if r.CacheReadPrice != nil { cacheR[i] = *r.CacheReadPrice } cacheW[i] = -1 if r.CacheWritePrice != nil { cacheW[i] = *r.CacheWritePrice } } return db.UpsertAIModelPrices(ctx, database.UpsertAIModelPricesParams{ Providers: providers, Models: models, InputPrices: inputs, OutputPrices: outputs, CacheReadPrices: cacheR, CacheWritePrices: cacheW, }) ``` ### Pros - Single round-trip. ### Cons - The generated `sqlc` params become plain `[]int64`, which can't represent `NULL`. --- ## Approach 3 — `jsonb_array_elements` over a single `@seed::jsonb` (chosen) Pass the raw seed JSON as one parameter; let Postgres expand and parse it. ```sql INSERT INTO ai_model_prices ( provider, model, input_price, output_price, cache_read_price, cache_write_price ) SELECT elem->>'provider', elem->>'model', (elem->>'input_price')::bigint, (elem->>'output_price')::bigint, (elem->>'cache_read_price')::bigint, (elem->>'cache_write_price')::bigint FROM jsonb_array_elements(@seed::jsonb) AS elem ON CONFLICT (provider, model) DO UPDATE SET input_price = EXCLUDED.input_price, output_price = EXCLUDED.output_price, cache_read_price = EXCLUDED.cache_read_price, cache_write_price = EXCLUDED.cache_write_price, updated_at = NOW(); ``` Go side reduces to: ```go return db.UpsertAIModelPrices(ctx, seedJSON) ``` ### Pros - Single round-trip. - NULLs fall out naturally: - `(elem->>'cache_write_price')::bigint` becomes `NULL` - no sentinels - The seed is already JSON: - Existing precedent: - `jsonb_array_elements` is already used elsewhere in the codebase ### Cons - Less type-safe at the SQL boundary than `UNNEST` - Slightly less standard than `UNNEST` - Readers need familiarity with: - `jsonb_array_elements` - `->>` extraction syntax - Postgres pays JSON parse cost - negligible at our scale --- --- # Decision We picked Approach 3. It collapses the round-trips like `UNNEST` does, but without: - nullable-array workarounds - sentinel values
This commit is contained in:
Generated
+24
-1
@@ -223,7 +223,10 @@ CREATE TYPE api_key_scope AS ENUM (
|
||||
'chat:*',
|
||||
'ai_seat:*',
|
||||
'ai_seat:create',
|
||||
'ai_seat:read'
|
||||
'ai_seat:read',
|
||||
'ai_model_price:*',
|
||||
'ai_model_price:read',
|
||||
'ai_model_price:update'
|
||||
);
|
||||
|
||||
CREATE TYPE app_sharing_level AS ENUM (
|
||||
@@ -1061,6 +1064,23 @@ BEGIN
|
||||
END;
|
||||
$$;
|
||||
|
||||
CREATE TABLE ai_model_prices (
|
||||
provider text NOT NULL,
|
||||
model text NOT NULL,
|
||||
input_price bigint,
|
||||
output_price bigint,
|
||||
cache_read_price bigint,
|
||||
cache_write_price bigint,
|
||||
created_at timestamp with time zone DEFAULT now() NOT NULL,
|
||||
updated_at timestamp with time zone DEFAULT now() NOT NULL,
|
||||
CONSTRAINT ai_model_prices_cache_read_price_check CHECK ((cache_read_price >= 0)),
|
||||
CONSTRAINT ai_model_prices_cache_write_price_check CHECK ((cache_write_price >= 0)),
|
||||
CONSTRAINT ai_model_prices_input_price_check CHECK ((input_price >= 0)),
|
||||
CONSTRAINT ai_model_prices_output_price_check CHECK ((output_price >= 0))
|
||||
);
|
||||
|
||||
COMMENT ON TABLE ai_model_prices IS 'Per-model token prices used by AI Bridge to compute interception cost.';
|
||||
|
||||
CREATE TABLE ai_seat_state (
|
||||
user_id uuid NOT NULL,
|
||||
first_used_at timestamp with time zone NOT NULL,
|
||||
@@ -3358,6 +3378,9 @@ ALTER TABLE ONLY workspace_resource_metadata ALTER COLUMN id SET DEFAULT nextval
|
||||
ALTER TABLE ONLY workspace_agent_stats
|
||||
ADD CONSTRAINT agent_stats_pkey PRIMARY KEY (id);
|
||||
|
||||
ALTER TABLE ONLY ai_model_prices
|
||||
ADD CONSTRAINT ai_model_prices_pkey PRIMARY KEY (provider, model);
|
||||
|
||||
ALTER TABLE ONLY ai_seat_state
|
||||
ADD CONSTRAINT ai_seat_state_pkey PRIMARY KEY (user_id);
|
||||
|
||||
|
||||
Reference in New Issue
Block a user