package pricing import ( "context" "encoding/json" "errors" "fmt" "log/slog" "math" "regexp" "sort" "strings" "sync/atomic" "time" platformid "aigateway.local/core/internal/platform/id" "github.com/jackc/pgx/v5" "github.com/jackc/pgx/v5/pgxpool" ) var ErrNotFound = errors.New("model price not found") var codePattern = regexp.MustCompile(`^[a-z][a-z0-9_-]{1,63}$`) type Price struct { ID string `json:"id"` ProviderCode string `json:"provider_code"` ModelPattern string `json:"model_pattern"` InputMicrounitsPerMillion int64 `json:"input_microunits_per_million"` OutputMicrounitsPerMillion int64 `json:"output_microunits_per_million"` Currency string `json:"currency"` EffectiveFrom time.Time `json:"effective_from"` EffectiveTo *time.Time `json:"effective_to"` Enabled bool `json:"enabled"` Revision int64 `json:"revision"` CreatedAt time.Time `json:"created_at"` UpdatedAt time.Time `json:"updated_at"` } type Cost struct { Microunits int64 PriceID, Currency string } type priceSnapshot struct{ prices []Price } type Service struct { pool *pgxpool.Pool refresh time.Duration logger *slog.Logger current atomic.Pointer[priceSnapshot] } func NewService(pool *pgxpool.Pool, refresh time.Duration, logger *slog.Logger) *Service { if refresh <= 0 { refresh = 30 * time.Second } s := &Service{pool: pool, refresh: refresh, logger: logger} s.current.Store(&priceSnapshot{}) return s } func (s *Service) Run(ctx context.Context) { _ = s.Reload(ctx) ticker := time.NewTicker(s.refresh) defer ticker.Stop() for { select { case <-ticker.C: if err := s.Reload(ctx); err != nil && s.logger != nil { s.logger.Warn("model price refresh failed", "error", err) } case <-ctx.Done(): return } } } func (s *Service) Reload(ctx context.Context) error { prices, err := s.List(ctx) if err != nil { return err } active := prices[:0] for _, p := range prices { if p.Enabled { active = append(active, p) } } sort.SliceStable(active, func(i, j int) bool { exactI := !strings.HasSuffix(active[i].ModelPattern, "*") exactJ := !strings.HasSuffix(active[j].ModelPattern, "*") if exactI != exactJ { return exactI } // 通配符之间按前缀长度降序:更具体的模式(gpt-4o*)必须先于宽泛模式 // (gpt-4*)命中,否则 gpt-4o-mini 会按 gpt-4* 的价格错误计费。 // 同长度再按生效时间(新价格优先)。 lengthI, lengthJ := len(active[i].ModelPattern), len(active[j].ModelPattern) if !exactI && lengthI != lengthJ { return lengthI > lengthJ } return active[i].EffectiveFrom.After(active[j].EffectiveFrom) }) s.current.Store(&priceSnapshot{prices: active}) return nil } func (s *Service) Calculate(providerCode, model string, input, output int64, at time.Time) Cost { if s == nil { return Cost{} } providerCode = strings.ToLower(providerCode) for _, p := range s.current.Load().prices { if p.ProviderCode != providerCode || at.Before(p.EffectiveFrom) || (p.EffectiveTo != nil && !at.Before(*p.EffectiveTo)) || !modelMatches(p.ModelPattern, model) { continue } inputCost := ceilMillion(max(input, 0), p.InputMicrounitsPerMillion) outputCost := ceilMillion(max(output, 0), p.OutputMicrounitsPerMillion) cost := inputCost + outputCost if inputCost > math.MaxInt64-outputCost { cost = math.MaxInt64 } return Cost{Microunits: cost, PriceID: p.ID, Currency: p.Currency} } return Cost{} } func ceilMillion(tokens, rate int64) int64 { if tokens <= 0 || rate <= 0 { return 0 } if tokens > (math.MaxInt64-999999)/rate { return math.MaxInt64 } return (tokens*rate + 999999) / 1000000 } func modelMatches(pattern, model string) bool { if strings.HasSuffix(pattern, "*") { return strings.HasPrefix(model, strings.TrimSuffix(pattern, "*")) } return pattern == model } func (s *Service) List(ctx context.Context) ([]Price, error) { rows, err := s.pool.Query(ctx, `SELECT id::text,provider_code,model_pattern,input_microunits_per_million,output_microunits_per_million,currency,effective_from,effective_to,enabled,revision,created_at,updated_at FROM gateway.model_prices ORDER BY provider_code,model_pattern,effective_from DESC`) if err != nil { return nil, err } defer rows.Close() result := []Price{} for rows.Next() { var p Price if err := rows.Scan(&p.ID, &p.ProviderCode, &p.ModelPattern, &p.InputMicrounitsPerMillion, &p.OutputMicrounitsPerMillion, &p.Currency, &p.EffectiveFrom, &p.EffectiveTo, &p.Enabled, &p.Revision, &p.CreatedAt, &p.UpdatedAt); err != nil { return nil, err } result = append(result, p) } return result, rows.Err() } func Validate(p *Price) error { p.ProviderCode = strings.ToLower(strings.TrimSpace(p.ProviderCode)) p.ModelPattern = strings.TrimSpace(p.ModelPattern) p.Currency = strings.ToUpper(strings.TrimSpace(p.Currency)) if !codePattern.MatchString(p.ProviderCode) { return errors.New("provider_code 格式无效") } if p.ModelPattern == "" || len(p.ModelPattern) > 255 || strings.Count(p.ModelPattern, "*") > 1 || (strings.Contains(p.ModelPattern, "*") && !strings.HasSuffix(p.ModelPattern, "*")) { return errors.New("model_pattern 只能是精确模型名或末尾带 * 的前缀") } if p.InputMicrounitsPerMillion < 0 || p.OutputMicrounitsPerMillion < 0 { return errors.New("价格不能为负数") } if p.InputMicrounitsPerMillion > 9_000_000_000_000 || p.OutputMicrounitsPerMillion > 9_000_000_000_000 { return errors.New("价格超过安全上限") } if p.Currency == "" { p.Currency = "USD" } if len(p.Currency) != 3 { return errors.New("currency 必须是 3 位代码") } if p.EffectiveFrom.IsZero() { return errors.New("effective_from 不能为空") } if p.EffectiveTo != nil && !p.EffectiveTo.After(p.EffectiveFrom) { return errors.New("effective_to 必须晚于 effective_from") } return nil } func (s *Service) Save(ctx context.Context, p Price, actorID string, create bool) (Price, error) { if err := Validate(&p); err != nil { return Price{}, err } eventID, err := platformid.NewUUID() if err != nil { return Price{}, err } tx, err := s.pool.Begin(ctx) if err != nil { return Price{}, err } defer func() { _ = tx.Rollback(ctx) }() if create { p.ID, err = platformid.NewUUID() if err != nil { return Price{}, err } err = tx.QueryRow(ctx, `INSERT INTO gateway.model_prices(id,provider_code,model_pattern,input_microunits_per_million,output_microunits_per_million,currency,effective_from,effective_to,enabled,created_by) VALUES($1,$2,$3,$4,$5,$6,$7,$8,$9,$10) RETURNING revision,created_at,updated_at`, p.ID, p.ProviderCode, p.ModelPattern, p.InputMicrounitsPerMillion, p.OutputMicrounitsPerMillion, p.Currency, p.EffectiveFrom, p.EffectiveTo, p.Enabled, actorID).Scan(&p.Revision, &p.CreatedAt, &p.UpdatedAt) } else { err = tx.QueryRow(ctx, `UPDATE gateway.model_prices SET provider_code=$2,model_pattern=$3,input_microunits_per_million=$4,output_microunits_per_million=$5,currency=$6,effective_from=$7,effective_to=$8,enabled=$9,revision=revision+1,updated_at=clock_timestamp() WHERE id=$1 RETURNING revision,created_at,updated_at`, p.ID, p.ProviderCode, p.ModelPattern, p.InputMicrounitsPerMillion, p.OutputMicrounitsPerMillion, p.Currency, p.EffectiveFrom, p.EffectiveTo, p.Enabled).Scan(&p.Revision, &p.CreatedAt, &p.UpdatedAt) } if errors.Is(err, pgx.ErrNoRows) { return Price{}, ErrNotFound } if err != nil { return Price{}, fmt.Errorf("save model price: %w", err) } eventType := "model_price.updated" if create { eventType = "model_price.created" } payload, _ := json.Marshal(map[string]any{"model_price_id": p.ID, "actor_id": actorID, "revision": p.Revision}) if _, err = tx.Exec(ctx, `INSERT INTO gateway.outbox_events(event_id,event_type,event_version,aggregate_type,aggregate_id,payload) VALUES($1,$2,1,'model_price',$3,$4)`, eventID, eventType, p.ID, payload); err != nil { return Price{}, err } if err = tx.Commit(ctx); err != nil { return Price{}, err } if err = s.Reload(ctx); err != nil && s.logger != nil { s.logger.Warn("price reload after save failed; snapshot may be stale", "error", err) } return p, nil } func (s *Service) Delete(ctx context.Context, id, actorID string) error { eventID, err := platformid.NewUUID() if err != nil { return err } tx, err := s.pool.Begin(ctx) if err != nil { return err } defer func() { _ = tx.Rollback(ctx) }() tag, err := tx.Exec(ctx, `DELETE FROM gateway.model_prices WHERE id=$1`, id) if err != nil { return err } if tag.RowsAffected() == 0 { return ErrNotFound } payload, _ := json.Marshal(map[string]string{"model_price_id": id, "actor_id": actorID}) if _, err = tx.Exec(ctx, `INSERT INTO gateway.outbox_events(event_id,event_type,event_version,aggregate_type,aggregate_id,payload) VALUES($1,'model_price.deleted',1,'model_price',$2,$3)`, eventID, id, payload); err != nil { return err } if err = tx.Commit(ctx); err == nil { if reloadErr := s.Reload(ctx); reloadErr != nil && s.logger != nil { s.logger.Warn("price reload after delete failed; snapshot may be stale", "error", reloadErr) } } return err }