package store import ( "context" "encoding/json" "fmt" "strconv" "strings" "time" "github.com/jackc/pgx/v5" ) const pricingRuleSetColumns = ` id::text, rule_set_key, name, COALESCE(description, ''), category, currency, status, metadata, created_at, updated_at` const pricingRuleColumns = ` id::text, COALESCE(rule_set_id::text, ''), rule_key, display_name, scope_type, COALESCE(scope_id::text, ''), resource_type, unit, base_price::float8, is_free, currency, base_weight, dynamic_weight, calculator_type, dimension_schema, formula_config, priority, status, metadata, COALESCE(effective_from::text, ''), COALESCE(effective_to::text, ''), created_at, updated_at` type PricingRuleInput struct { RuleKey string `json:"ruleKey"` DisplayName string `json:"displayName"` ResourceType string `json:"resourceType"` Unit string `json:"unit"` BasePrice float64 `json:"basePrice"` IsFree bool `json:"isFree"` Currency string `json:"currency"` BaseWeight map[string]any `json:"baseWeight"` DynamicWeight map[string]any `json:"dynamicWeight"` CalculatorType string `json:"calculatorType"` DimensionSchema map[string]any `json:"dimensionSchema"` FormulaConfig map[string]any `json:"formulaConfig"` Priority int `json:"priority"` Status string `json:"status"` Metadata map[string]any `json:"metadata"` EffectiveFrom string `json:"effectiveFrom"` EffectiveTo string `json:"effectiveTo"` } type EffectivePricingConfig struct { RuleSetID string RuleSetKey string Currency string Config map[string]any FreeResource map[string]bool Snapshot map[string]any } type PricingRuleSetInput struct { RuleSetKey string `json:"ruleSetKey"` Name string `json:"name"` Description string `json:"description"` Category string `json:"category"` Currency string `json:"currency"` Status string `json:"status"` Metadata map[string]any `json:"metadata"` Rules []PricingRuleInput `json:"rules"` } type pricingScanner interface { Scan(dest ...any) error } func (s *Store) ListPricingRuleSets(ctx context.Context) ([]PricingRuleSet, error) { rows, err := s.pool.Query(ctx, `SELECT `+pricingRuleSetColumns+` FROM model_pricing_rule_sets ORDER BY category ASC, name ASC`) if err != nil { return nil, err } defer rows.Close() items := make([]PricingRuleSet, 0) byID := map[string]int{} for rows.Next() { item, err := scanPricingRuleSet(rows) if err != nil { return nil, err } byID[item.ID] = len(items) items = append(items, item) } if err := rows.Err(); err != nil { return nil, err } ruleRows, err := s.pool.Query(ctx, ` SELECT `+pricingRuleColumns+` FROM model_pricing_rules WHERE rule_set_id IS NOT NULL ORDER BY rule_set_id, priority ASC, resource_type ASC, rule_key ASC`) if err != nil { return nil, err } defer ruleRows.Close() for ruleRows.Next() { rule, err := scanPricingRule(ruleRows) if err != nil { return nil, err } if index, ok := byID[rule.RuleSetID]; ok { items[index].Rules = append(items[index].Rules, rule) } } return items, ruleRows.Err() } func (s *Store) CreatePricingRuleSet(ctx context.Context, input PricingRuleSetInput) (PricingRuleSet, error) { input = normalizePricingRuleSet(input) tx, err := s.pool.Begin(ctx) if err != nil { return PricingRuleSet{}, err } defer tx.Rollback(ctx) metadata, _ := json.Marshal(emptyObjectIfNil(input.Metadata)) item, err := scanPricingRuleSet(tx.QueryRow(ctx, ` INSERT INTO model_pricing_rule_sets (rule_set_key, name, description, category, currency, status, metadata) VALUES ($1, $2, NULLIF($3, ''), $4, $5, $6, $7) RETURNING `+pricingRuleSetColumns, input.RuleSetKey, input.Name, input.Description, input.Category, input.Currency, input.Status, metadata, )) if err != nil { return PricingRuleSet{}, err } if err := insertPricingRules(ctx, tx, item.ID, input.Currency, input.Rules); err != nil { return PricingRuleSet{}, err } if err := tx.Commit(ctx); err != nil { return PricingRuleSet{}, err } item.Rules = pricingInputsToRules(item.ID, input.Currency, input.Rules) return item, nil } func (s *Store) UpdatePricingRuleSet(ctx context.Context, id string, input PricingRuleSetInput) (PricingRuleSet, error) { input = normalizePricingRuleSet(input) tx, err := s.pool.Begin(ctx) if err != nil { return PricingRuleSet{}, err } defer tx.Rollback(ctx) metadata, _ := json.Marshal(emptyObjectIfNil(input.Metadata)) item, err := scanPricingRuleSet(tx.QueryRow(ctx, ` UPDATE model_pricing_rule_sets SET rule_set_key = $2, name = $3, description = NULLIF($4, ''), category = $5, currency = $6, status = $7, metadata = $8, updated_at = now() WHERE id = $1::uuid RETURNING `+pricingRuleSetColumns, id, input.RuleSetKey, input.Name, input.Description, input.Category, input.Currency, input.Status, metadata, )) if err != nil { return PricingRuleSet{}, err } if _, err := tx.Exec(ctx, `DELETE FROM model_pricing_rules WHERE rule_set_id = $1::uuid`, id); err != nil { return PricingRuleSet{}, err } if err := insertPricingRules(ctx, tx, item.ID, input.Currency, input.Rules); err != nil { return PricingRuleSet{}, err } if err := tx.Commit(ctx); err != nil { return PricingRuleSet{}, err } item.Rules = pricingInputsToRules(item.ID, input.Currency, input.Rules) return item, nil } func (s *Store) DeletePricingRuleSet(ctx context.Context, id string) error { var ruleSetKey string if err := s.pool.QueryRow(ctx, `SELECT rule_set_key FROM model_pricing_rule_sets WHERE id = $1::uuid`, id).Scan(&ruleSetKey); err != nil { return err } if ruleSetKey == "default-multimodal-v1" { return ErrProtectedDefault } result, err := s.pool.Exec(ctx, `DELETE FROM model_pricing_rule_sets WHERE id = $1::uuid`, id) if err != nil { return err } if result.RowsAffected() == 0 { return pgx.ErrNoRows } return nil } func (s *Store) PricingRuleSetBillingConfig(ctx context.Context, id string) (map[string]any, error) { id = strings.TrimSpace(id) if id == "" { return nil, nil } rows, err := s.pool.Query(ctx, ` SELECT resource_type, base_price::float8, dynamic_weight, formula_config FROM model_pricing_rules WHERE rule_set_id = $1::uuid AND status = 'active' ORDER BY priority ASC, resource_type ASC`, id) if err != nil { return nil, err } defer rows.Close() config := map[string]any{} for rows.Next() { var resourceType string var basePrice float64 var dynamicWeightBytes []byte var formulaConfigBytes []byte if err := rows.Scan(&resourceType, &basePrice, &dynamicWeightBytes, &formulaConfigBytes); err != nil { return nil, err } dynamicWeight := decodeObject(dynamicWeightBytes) formulaConfig := decodeObject(formulaConfigBytes) switch resourceType { case "text_input": config["textInputPer1k"] = basePrice case "text_cached_input": config["textCachedInputPer1k"] = basePrice case "text_output": config["textOutputPer1k"] = basePrice case "text_total": inputPrice := basePrice if value, ok := pricingRuleNumberFromKeys(formulaConfig, "inputTokenPrice", "input_token_price", "textInputPer1k", "text_input"); ok { inputPrice = value } config["textInputPer1k"] = inputPrice if cachedInputPrice, ok := pricingRuleNumberFromKeys(formulaConfig, "cachedInputTokenPrice", "cached_input_token_price", "textCachedInputPer1k", "text_cached_input", "inputCacheHitTokenPrice", "input_cache_hit_token_price"); ok { config["textCachedInputPer1k"] = cachedInputPrice } if outputPrice, ok := pricingRuleNumberFromKeys(formulaConfig, "outputTokenPrice", "output_token_price", "textOutputPer1k", "text_output"); ok { config["textOutputPer1k"] = outputPrice } resourceConfig := pricingResourceConfig(basePrice, dynamicWeight) if len(formulaConfig) > 0 { resourceConfig["formulaConfig"] = formulaConfig } config["text_total"] = resourceConfig case "image": config["imageBase"] = basePrice config["image"] = pricingResourceConfig(basePrice, dynamicWeight) case "image_edit": config["editBase"] = basePrice config["image_edit"] = pricingResourceConfig(basePrice, dynamicWeight) case "video": config["videoBase"] = basePrice config["video"] = pricingResourceConfig(basePrice, dynamicWeight) default: config[resourceType] = pricingResourceConfig(basePrice, dynamicWeight) } } if err := rows.Err(); err != nil { return nil, err } return config, nil } func (s *Store) PricingRuleSetBillingConfigV2(ctx context.Context, id string) (EffectivePricingConfig, error) { id = strings.TrimSpace(id) if id == "" { return EffectivePricingConfig{}, fmt.Errorf("pricing rule set id is required") } var ruleSetKey string var currency string var status string if err := s.pool.QueryRow(ctx, ` SELECT rule_set_key, currency, status FROM model_pricing_rule_sets WHERE id = $1::uuid`, id).Scan(&ruleSetKey, ¤cy, &status); err != nil { return EffectivePricingConfig{}, err } if status != "active" { return EffectivePricingConfig{}, fmt.Errorf("pricing rule set %s is not active", id) } if currency != "resource" { return EffectivePricingConfig{}, fmt.Errorf("pricing rule set %s uses unsupported currency %s", id, currency) } rows, err := s.pool.Query(ctx, ` SELECT rule_key, resource_type, unit, base_price::text, currency, base_weight, dynamic_weight, calculator_type, dimension_schema, formula_config, priority, is_free, COALESCE(effective_from::text, ''), COALESCE(effective_to::text, '') FROM model_pricing_rules WHERE rule_set_id = $1::uuid AND status = 'active' AND (effective_from IS NULL OR effective_from <= now()) AND (effective_to IS NULL OR effective_to > now()) ORDER BY priority ASC, resource_type ASC, rule_key ASC`, id) if err != nil { return EffectivePricingConfig{}, err } defer rows.Close() config := map[string]any{} freeResource := map[string]bool{} rules := make([]any, 0) seen := map[string]bool{} for rows.Next() { var ruleKey string var resourceType string var unit string var basePrice string var ruleCurrency string var baseWeightBytes []byte var dynamicWeightBytes []byte var calculatorType string var dimensionSchemaBytes []byte var formulaConfigBytes []byte var priority int var isFree bool var effectiveFrom string var effectiveTo string if err := rows.Scan( &ruleKey, &resourceType, &unit, &basePrice, &ruleCurrency, &baseWeightBytes, &dynamicWeightBytes, &calculatorType, &dimensionSchemaBytes, &formulaConfigBytes, &priority, &isFree, &effectiveFrom, &effectiveTo, ); err != nil { return EffectivePricingConfig{}, err } if seen[resourceType] { continue } if ruleCurrency != currency { return EffectivePricingConfig{}, fmt.Errorf("pricing rule %s currency does not match its rule set", ruleKey) } switch calculatorType { case "token_usage", "unit_weight", "duration_weight": default: return EffectivePricingConfig{}, fmt.Errorf("pricing rule %s uses unsupported calculator %s", ruleKey, calculatorType) } if err := ValidateEffectivePricingRuleShape(resourceType, unit, calculatorType); err != nil { return EffectivePricingConfig{}, fmt.Errorf("pricing rule %s: %w", ruleKey, err) } if strings.HasPrefix(strings.TrimSpace(basePrice), "-") { return EffectivePricingConfig{}, fmt.Errorf("pricing rule %s has a negative base price", ruleKey) } dynamicWeight := decodeObject(dynamicWeightBytes) formulaConfig := decodeObject(formulaConfigBytes) addPricingRuleToConfig(config, resourceType, basePrice, dynamicWeight, formulaConfig) freeResource[resourceType] = isFree seen[resourceType] = true rules = append(rules, map[string]any{ "ruleKey": ruleKey, "resourceType": resourceType, "unit": unit, "basePrice": basePrice, "currency": ruleCurrency, "baseWeight": decodeObject(baseWeightBytes), "dynamicWeight": dynamicWeight, "calculatorType": calculatorType, "dimensionSchema": decodeObject(dimensionSchemaBytes), "formulaConfig": formulaConfig, "priority": priority, "isFree": isFree, "effectiveFrom": effectiveFrom, "effectiveTo": effectiveTo, }) } if err := rows.Err(); err != nil { return EffectivePricingConfig{}, err } if len(rules) == 0 { return EffectivePricingConfig{}, fmt.Errorf("pricing rule set %s has no effective rules", id) } return EffectivePricingConfig{ RuleSetID: id, RuleSetKey: ruleSetKey, Currency: currency, Config: config, FreeResource: freeResource, Snapshot: map[string]any{ "pricingVersion": "effective-pricing-v2", "ruleSetId": id, "ruleSetKey": ruleSetKey, "currency": currency, "resolvedAt": time.Now().UTC().Format(time.RFC3339Nano), "rules": rules, }, }, nil } func ValidateEffectivePricingRuleShape(resourceType string, unit string, calculatorType string) error { unit = NormalizeEffectivePricingRuleUnit(resourceType, unit) if strings.TrimSpace(calculatorType) == "" { calculatorType = DefaultEffectivePricingCalculator(resourceType) } allowed := false switch resourceType { case "text_input", "text_cached_input", "text_output", "text_total": allowed = unit == "1k_tokens" && calculatorType == "token_usage" case "image", "image_edit": allowed = unit == "image" && calculatorType == "unit_weight" case "video": allowed = unit == "5s" && calculatorType == "duration_weight" case "music": allowed = (unit == "song" || unit == "item") && calculatorType == "unit_weight" case "audio": allowed = unit == "character" && calculatorType == "unit_weight" default: return nil } if !allowed { return fmt.Errorf("unit %q and calculator %q do not match resource %q", unit, calculatorType, resourceType) } return nil } func NormalizeEffectivePricingRuleUnit(resourceType string, unit string) string { unit = strings.TrimSpace(unit) if resourceType == "video" && unit == "video" { return "5s" } return unit } func DefaultEffectivePricingCalculator(resourceType string) string { switch strings.TrimSpace(resourceType) { case "text_input", "text_cached_input", "text_output", "text_total": return "token_usage" case "video": return "duration_weight" default: return "unit_weight" } } func addPricingRuleToConfig(config map[string]any, resourceType string, basePrice string, dynamicWeight map[string]any, formulaConfig map[string]any) { resourceConfig := map[string]any{"basePrice": basePrice} if len(dynamicWeight) > 0 { resourceConfig["dynamicWeight"] = dynamicWeight } if len(formulaConfig) > 0 { resourceConfig["formulaConfig"] = formulaConfig } switch resourceType { case "text_input": config["textInputPer1k"] = basePrice case "text_cached_input": config["textCachedInputPer1k"] = basePrice case "text_output": config["textOutputPer1k"] = basePrice case "text_total": inputPrice := any(basePrice) if value, ok := pricingRuleValueFromKeys(formulaConfig, "inputTokenPrice", "input_token_price", "textInputPer1k", "text_input"); ok { inputPrice = value } config["textInputPer1k"] = inputPrice if value, ok := pricingRuleValueFromKeys(formulaConfig, "cachedInputTokenPrice", "cached_input_token_price", "textCachedInputPer1k", "text_cached_input", "inputCacheHitTokenPrice", "input_cache_hit_token_price"); ok { config["textCachedInputPer1k"] = value } if value, ok := pricingRuleValueFromKeys(formulaConfig, "outputTokenPrice", "output_token_price", "textOutputPer1k", "text_output"); ok { config["textOutputPer1k"] = value } config["text_total"] = resourceConfig case "image": config["imageBase"] = basePrice config["image"] = resourceConfig case "image_edit": config["editBase"] = basePrice config["image_edit"] = resourceConfig case "video": config["videoBase"] = basePrice config["video"] = resourceConfig case "music": config["musicBase"] = basePrice config["music"] = resourceConfig case "audio": config["audioBase"] = basePrice config["audio"] = resourceConfig default: config[resourceType] = resourceConfig } } func pricingRuleValueFromKeys(config map[string]any, keys ...string) (any, bool) { for _, key := range keys { if value, ok := config[key]; ok { return value, true } } return nil, false } func pricingRuleNumberFromKeys(config map[string]any, keys ...string) (float64, bool) { if len(config) == 0 { return 0, false } for _, key := range keys { if value, ok := pricingRuleNumberValue(config[key]); ok { return value, true } } return 0, false } func pricingRuleNumberValue(value any) (float64, bool) { switch typed := value.(type) { case float64: return typed, true case float32: return float64(typed), true case int: return float64(typed), true case int64: return float64(typed), true case int32: return float64(typed), true case json.Number: number, err := typed.Float64() return number, err == nil case string: trimmed := strings.TrimSpace(typed) if trimmed == "" { return 0, false } number, err := strconv.ParseFloat(trimmed, 64) return number, err == nil default: return 0, false } } func pricingResourceConfig(basePrice float64, dynamicWeight map[string]any) map[string]any { config := map[string]any{"basePrice": basePrice} if len(dynamicWeight) > 0 { config["dynamicWeight"] = dynamicWeight } return config } func insertPricingRules(ctx context.Context, tx pgx.Tx, ruleSetID string, defaultCurrency string, rules []PricingRuleInput) error { for index, rule := range rules { rule = normalizePricingRule(rule, index, defaultCurrency) baseWeight, _ := json.Marshal(emptyObjectIfNil(rule.BaseWeight)) dynamicWeight, _ := json.Marshal(emptyObjectIfNil(rule.DynamicWeight)) dimensionSchema, _ := json.Marshal(emptyObjectIfNil(rule.DimensionSchema)) formulaConfig, _ := json.Marshal(emptyObjectIfNil(rule.FormulaConfig)) metadata, _ := json.Marshal(emptyObjectIfNil(rule.Metadata)) if _, err := tx.Exec(ctx, ` INSERT INTO model_pricing_rules ( rule_set_id, rule_key, display_name, scope_type, scope_id, resource_type, unit, base_price, is_free, currency, base_weight, dynamic_weight, calculator_type, dimension_schema, formula_config, priority, status, metadata, effective_from, effective_to ) VALUES ( $1::uuid, $2, $3, 'rule_set', $1::uuid, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, NULLIF($17, '')::timestamptz, NULLIF($18, '')::timestamptz )`, ruleSetID, rule.RuleKey, rule.DisplayName, rule.ResourceType, rule.Unit, rule.BasePrice, rule.IsFree, rule.Currency, baseWeight, dynamicWeight, rule.CalculatorType, dimensionSchema, formulaConfig, rule.Priority, rule.Status, metadata, rule.EffectiveFrom, rule.EffectiveTo, ); err != nil { return err } } return nil } func scanPricingRuleSet(scanner pricingScanner) (PricingRuleSet, error) { var item PricingRuleSet var metadata []byte if err := scanner.Scan( &item.ID, &item.RuleSetKey, &item.Name, &item.Description, &item.Category, &item.Currency, &item.Status, &metadata, &item.CreatedAt, &item.UpdatedAt, ); err != nil { return PricingRuleSet{}, err } item.Metadata = decodeObject(metadata) item.Rules = []PricingRule{} return item, nil } func scanPricingRule(scanner pricingScanner) (PricingRule, error) { var item PricingRule var baseWeight []byte var dynamicWeight []byte var dimensionSchema []byte var formulaConfig []byte var metadata []byte if err := scanner.Scan( &item.ID, &item.RuleSetID, &item.RuleKey, &item.DisplayName, &item.ScopeType, &item.ScopeID, &item.ResourceType, &item.Unit, &item.BasePrice, &item.IsFree, &item.Currency, &baseWeight, &dynamicWeight, &item.CalculatorType, &dimensionSchema, &formulaConfig, &item.Priority, &item.Status, &metadata, &item.EffectiveFrom, &item.EffectiveTo, &item.CreatedAt, &item.UpdatedAt, ); err != nil { return PricingRule{}, err } item.BaseWeight = decodeObject(baseWeight) item.DynamicWeight = decodeObject(dynamicWeight) item.DimensionSchema = decodeObject(dimensionSchema) item.FormulaConfig = decodeObject(formulaConfig) item.Metadata = decodeObject(metadata) return item, nil } func normalizePricingRuleSet(input PricingRuleSetInput) PricingRuleSetInput { input.RuleSetKey = strings.TrimSpace(input.RuleSetKey) input.Name = strings.TrimSpace(input.Name) input.Description = strings.TrimSpace(input.Description) input.Category = strings.TrimSpace(input.Category) input.Currency = strings.TrimSpace(input.Currency) input.Status = strings.TrimSpace(input.Status) if input.Category == "" { input.Category = "custom" } if input.Currency == "" { input.Currency = "resource" } if input.Status == "" { input.Status = "active" } return input } func normalizePricingRule(input PricingRuleInput, index int, defaultCurrency string) PricingRuleInput { input.RuleKey = strings.TrimSpace(input.RuleKey) input.DisplayName = strings.TrimSpace(input.DisplayName) input.ResourceType = strings.TrimSpace(input.ResourceType) input.Unit = NormalizeEffectivePricingRuleUnit(input.ResourceType, input.Unit) input.Currency = strings.TrimSpace(input.Currency) input.CalculatorType = strings.TrimSpace(input.CalculatorType) input.Status = strings.TrimSpace(input.Status) if input.RuleKey == "" { input.RuleKey = "rule_" + strings.ReplaceAll(input.ResourceType+"_"+input.Unit, " ", "_") } if input.DisplayName == "" { input.DisplayName = input.ResourceType } if input.Unit == "" { input.Unit = "item" } if input.Currency == "" { input.Currency = defaultCurrency } if input.CalculatorType == "" { input.CalculatorType = DefaultEffectivePricingCalculator(input.ResourceType) } if input.Priority == 0 { input.Priority = (index + 1) * 10 } if input.Status == "" { input.Status = "active" } return input } func pricingInputsToRules(ruleSetID string, defaultCurrency string, rules []PricingRuleInput) []PricingRule { items := make([]PricingRule, 0, len(rules)) for index, input := range rules { input = normalizePricingRule(input, index, defaultCurrency) items = append(items, PricingRule{ RuleSetID: ruleSetID, RuleKey: input.RuleKey, DisplayName: input.DisplayName, ScopeType: "rule_set", ScopeID: ruleSetID, ResourceType: input.ResourceType, Unit: input.Unit, BasePrice: input.BasePrice, IsFree: input.IsFree, Currency: input.Currency, BaseWeight: emptyObjectIfNil(input.BaseWeight), DynamicWeight: emptyObjectIfNil(input.DynamicWeight), CalculatorType: input.CalculatorType, DimensionSchema: emptyObjectIfNil(input.DimensionSchema), FormulaConfig: emptyObjectIfNil(input.FormulaConfig), Priority: input.Priority, Status: input.Status, Metadata: emptyObjectIfNil(input.Metadata), EffectiveFrom: input.EffectiveFrom, EffectiveTo: input.EffectiveTo, }) } return items }