345 lines
12 KiB
Go
345 lines
12 KiB
Go
package orchestration
|
||
|
||
import (
|
||
"context"
|
||
"encoding/json"
|
||
"errors"
|
||
"strconv"
|
||
"strings"
|
||
"time"
|
||
|
||
"git.nianxx.cn/wangxuming/NianAIGC/backend/internal/billing"
|
||
"git.nianxx.cn/wangxuming/NianAIGC/backend/internal/jobs"
|
||
)
|
||
|
||
type SettlementLedger interface {
|
||
Settle(context.Context, billing.SettlementRequest) (billing.WalletPosting, error)
|
||
}
|
||
|
||
type SettlementStateWriter interface {
|
||
WriteBilling(context.Context, string, json.RawMessage) error
|
||
}
|
||
|
||
type FencedSettlementStateWriter interface {
|
||
WriteBillingFenced(context.Context, string, json.RawMessage, jobs.Status, string) error
|
||
}
|
||
|
||
// SettlementProcessor adds actual-usage settlement for Seedance videos and
|
||
// Seedream layer decomposition. Other providers and non-successful advances
|
||
// pass through.
|
||
type SettlementProcessor struct {
|
||
inner jobs.Processor
|
||
ledger SettlementLedger
|
||
state SettlementStateWriter
|
||
now func() time.Time
|
||
}
|
||
|
||
func NewSettlementProcessor(inner jobs.Processor, ledger SettlementLedger, state SettlementStateWriter, now func() time.Time) *SettlementProcessor {
|
||
if now == nil {
|
||
now = time.Now
|
||
}
|
||
return &SettlementProcessor{inner: inner, ledger: ledger, state: state, now: now}
|
||
}
|
||
|
||
func (p *SettlementProcessor) Advance(ctx context.Context, job jobs.Job) (jobs.Job, error) {
|
||
if p == nil || p.inner == nil {
|
||
return jobs.Job{}, errors.New("advance generation job")
|
||
}
|
||
advanced, err := p.inner.Advance(ctx, job)
|
||
if err != nil || advanced.Status != jobs.StatusSucceeded || len(advanced.Billing) == 0 {
|
||
return advanced, err
|
||
}
|
||
if advanced.Provider == "seedream" && seedreamLayerRequest(advanced.RequestPayload) {
|
||
return p.settleSeedreamLayers(ctx, advanced)
|
||
}
|
||
if advanced.Provider != "seedance" {
|
||
return advanced, nil
|
||
}
|
||
|
||
var charge billingState
|
||
if json.Unmarshal(advanced.Billing, &charge) != nil {
|
||
return jobs.Job{}, errors.New("settle seedance generation charge")
|
||
}
|
||
status, _ := charge.raw["settlementStatus"].(string)
|
||
if status == "settled" || status == "estimated" {
|
||
return advanced, nil
|
||
}
|
||
var use usageContext
|
||
if len(advanced.UsageContext) != 0 && json.Unmarshal(advanced.UsageContext, &use) != nil {
|
||
return jobs.Job{}, errors.New("settle seedance generation charge")
|
||
}
|
||
quotaExempt := charge.QuotaExempt || use.Source == "platform" && use.Role == "super_admin"
|
||
chargeReady := charge.Status == "charged" || quotaExempt && charge.Status == "not_charged"
|
||
if !chargeReady || !quotaExempt && use.OrganizationID == "" {
|
||
return advanced, nil
|
||
}
|
||
|
||
completionTokens := seedanceCompletionTokens(advanced.ResponsePayload)
|
||
if completionTokens <= 0 {
|
||
charge.raw["settlementStatus"] = "estimated"
|
||
charge.raw["settlementReason"] = "provider_usage_unavailable"
|
||
charge.raw["settledAt"] = p.now().UTC().Format(time.RFC3339Nano)
|
||
return p.write(ctx, advanced, charge.raw)
|
||
}
|
||
|
||
resolution := seedanceResolution(charge.raw)
|
||
inputVideo := seedanceRequestHasInputVideo(advanced.RequestPayload)
|
||
actualAmount, err := billing.CalculateSeedanceActualAmountFen(billing.SeedanceActualAmountInput{
|
||
Model: advanced.ReqKey, Resolution: resolution, InputVideo: inputVideo, CompletionTokens: completionTokens,
|
||
MarkupMultiplier: numberOrZero(charge.raw["markupMultiplier"]),
|
||
})
|
||
if err != nil {
|
||
return jobs.Job{}, errors.New("settle seedance generation charge")
|
||
}
|
||
delta := actualAmount - charge.AmountFen
|
||
settledAt := p.now().UTC()
|
||
var settlementLedgerID string
|
||
if delta != 0 && !quotaExempt {
|
||
if p.ledger == nil {
|
||
return jobs.Job{}, errors.New("settle seedance generation charge")
|
||
}
|
||
description := capabilityLabel(advanced.Capability) + "实际用量差额退回"
|
||
if delta > 0 {
|
||
description = capabilityLabel(advanced.Capability) + "实际用量补扣"
|
||
}
|
||
reservedAmount := charge.AmountFen
|
||
if value, ok := integer(charge.raw["reservedAmountFen"]); ok {
|
||
reservedAmount = value
|
||
}
|
||
posting, postErr := p.ledger.Settle(ctx, billing.SettlementRequest{
|
||
OrganizationID: use.OrganizationID, AccountID: use.AccountID, JobID: advanced.ID,
|
||
DeltaFen: delta, Description: description,
|
||
Metadata: map[string]any{
|
||
"operation": "seedance_actual_settlement", "reservedAmountFen": reservedAmount,
|
||
"chargedAmountFen": charge.AmountFen, "actualAmountFen": actualAmount,
|
||
"completionTokens": completionTokens, "inputVideo": inputVideo, "resolution": resolution, "model": advanced.ReqKey,
|
||
},
|
||
})
|
||
if postErr != nil {
|
||
return jobs.Job{}, errors.New("settle seedance generation charge")
|
||
}
|
||
settlementLedgerID = posting.LedgerID
|
||
if !posting.CreatedAt.IsZero() {
|
||
settledAt = posting.CreatedAt.UTC()
|
||
}
|
||
}
|
||
|
||
charge.raw["amountFen"] = actualAmount
|
||
charge.raw["settlementStatus"] = "settled"
|
||
charge.raw["settledAt"] = settledAt.Format(time.RFC3339Nano)
|
||
if settlementLedgerID != "" {
|
||
charge.raw["settlementLedgerEntryId"] = settlementLedgerID
|
||
} else {
|
||
delete(charge.raw, "settlementLedgerEntryId")
|
||
}
|
||
charge.raw["providerUsage"] = map[string]any{
|
||
"completionTokens": completionTokens, "resolution": resolution, "inputVideo": inputVideo, "model": advanced.ReqKey,
|
||
"tokenPriceFenPerMillion": billing.SeedanceTokenPriceFenPerMillion(advanced.ReqKey, resolution, inputVideo),
|
||
}
|
||
return p.write(ctx, advanced, charge.raw)
|
||
}
|
||
|
||
func (p *SettlementProcessor) settleSeedreamLayers(ctx context.Context, job jobs.Job) (jobs.Job, error) {
|
||
var charge billingState
|
||
if json.Unmarshal(job.Billing, &charge) != nil {
|
||
return jobs.Job{}, errors.New("settle seedream layer charge")
|
||
}
|
||
status, _ := charge.raw["settlementStatus"].(string)
|
||
if status == "settled" || status == "estimated" {
|
||
return job, nil
|
||
}
|
||
var use usageContext
|
||
if len(job.UsageContext) != 0 && json.Unmarshal(job.UsageContext, &use) != nil {
|
||
return jobs.Job{}, errors.New("settle seedream layer charge")
|
||
}
|
||
quotaExempt := charge.QuotaExempt || use.Source == "platform" && use.Role == "super_admin"
|
||
chargeReady := charge.Status == "charged" || quotaExempt && charge.Status == "not_charged"
|
||
if !chargeReady || !quotaExempt && use.OrganizationID == "" {
|
||
return job, nil
|
||
}
|
||
|
||
images := seedreamLayerImages(job.ResponsePayload)
|
||
if len(images) == 0 {
|
||
charge.raw["settlementStatus"] = "estimated"
|
||
charge.raw["settlementReason"] = "provider_layer_metadata_unavailable"
|
||
charge.raw["settledAt"] = p.now().UTC().Format(time.RFC3339Nano)
|
||
return p.write(ctx, job, charge.raw)
|
||
}
|
||
actualAmount, err := billing.CalculateSeedreamLayerAmountFen(images, numberOrZero(charge.raw["markupMultiplier"]))
|
||
if err != nil {
|
||
return jobs.Job{}, errors.New("settle seedream layer charge")
|
||
}
|
||
delta := actualAmount - charge.AmountFen
|
||
settledAt := p.now().UTC()
|
||
var settlementLedgerID string
|
||
if delta != 0 && !quotaExempt {
|
||
if p.ledger == nil {
|
||
return jobs.Job{}, errors.New("settle seedream layer charge")
|
||
}
|
||
description := capabilityLabel(job.Capability) + "图层拆分实际费用退回"
|
||
if delta > 0 {
|
||
description = capabilityLabel(job.Capability) + "图层拆分实际费用补扣"
|
||
}
|
||
posting, postErr := p.ledger.Settle(ctx, billing.SettlementRequest{
|
||
OrganizationID: use.OrganizationID, AccountID: use.AccountID, JobID: job.ID,
|
||
DeltaFen: delta, Description: description,
|
||
Metadata: map[string]any{
|
||
"operation": "seedream_layer_actual_settlement", "reservedAmountFen": charge.AmountFen,
|
||
"actualAmountFen": actualAmount, "outputCount": len(images), "model": job.ReqKey,
|
||
},
|
||
})
|
||
if postErr != nil {
|
||
return jobs.Job{}, errors.New("settle seedream layer charge")
|
||
}
|
||
settlementLedgerID = posting.LedgerID
|
||
if !posting.CreatedAt.IsZero() {
|
||
settledAt = posting.CreatedAt.UTC()
|
||
}
|
||
}
|
||
charge.raw["amountFen"] = actualAmount
|
||
charge.raw["settlementStatus"] = "settled"
|
||
charge.raw["settledAt"] = settledAt.Format(time.RFC3339Nano)
|
||
if settlementLedgerID != "" {
|
||
charge.raw["settlementLedgerEntryId"] = settlementLedgerID
|
||
} else {
|
||
delete(charge.raw, "settlementLedgerEntryId")
|
||
}
|
||
charge.raw["providerUsage"] = map[string]any{
|
||
"model": job.ReqKey, "layerDecomposition": true, "generatedImages": len(images),
|
||
}
|
||
return p.write(ctx, job, charge.raw)
|
||
}
|
||
|
||
func (p *SettlementProcessor) write(ctx context.Context, job jobs.Job, snapshot map[string]any) (jobs.Job, error) {
|
||
if p.state == nil {
|
||
return jobs.Job{}, errors.New("persist generation settlement")
|
||
}
|
||
encoded, err := json.Marshal(snapshot)
|
||
if err != nil {
|
||
return jobs.Job{}, errors.New("persist generation settlement")
|
||
}
|
||
var writeErr error
|
||
if fenced, ok := p.state.(FencedSettlementStateWriter); ok && job.LockedBy != "" {
|
||
writeErr = fenced.WriteBillingFenced(ctx, job.ID, encoded, job.Status, job.LockedBy)
|
||
} else {
|
||
writeErr = p.state.WriteBilling(ctx, job.ID, encoded)
|
||
}
|
||
if writeErr != nil {
|
||
return jobs.Job{}, errors.New("persist seedance generation settlement")
|
||
}
|
||
job.Billing = encoded
|
||
return job, nil
|
||
}
|
||
|
||
func seedanceCompletionTokens(payload json.RawMessage) int64 {
|
||
var response struct {
|
||
Usage map[string]any `json:"usage"`
|
||
}
|
||
if json.Unmarshal(payload, &response) != nil {
|
||
return 0
|
||
}
|
||
value, _ := integer(response.Usage["completionTokens"])
|
||
return value
|
||
}
|
||
|
||
func seedreamLayerRequest(payload json.RawMessage) bool {
|
||
var request struct {
|
||
Settings map[string]any `json:"settings"`
|
||
}
|
||
if json.Unmarshal(payload, &request) != nil {
|
||
return false
|
||
}
|
||
value, _ := request.Settings["layerDecomposition"].(bool)
|
||
creationMode, _ := request.Settings["creationMode"].(string)
|
||
return value || strings.EqualFold(strings.TrimSpace(creationMode), "layers")
|
||
}
|
||
|
||
func seedreamLayerImages(payload json.RawMessage) []billing.SeedreamLayerImage {
|
||
var response struct {
|
||
Raw json.RawMessage `json:"raw"`
|
||
}
|
||
if json.Unmarshal(payload, &response) != nil || len(response.Raw) == 0 {
|
||
return nil
|
||
}
|
||
var providerResponse struct {
|
||
Data []struct {
|
||
Size string `json:"size"`
|
||
} `json:"data"`
|
||
}
|
||
if json.Unmarshal(response.Raw, &providerResponse) != nil || len(providerResponse.Data) == 0 || len(providerResponse.Data) > 17 {
|
||
return nil
|
||
}
|
||
images := make([]billing.SeedreamLayerImage, 0, len(providerResponse.Data))
|
||
for _, output := range providerResponse.Data {
|
||
parts := strings.FieldsFunc(strings.TrimSpace(output.Size), func(character rune) bool {
|
||
return character == 'x' || character == 'X' || character == '×' || character == '*'
|
||
})
|
||
if len(parts) != 2 {
|
||
return nil
|
||
}
|
||
width, widthErr := strconv.Atoi(strings.TrimSpace(parts[0]))
|
||
height, heightErr := strconv.Atoi(strings.TrimSpace(parts[1]))
|
||
if widthErr != nil || heightErr != nil || width <= 0 || height <= 0 {
|
||
return nil
|
||
}
|
||
images = append(images, billing.SeedreamLayerImage{Width: width, Height: height})
|
||
}
|
||
return images
|
||
}
|
||
|
||
func seedanceResolution(snapshot map[string]any) string {
|
||
parameters, _ := snapshot["parameters"].(map[string]any)
|
||
resolution, _ := parameters["resolution"].(string)
|
||
resolution = strings.ToLower(strings.TrimSpace(resolution))
|
||
if resolution == "480p" || resolution == "1080p" || resolution == "4k" {
|
||
return resolution
|
||
}
|
||
return "720p"
|
||
}
|
||
|
||
func seedanceRequestHasInputVideo(payload json.RawMessage) bool {
|
||
var value any
|
||
if json.Unmarshal(payload, &value) != nil {
|
||
return false
|
||
}
|
||
return hasVideoMaterial(value)
|
||
}
|
||
|
||
func hasVideoMaterial(value any) bool {
|
||
switch typed := value.(type) {
|
||
case map[string]any:
|
||
if materialType, _ := typed["type"].(string); strings.EqualFold(strings.TrimSpace(materialType), "video") {
|
||
return true
|
||
}
|
||
for _, child := range typed {
|
||
if hasVideoMaterial(child) {
|
||
return true
|
||
}
|
||
}
|
||
case []any:
|
||
for _, child := range typed {
|
||
if hasVideoMaterial(child) {
|
||
return true
|
||
}
|
||
}
|
||
}
|
||
return false
|
||
}
|
||
|
||
func numberOrZero(value any) float64 {
|
||
if number, ok := value.(float64); ok {
|
||
return number
|
||
}
|
||
return 0
|
||
}
|
||
|
||
func integer(value any) (int64, bool) {
|
||
number, ok := value.(float64)
|
||
if !ok || number <= 0 {
|
||
return 0, false
|
||
}
|
||
return int64(number), true
|
||
}
|
||
|
||
var _ jobs.Processor = (*SettlementProcessor)(nil)
|