AI Gateway Go 0.10.0 源码快照 + 旗舰版需求规划报告

M0-M7 已完成:核心网关(身份/RBAC/TOTP/OIDC/SAML/Provider/配额/路由/内容策略/审计/定价)+ 资源市场(MCP/Skills/数字员工)。含 22 个 PostgreSQL 迁移、管理端/门户端前端源码、OpenAPI 契约、部署 compose。 Co-Authored-By: Claude <noreply@anthropic.com>
2026-08-12 11:45:54 +08:00
commit 5759c1862e
807 changed files with 114727 additions and 0 deletions
@@ -0,0 +1,252 @@
+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
+		}
+		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
+}