diff options
Diffstat (limited to 'internal/billing/service.go')
| -rw-r--r-- | internal/billing/service.go | 352 |
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(¤cy, &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, ¤cy, &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 +} |