BEGIN;

ALTER TABLE models DROP CONSTRAINT IF EXISTS models_text_billing_mode_check;
ALTER TABLE models ADD COLUMN IF NOT EXISTS max_output_tokens integer;
UPDATE models SET max_output_tokens = 4096 WHERE model_type = 'text' AND text_billing_mode = 'per_token';
ALTER TABLE models ADD CONSTRAINT models_text_billing_mode_check CHECK (
    (model_type = 'text' AND text_billing_mode IN ('per_request', 'per_token') AND (text_billing_mode <> 'per_token' OR max_output_tokens > 0))
    OR (model_type <> 'text' AND text_billing_mode IS NULL AND max_output_tokens IS NULL)
);

COMMIT;