add some scriptsmain

1 files changed, 352 insertions, 0 deletions

diff --git a/internal/billing/service.go b/internal/billing/service.go
new file mode 100644
index 0000000..b87a0bd
--- /dev/null
+++ b/internal/billing/service.go
@@ -0,0 +1,352 @@
+package billing
+
+import (
+	"context"
+	"crypto/rand"
+	"encoding/hex"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"math/big"
+	"strings"
+	"time"
+
+	"aigw/internal/domain"
+
+	"github.com/jackc/pgx/v5"
+	"github.com/jackc/pgx/v5/pgxpool"
+)
+
+const microsPerUnit = int64(1_000_000)
+
+type Service struct {
+	db                     *pgxpool.Pool
+	currency               string
+	defaultMaxOutputTokens int64
+	minTopUpMinor          int64
+	maxTopUpMinor          int64
+	stripeEnabled          bool
+	stripeWebhookSecret    string
+	stripeSuccessURL       string
+	stripeCancelURL        string
+	integrationIdentifier  string
+	createStripeCheckout   stripeCheckoutCreator
+}
+
+func New(ctx context.Context, options Options) (*Service, error) {
+	db, err := pgxpool.New(ctx, options.DatabaseURL)
+	if err != nil {
+		return nil, fmt.Errorf("configure billing PostgreSQL: %w", err)
+	}
+	if err := db.Ping(ctx); err != nil {
+		db.Close()
+		return nil, fmt.Errorf("connect billing PostgreSQL: %w", err)
+	}
+	service := &Service{
+		db: db, currency: options.Currency, defaultMaxOutputTokens: options.DefaultMaxOutputTokens,
+		minTopUpMinor: options.MinTopUpMinor, maxTopUpMinor: options.MaxTopUpMinor,
+		stripeEnabled: options.StripeEnabled, stripeWebhookSecret: options.StripeWebhookSecret,
+		stripeSuccessURL: options.StripeSuccessURL, stripeCancelURL: options.StripeCancelURL,
+		integrationIdentifier: "aigw_balance_" + randomLetters(8),
+	}
+	if options.StripeEnabled {
+		service.createStripeCheckout = newStripeCheckoutCreator(options.StripeAPIKey)
+	}
+	return service, nil
+}
+
+func (s *Service) Close() {
+	s.db.Close()
+}
+
+func (s *Service) StripeEnabled() bool {
+	return s.stripeEnabled
+}
+
+func (s *Service) Currency() string {
+	return s.currency
+}
+
+func (s *Service) Authorize(ctx context.Context, input Authorization) error {
+	if input.RequestID == "" || input.Principal.TenantID == "" || input.Principal.ProjectID == "" || input.Principal.KeyID == "" {
+		return errors.New("billing authorization identity is incomplete")
+	}
+	reserved, err := reservationCost(input.Model, input.Body, s.defaultMaxOutputTokens)
+	if err != nil {
+		return err
+	}
+	tx, err := s.db.BeginTx(ctx, pgx.TxOptions{IsoLevel: pgx.ReadCommitted})
+	if err != nil {
+		return fmt.Errorf("begin billing authorization: %w", err)
+	}
+	defer tx.Rollback(ctx)
+	if _, err := tx.Exec(ctx, `
+		INSERT INTO tenant_wallets (tenant_id, currency) VALUES ($1, $2)
+		ON CONFLICT (tenant_id) DO NOTHING`, input.Principal.TenantID, s.currency); err != nil {
+		return fmt.Errorf("ensure tenant wallet: %w", err)
+	}
+	var currency string
+	var balance, held int64
+	if err := tx.QueryRow(ctx, `
+		SELECT currency, balance_micros, reserved_micros FROM tenant_wallets
+		WHERE tenant_id = $1 FOR UPDATE`, input.Principal.TenantID).Scan(&currency, &balance, &held); err != nil {
+		return fmt.Errorf("lock tenant wallet: %w", err)
+	}
+	if currency != s.currency {
+		return fmt.Errorf("tenant wallet currency %s does not match billing currency %s", currency, s.currency)
+	}
+	if input.Policy.MonthlySpendMicros > 0 {
+		period := time.Date(time.Now().UTC().Year(), time.Now().UTC().Month(), 1, 0, 0, 0, 0, time.UTC)
+		nextPeriod := period.AddDate(0, 1, 0)
+		var used, pending int64
+		if err := tx.QueryRow(ctx, `SELECT
+			COALESCE((SELECT cost_micros FROM usage_monthly_rollups WHERE project_id=$1 AND period_start=$2),0),
+			COALESCE((SELECT sum(reserved_micros) FROM billing_reservations WHERE project_id=$1 AND status='pending' AND created_at >= $2 AND created_at < $3),0)`,
+			input.Principal.ProjectID, period, nextPeriod).Scan(&used, &pending); err != nil {
+			return fmt.Errorf("read monthly spend quota: %w", err)
+		}
+		if reserved > input.Policy.MonthlySpendMicros || used > input.Policy.MonthlySpendMicros-reserved || pending > input.Policy.MonthlySpendMicros-used-reserved {
+			return ErrQuotaExceeded
+		}
+	}
+	if balance-held < reserved {
+		return ErrInsufficientBalance
+	}
+	if _, err := tx.Exec(ctx, `
+		INSERT INTO billing_reservations (
+			request_id, tenant_id, project_id, key_id, public_model, currency, reserved_micros,
+			input_price_micros_per_million, output_price_micros_per_million,
+			cache_read_price_micros_per_million, cache_write_price_micros_per_million)
+		VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11)`,
+		input.RequestID, input.Principal.TenantID, input.Principal.ProjectID, input.Principal.KeyID,
+		input.Model.ID, s.currency, reserved, input.Model.InputPriceMicrosPerMillion,
+		input.Model.OutputPriceMicrosPerMillion, input.Model.CacheReadPriceMicrosPerMillion,
+		input.Model.CacheWritePriceMicrosPerMillion); err != nil {
+		return fmt.Errorf("create billing reservation: %w", err)
+	}
+	if _, err := tx.Exec(ctx, `
+		UPDATE tenant_wallets SET reserved_micros = reserved_micros + $2, updated_at = now()
+		WHERE tenant_id = $1`, input.Principal.TenantID, reserved); err != nil {
+		return fmt.Errorf("reserve tenant balance: %w", err)
+	}
+	if err := tx.Commit(ctx); err != nil {
+		return fmt.Errorf("commit billing authorization: %w", err)
+	}
+	return nil
+}
+
+func (s *Service) Settle(ctx context.Context, event domain.UsageEvent) error {
+	tx, err := s.db.BeginTx(ctx, pgx.TxOptions{IsoLevel: pgx.ReadCommitted})
+	if err != nil {
+		return fmt.Errorf("begin usage settlement: %w", err)
+	}
+	defer tx.Rollback(ctx)
+
+	var tenantID, projectID, keyID, modelID, currency, status string
+	var reserved, inputPrice, outputPrice, cacheReadPrice, cacheWritePrice int64
+	err = tx.QueryRow(ctx, `
+		SELECT tenant_id::text, project_id::text, key_id::text, public_model, currency, reserved_micros, status,
+		       input_price_micros_per_million, output_price_micros_per_million,
+		       cache_read_price_micros_per_million, cache_write_price_micros_per_million
+		FROM billing_reservations WHERE request_id = $1 FOR UPDATE`, event.RequestID,
+	).Scan(&tenantID, &projectID, &keyID, &modelID, &currency, &reserved, &status,
+		&inputPrice, &outputPrice, &cacheReadPrice, &cacheWritePrice)
+	if err != nil {
+		return fmt.Errorf("load billing reservation: %w", err)
+	}
+	if status != "pending" {
+		return tx.Commit(ctx)
+	}
+
+	actualCost := int64(0)
+	if event.StatusCode >= 200 && event.StatusCode < 300 {
+		actualCost, err = usageCost(event.Usage, inputPrice, outputPrice, cacheReadPrice, cacheWritePrice)
+		if err != nil {
+			return err
+		}
+	}
+	var balance, held int64
+	if err := tx.QueryRow(ctx, `SELECT balance_micros, reserved_micros FROM tenant_wallets WHERE tenant_id = $1 FOR UPDATE`, tenantID).Scan(&balance, &held); err != nil {
+		return fmt.Errorf("lock wallet for settlement: %w", err)
+	}
+	charged, err := collectibleCharge(actualCost, balance, held, reserved)
+	if err != nil {
+		return err
+	}
+	uncollected := actualCost - charged
+	newBalance := balance - charged
+	if _, err := tx.Exec(ctx, `
+		UPDATE tenant_wallets SET balance_micros = $2, reserved_micros = reserved_micros - $3, updated_at = now()
+		WHERE tenant_id = $1`, tenantID, newBalance, reserved); err != nil {
+		return fmt.Errorf("settle tenant wallet: %w", err)
+	}
+	reservationStatus := "released"
+	if actualCost > 0 {
+		reservationStatus = "settled"
+	}
+	if _, err := tx.Exec(ctx, `
+		UPDATE billing_reservations
+		SET status = $2, actual_cost_micros = $3, charged_micros = $4, uncollected_micros = $5, settled_at = now()
+		WHERE request_id = $1`, event.RequestID, reservationStatus, actualCost, charged, uncollected); err != nil {
+		return fmt.Errorf("update billing reservation: %w", err)
+	}
+	usageAlreadyRecorded := false
+	if err := tx.QueryRow(ctx, `SELECT true FROM usage_events WHERE request_id=$1 FOR UPDATE`, event.RequestID).Scan(&usageAlreadyRecorded); err != nil && !errors.Is(err, pgx.ErrNoRows) {
+		return fmt.Errorf("lock existing usage event: %w", err)
+	}
+	if usageAlreadyRecorded {
+		if _, err := tx.Exec(ctx, `UPDATE usage_events SET cost_micros=$2, charged_micros=$3, uncollected_micros=$4 WHERE request_id=$1`,
+			event.RequestID, actualCost, charged, uncollected); err != nil {
+			return fmt.Errorf("apply usage charge: %w", err)
+		}
+	} else if _, err := tx.Exec(ctx, `
+		INSERT INTO usage_events (
+			request_id, tenant_id, project_id, key_id, public_model, provider_id, upstream_model,
+			protocol, stream, status_code, success, error_type, attempts, started_at, duration_ms,
+			input_tokens, output_tokens, total_tokens, cache_creation_input_tokens, cache_read_input_tokens,
+			cost_micros, charged_micros, uncollected_micros)
+		VALUES ($1,$2,$3,$4,$5,NULLIF($6,''),NULLIF($7,''),$8,$9,$10,$11,$12,$13,$14,$15,$16,$17,$18,$19,$20,$21,$22,$23)
+		ON CONFLICT (request_id) DO NOTHING`,
+		event.RequestID, tenantID, projectID, keyID, modelID, event.ProviderID, event.UpstreamModel,
+		string(event.Protocol), event.Stream, event.StatusCode, event.Success, event.ErrorType, event.Attempts,
+		event.StartedAt, event.DurationMS, event.Usage.InputTokens, event.Usage.OutputTokens,
+		event.Usage.TotalTokens, event.Usage.CacheCreationInputTokens, event.Usage.CacheReadInputTokens,
+		actualCost, charged, uncollected); err != nil {
+		return fmt.Errorf("persist usage event: %w", err)
+	}
+	if charged > 0 {
+		if _, err := tx.Exec(ctx, `
+			INSERT INTO billing_ledger (tenant_id, project_id, currency, amount_micros, balance_after_micros, kind, source_type, source_id, description)
+			VALUES ($1,$2,$3,$4,$5,'usage','request',$6,$7)
+			ON CONFLICT (source_type, source_id) DO NOTHING`, tenantID, projectID, currency, -charged, newBalance, event.RequestID, modelID); err != nil {
+			return fmt.Errorf("write usage ledger entry: %w", err)
+		}
+	}
+	period := time.Date(event.StartedAt.UTC().Year(), event.StartedAt.UTC().Month(), 1, 0, 0, 0, 0, time.UTC)
+	if usageAlreadyRecorded {
+		if _, err := tx.Exec(ctx, `UPDATE usage_monthly_rollups SET cost_micros=cost_micros+$3,
+			charged_micros=charged_micros+$4, uncollected_micros=uncollected_micros+$5, updated_at=now()
+			WHERE project_id=$1 AND period_start=$2`, projectID, period, actualCost, charged, uncollected); err != nil {
+			return fmt.Errorf("apply usage rollup charge: %w", err)
+		}
+	} else if _, err := tx.Exec(ctx, `
+		INSERT INTO usage_monthly_rollups
+		(period_start, tenant_id, project_id, request_count, successful_requests, input_tokens, output_tokens, total_tokens, cost_micros, charged_micros, uncollected_micros)
+		VALUES ($1,$2,$3,1,$4,$5,$6,$7,$8,$9,$10)
+		ON CONFLICT (project_id, period_start) DO UPDATE SET request_count=usage_monthly_rollups.request_count+1,
+			successful_requests=usage_monthly_rollups.successful_requests+EXCLUDED.successful_requests,
+			input_tokens=usage_monthly_rollups.input_tokens+EXCLUDED.input_tokens,
+			output_tokens=usage_monthly_rollups.output_tokens+EXCLUDED.output_tokens,
+			total_tokens=usage_monthly_rollups.total_tokens+EXCLUDED.total_tokens,
+			cost_micros=usage_monthly_rollups.cost_micros+EXCLUDED.cost_micros,
+			charged_micros=usage_monthly_rollups.charged_micros+EXCLUDED.charged_micros,
+			uncollected_micros=usage_monthly_rollups.uncollected_micros+EXCLUDED.uncollected_micros,
+			updated_at=now()`, period, tenantID, projectID, boolToInt(event.Success), event.Usage.InputTokens,
+		event.Usage.OutputTokens, event.Usage.TotalTokens, actualCost, charged, uncollected); err != nil {
+		return fmt.Errorf("update usage monthly rollup: %w", err)
+	}
+	if err := tx.Commit(ctx); err != nil {
+		return fmt.Errorf("commit usage settlement: %w", err)
+	}
+	return nil
+}
+
+func boolToInt(value bool) int {
+	if value {
+		return 1
+	}
+	return 0
+}
+
+func reservationCost(model domain.Model, body []byte, defaultMaxOutput int64) (int64, error) {
+	maxOutput := defaultMaxOutput
+	var limits struct {
+		MaxTokens           int64 `json:"max_tokens"`
+		MaxCompletionTokens int64 `json:"max_completion_tokens"`
+		MaxOutputTokens     int64 `json:"max_output_tokens"`
+	}
+	if json.Unmarshal(body, &limits) == nil {
+		explicitMax := int64(0)
+		for _, value := range []int64{limits.MaxTokens, limits.MaxCompletionTokens, limits.MaxOutputTokens} {
+			if value > explicitMax {
+				explicitMax = value
+			}
+		}
+		if explicitMax > 0 {
+			maxOutput = explicitMax
+		}
+	}
+	cacheReservePrice := model.CacheReadPriceMicrosPerMillion
+	if model.CacheWritePriceMicrosPerMillion > cacheReservePrice {
+		cacheReservePrice = model.CacheWritePriceMicrosPerMillion
+	}
+	return calculateCost(int64(len(body)), maxOutput, 0, int64(len(body)),
+		model.InputPriceMicrosPerMillion, model.OutputPriceMicrosPerMillion,
+		0, cacheReservePrice)
+}
+
+func usageCost(usage domain.Usage, inputPrice, outputPrice, cacheReadPrice, cacheWritePrice int64) (int64, error) {
+	return calculateCost(usage.InputTokens, usage.OutputTokens, usage.CacheReadInputTokens,
+		usage.CacheCreationInputTokens, inputPrice, outputPrice, cacheReadPrice, cacheWritePrice)
+}
+
+func calculateCost(input, output, cacheRead, cacheWrite, inputPrice, outputPrice, cacheReadPrice, cacheWritePrice int64) (int64, error) {
+	values := []int64{input, output, cacheRead, cacheWrite, inputPrice, outputPrice, cacheReadPrice, cacheWritePrice}
+	for _, value := range values {
+		if value < 0 {
+			return 0, errors.New("billing values cannot be negative")
+		}
+	}
+	total := new(big.Int)
+	for _, pair := range [][2]int64{{input, inputPrice}, {output, outputPrice}, {cacheRead, cacheReadPrice}, {cacheWrite, cacheWritePrice}} {
+		total.Add(total, new(big.Int).Mul(big.NewInt(pair[0]), big.NewInt(pair[1])))
+	}
+	if total.Sign() == 0 {
+		return 0, nil
+	}
+	total.Add(total, big.NewInt(microsPerUnit-1))
+	total.Div(total, big.NewInt(microsPerUnit))
+	if !total.IsInt64() {
+		return 0, errors.New("calculated charge exceeds supported range")
+	}
+	return total.Int64(), nil
+}
+
+func collectibleCharge(actualCost, balance, held, reservation int64) (int64, error) {
+	if actualCost < 0 || balance < 0 || held < 0 || reservation < 0 || held > balance || reservation > held {
+		return 0, errors.New("wallet reservation invariant violated")
+	}
+	spendable := balance - (held - reservation)
+	if actualCost > spendable {
+		return spendable, nil
+	}
+	return actualCost, nil
+}
+
+func randomHex(bytes int) string {
+	buffer := make([]byte, bytes)
+	if _, err := rand.Read(buffer); err != nil {
+		return fmt.Sprintf("%08x", time.Now().UnixNano())[:bytes*2]
+	}
+	return hex.EncodeToString(buffer)
+}
+
+func randomLetters(length int) string {
+	const letters = "abcdefghijklmnopqrstuvwxyz"
+	buffer := make([]byte, length)
+	if _, err := rand.Read(buffer); err != nil {
+		return strings.Repeat("a", length)
+	}
+	for i := range buffer {
+		buffer[i] = letters[int(buffer[i])%len(letters)]
+	}
+	return string(buffer)
+}
+
+func normalizeDescription(value string) string {
+	value = strings.TrimSpace(value)
+	if len(value) > 240 {
+		value = value[:240]
+	}
+	return value
+}