ai-gateway-go/internal/pricing/service.go

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
}