初始化
This commit is contained in:
@@ -0,0 +1,59 @@
|
||||
package billing
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"strconv"
|
||||
|
||||
"juhe-factory/api/internal/model"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
// PrechargeGenerationTask creates a task and reserves its full quoted cost in
|
||||
// the same transaction. A zero price is allowed only when the model has an
|
||||
// explicit zero-price record; callers must reject a missing price record.
|
||||
func PrechargeGenerationTask(tx *gorm.DB, task *model.GenerationTask, unitPrice string, quantity int, remark string) error {
|
||||
price, err := strconv.ParseFloat(unitPrice, 64)
|
||||
if err != nil || price < 0 || quantity <= 0 {
|
||||
return errors.New("生成任务价格无效")
|
||||
}
|
||||
var amount string
|
||||
if err := tx.Raw("SELECT round(?::numeric * ?::numeric,2)::text", unitPrice, quantity).Scan(&amount).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if amount == "" {
|
||||
return errors.New("生成任务价格无效")
|
||||
}
|
||||
task.EstimatedPoints = amount
|
||||
task.PrepaidPoints = amount
|
||||
if task.ID == uuid.Nil {
|
||||
task.ID = uuid.New()
|
||||
}
|
||||
|
||||
if amountValue, _ := strconv.ParseFloat(amount, 64); amountValue > 0 {
|
||||
if _, err := DebitPoints(tx, task.UserID, amount, "generation_hold", task.ID.String(), remark, "generation:hold:"+task.ID.String()); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
if err := tx.Create(task).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// RefundGenerationTask refunds a prepaid task exactly once. The caller must
|
||||
// hold a row lock on task inside the same transaction.
|
||||
func RefundGenerationTask(tx *gorm.DB, task *model.GenerationTask, remark string) (bool, error) {
|
||||
amount, _ := strconv.ParseFloat(task.PrepaidPoints, 64)
|
||||
if task.CostRefunded || amount <= 0 {
|
||||
return false, nil
|
||||
}
|
||||
refunded, err := RefundDebit(tx, task.UserID, "generation:hold:"+task.ID.String(), "generation:refund:"+task.ID.String(),
|
||||
"generation_refund", task.ID.String(), remark)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
task.CostRefunded = true
|
||||
return refunded, nil
|
||||
}
|
||||
@@ -0,0 +1,250 @@
|
||||
// 积分账务模块,负责积分入账、扣减、退款及批次分配记录维护。
|
||||
package billing
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
var ErrInsufficientPoints = errors.New("积分余额不足,请前往充值")
|
||||
|
||||
// pointGrant 表示参与扣减的可用积分批次。
|
||||
type pointGrant struct {
|
||||
ID uuid.UUID
|
||||
AvailableAmount string
|
||||
}
|
||||
|
||||
// pointAllocation 表示原扣减流水关联的积分批次和分配金额。
|
||||
type pointAllocation struct {
|
||||
GrantID uuid.UUID
|
||||
Amount string
|
||||
}
|
||||
|
||||
// parsePointCents 将最多两位小数的积分字符串转换为整数分值。
|
||||
func parsePointCents(value string) (int64, error) {
|
||||
value = strings.TrimSpace(value)
|
||||
if value == "" {
|
||||
return 0, errors.New("积分数量为空")
|
||||
}
|
||||
parts := strings.Split(value, ".")
|
||||
if len(parts) > 2 || strings.HasPrefix(parts[0], "-") {
|
||||
return 0, fmt.Errorf("无效积分数量 %q", value)
|
||||
}
|
||||
whole, err := strconv.ParseInt(parts[0], 10, 64)
|
||||
if err != nil {
|
||||
return 0, fmt.Errorf("无效积分数量 %q", value)
|
||||
}
|
||||
fraction := ""
|
||||
if len(parts) == 2 {
|
||||
fraction = parts[1]
|
||||
}
|
||||
if len(fraction) > 2 {
|
||||
if strings.Trim(fraction[2:], "0") != "" {
|
||||
return 0, fmt.Errorf("积分数量最多保留两位小数")
|
||||
}
|
||||
fraction = fraction[:2]
|
||||
}
|
||||
fraction += strings.Repeat("0", 2-len(fraction))
|
||||
fractionValue, err := strconv.ParseInt(fraction, 10, 64)
|
||||
if err != nil || whole > (int64(^uint64(0)>>1)-fractionValue)/100 {
|
||||
return 0, fmt.Errorf("无效积分数量 %q", value)
|
||||
}
|
||||
return whole*100 + fractionValue, nil
|
||||
}
|
||||
|
||||
// formatPointCents 将整数分值格式化为两位小数的积分字符串。
|
||||
func formatPointCents(value int64) string {
|
||||
return fmt.Sprintf("%d.%02d", value/100, value%100)
|
||||
}
|
||||
|
||||
// lockUser 锁定用户积分余额,避免并发账务操作导致余额和批次不一致。
|
||||
func lockUser(tx *gorm.DB, userID uuid.UUID, requireActive bool) (string, error) {
|
||||
var balance string
|
||||
userQuery := "SELECT point_balance::text FROM web_users WHERE id=? FOR UPDATE"
|
||||
if requireActive {
|
||||
userQuery = "SELECT point_balance::text FROM web_users WHERE id=? AND enabled=true AND deleted_at IS NULL FOR UPDATE"
|
||||
}
|
||||
if err := tx.Raw(userQuery, userID).Scan(&balance).Error; err != nil {
|
||||
return "", err
|
||||
}
|
||||
if balance == "" {
|
||||
return "", gorm.ErrRecordNotFound
|
||||
}
|
||||
return balance, nil
|
||||
}
|
||||
|
||||
// CreditPoints 创建永久有效的积分批次,并原子写入用户余额和积分流水。
|
||||
func CreditPoints(tx *gorm.DB, userID uuid.UUID, sourceType, sourceID, points, remark string) (string, bool, error) {
|
||||
amount, err := parsePointCents(points)
|
||||
if err != nil || amount <= 0 {
|
||||
return "", false, errors.New("入账积分必须大于零")
|
||||
}
|
||||
balance, err := lockUser(tx, userID, true)
|
||||
if err != nil {
|
||||
return "", false, err
|
||||
}
|
||||
|
||||
grantID := uuid.New()
|
||||
result := tx.Exec(`INSERT INTO point_grants(id,user_id,source_type,source_id,granted_amount,available_amount)
|
||||
VALUES(?,?,?,?,?::numeric,?::numeric) ON CONFLICT(source_type,source_id) DO NOTHING`,
|
||||
grantID, userID, sourceType, sourceID, points, points)
|
||||
if result.Error != nil {
|
||||
return "", false, result.Error
|
||||
}
|
||||
if result.RowsAffected == 0 {
|
||||
return balance, false, nil
|
||||
}
|
||||
if err := tx.Raw(`UPDATE web_users SET point_balance=point_balance+?::numeric
|
||||
WHERE id=? RETURNING point_balance::text`, points, userID).Scan(&balance).Error; err != nil {
|
||||
return "", false, err
|
||||
}
|
||||
if err := tx.Exec(`INSERT INTO point_ledger(user_id,change_amount,balance_after,business_type,business_id,remark,idempotency_key)
|
||||
VALUES(?,?::numeric,?::numeric,?,?,?,?)`, userID, points, balance, sourceType, sourceID, remark, sourceType+":"+sourceID).Error; err != nil {
|
||||
return "", false, err
|
||||
}
|
||||
return balance, true, nil
|
||||
}
|
||||
|
||||
// CreditRedemptionPoints 将兑换码对应积分以永久有效批次计入用户账户。
|
||||
func CreditRedemptionPoints(tx *gorm.DB, userID, codeID uuid.UUID, points string) (string, bool, error) {
|
||||
return CreditPoints(tx, userID, "redemption", codeID.String(), points, "兑换码兑换")
|
||||
}
|
||||
|
||||
// DebitPoints 按批次创建顺序扣减积分,并记录分配明细供后续退款恢复。
|
||||
func DebitPoints(tx *gorm.DB, userID uuid.UUID, points, businessType, businessID, remark, idempotencyKey string) (string, error) {
|
||||
requested, err := parsePointCents(points)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
balance, err := lockUser(tx, userID, true)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if requested == 0 {
|
||||
return balance, nil
|
||||
}
|
||||
|
||||
var grants []pointGrant
|
||||
if err := tx.Raw(`SELECT id,available_amount::text AS available_amount
|
||||
FROM point_grants WHERE user_id=? AND available_amount>0
|
||||
ORDER BY created_at,id FOR UPDATE`, userID).Scan(&grants).Error; err != nil {
|
||||
return "", err
|
||||
}
|
||||
remaining := requested
|
||||
type usedGrant struct {
|
||||
id uuid.UUID
|
||||
amount int64
|
||||
}
|
||||
used := make([]usedGrant, 0, len(grants))
|
||||
for _, grant := range grants {
|
||||
available, parseErr := parsePointCents(grant.AvailableAmount)
|
||||
if parseErr != nil {
|
||||
return "", parseErr
|
||||
}
|
||||
use := available
|
||||
if use > remaining {
|
||||
use = remaining
|
||||
}
|
||||
if use > 0 {
|
||||
used = append(used, usedGrant{id: grant.ID, amount: use})
|
||||
remaining -= use
|
||||
}
|
||||
if remaining == 0 {
|
||||
break
|
||||
}
|
||||
}
|
||||
if remaining > 0 {
|
||||
return "", ErrInsufficientPoints
|
||||
}
|
||||
for _, allocation := range used {
|
||||
if err := tx.Exec(`UPDATE point_grants SET available_amount=available_amount-?::numeric WHERE id=?`,
|
||||
formatPointCents(allocation.amount), allocation.id).Error; err != nil {
|
||||
return "", err
|
||||
}
|
||||
}
|
||||
if err := tx.Raw(`UPDATE web_users SET point_balance=point_balance-?::numeric
|
||||
WHERE id=? RETURNING point_balance::text`, points, userID).Scan(&balance).Error; err != nil {
|
||||
return "", err
|
||||
}
|
||||
var ledgerID int64
|
||||
if err := tx.Raw(`INSERT INTO point_ledger(user_id,change_amount,balance_after,business_type,business_id,remark,idempotency_key)
|
||||
VALUES(?,(-?::numeric),?::numeric,?,?,?,?) RETURNING id`,
|
||||
userID, points, balance, businessType, businessID, remark, idempotencyKey).Scan(&ledgerID).Error; err != nil {
|
||||
return "", err
|
||||
}
|
||||
for _, allocation := range used {
|
||||
if err := tx.Exec(`INSERT INTO point_ledger_allocations(ledger_id,grant_id,allocation_type,amount)
|
||||
VALUES(?,?,'consume',?::numeric)`, ledgerID, allocation.id, formatPointCents(allocation.amount)).Error; err != nil {
|
||||
return "", err
|
||||
}
|
||||
}
|
||||
return balance, nil
|
||||
}
|
||||
|
||||
// RefundDebit 将原扣减记录恢复到对应的永久积分批次,并写入退款流水。
|
||||
func RefundDebit(tx *gorm.DB, userID uuid.UUID, originalKey, refundKey, businessType, businessID, remark string) (bool, error) {
|
||||
balance, err := lockUser(tx, userID, false)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
var existing int64
|
||||
if err := tx.Raw("SELECT count(*) FROM point_ledger WHERE idempotency_key=?", refundKey).Scan(&existing).Error; err != nil {
|
||||
return false, err
|
||||
}
|
||||
if existing > 0 {
|
||||
return false, nil
|
||||
}
|
||||
|
||||
var allocations []pointAllocation
|
||||
if err := tx.Raw(`SELECT a.grant_id,a.amount::text AS amount
|
||||
FROM point_ledger original
|
||||
JOIN point_ledger_allocations a ON a.ledger_id=original.id AND a.allocation_type='consume'
|
||||
JOIN point_grants g ON g.id=a.grant_id
|
||||
WHERE original.idempotency_key=? ORDER BY g.created_at,g.id FOR UPDATE OF g`, originalKey).Scan(&allocations).Error; err != nil {
|
||||
return false, err
|
||||
}
|
||||
if len(allocations) == 0 {
|
||||
return false, errors.New("未找到原积分扣减分配记录")
|
||||
}
|
||||
restored := int64(0)
|
||||
restoredAllocations := make([]pointAllocation, 0, len(allocations))
|
||||
for _, allocation := range allocations {
|
||||
amount, parseErr := parsePointCents(allocation.Amount)
|
||||
if parseErr != nil {
|
||||
return false, parseErr
|
||||
}
|
||||
result := tx.Exec(`UPDATE point_grants SET available_amount=available_amount+?::numeric
|
||||
WHERE id=?`, allocation.Amount, allocation.GrantID)
|
||||
if result.Error != nil {
|
||||
return false, result.Error
|
||||
}
|
||||
if result.RowsAffected > 0 {
|
||||
restored += amount
|
||||
restoredAllocations = append(restoredAllocations, allocation)
|
||||
}
|
||||
}
|
||||
if restored > 0 {
|
||||
if err := tx.Raw(`UPDATE web_users SET point_balance=point_balance+?::numeric
|
||||
WHERE id=? RETURNING point_balance::text`, formatPointCents(restored), userID).Scan(&balance).Error; err != nil {
|
||||
return false, err
|
||||
}
|
||||
}
|
||||
var refundLedgerID int64
|
||||
if err := tx.Raw(`INSERT INTO point_ledger(user_id,change_amount,balance_after,business_type,business_id,remark,idempotency_key)
|
||||
VALUES(?,?::numeric,?::numeric,?,?,?,?) RETURNING id`, userID, formatPointCents(restored), balance,
|
||||
businessType, businessID, remark, refundKey).Scan(&refundLedgerID).Error; err != nil {
|
||||
return false, err
|
||||
}
|
||||
for _, allocation := range restoredAllocations {
|
||||
if err := tx.Exec(`INSERT INTO point_ledger_allocations(ledger_id,grant_id,allocation_type,amount)
|
||||
VALUES(?,?,'refund',?::numeric)`, refundLedgerID, allocation.GrantID, allocation.Amount).Error; err != nil {
|
||||
return false, err
|
||||
}
|
||||
}
|
||||
return true, nil
|
||||
}
|
||||
@@ -0,0 +1,30 @@
|
||||
// 积分账务单元测试,验证积分精度解析和非法输入处理。
|
||||
package billing
|
||||
|
||||
import "testing"
|
||||
|
||||
// TestParsePointCents 验证积分字符串能够稳定转换为整数分值。
|
||||
func TestParsePointCents(t *testing.T) {
|
||||
tests := map[string]int64{
|
||||
"0": 0,
|
||||
"1": 100,
|
||||
"1.2": 120,
|
||||
"1.23": 123,
|
||||
"1.2300": 123,
|
||||
"999.99": 99999,
|
||||
}
|
||||
for input, expected := range tests {
|
||||
actual, err := parsePointCents(input)
|
||||
if err != nil {
|
||||
t.Fatalf("parse %q: %v", input, err)
|
||||
}
|
||||
if actual != expected {
|
||||
t.Fatalf("parse %q: expected %d, got %d", input, expected, actual)
|
||||
}
|
||||
}
|
||||
for _, input := range []string{"", "-1", "1.234", "abc"} {
|
||||
if _, err := parsePointCents(input); err == nil {
|
||||
t.Fatalf("expected %q to be rejected", input)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,321 @@
|
||||
package billing
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"math"
|
||||
"math/big"
|
||||
"strings"
|
||||
"unicode"
|
||||
|
||||
"juhe-factory/api/internal/model"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
const (
|
||||
TextBillingPerRequest = "per_request"
|
||||
TextBillingPerToken = "per_token"
|
||||
)
|
||||
|
||||
type TextPricing struct {
|
||||
Mode string `json:"mode"`
|
||||
PerRequest string `json:"per_request,omitempty"`
|
||||
InputPerM string `json:"input_per_m,omitempty"`
|
||||
OutputPerM string `json:"output_per_m,omitempty"`
|
||||
}
|
||||
|
||||
type TextUsage struct {
|
||||
Input int64
|
||||
Output int64
|
||||
}
|
||||
|
||||
func CreateTextGenerationTask(tx *gorm.DB, task *model.GenerationTask, pricing TextPricing, remark string) error {
|
||||
if err := pricing.Validate(); err != nil {
|
||||
return err
|
||||
}
|
||||
task.BillingSnapshot = pricing.MarshalSnapshot()
|
||||
if pricing.Mode == TextBillingPerRequest {
|
||||
return PrechargeGenerationTask(tx, task, pricing.PerRequest, 1, remark)
|
||||
}
|
||||
task.EstimatedPoints = "0.00"
|
||||
task.PrepaidPoints = "0.00"
|
||||
if task.ID == uuid.Nil {
|
||||
task.ID = uuid.New()
|
||||
}
|
||||
return tx.Create(task).Error
|
||||
}
|
||||
|
||||
func ChargeTextRequest(tx *gorm.DB, userID uuid.UUID, referenceID string, pricing TextPricing, remark string) (string, error) {
|
||||
if pricing.Mode != TextBillingPerRequest {
|
||||
return "", errors.New("文本模型不是按次计费")
|
||||
}
|
||||
businessID := textCallBusinessID(referenceID, "request")
|
||||
idempotencyKey := "generation:hold:" + businessID
|
||||
if existing, found, err := debitAmount(tx, idempotencyKey); err != nil {
|
||||
return "", err
|
||||
} else if found {
|
||||
return existing, nil
|
||||
}
|
||||
if pricing.PerRequest == "0" || pricing.PerRequest == "0.00" || pricing.PerRequest == "0.0000" {
|
||||
return "0.00", nil
|
||||
}
|
||||
_, err := DebitPoints(tx, userID, pricing.PerRequest, "generation_hold", businessID, remark, idempotencyKey)
|
||||
return pricing.PerRequest, err
|
||||
}
|
||||
|
||||
func RefundTextGenerationTask(tx *gorm.DB, task *model.GenerationTask, remark string) (bool, error) {
|
||||
pricing, err := ParseTextPricingSnapshot(task.BillingSnapshot)
|
||||
if err != nil || pricing.Mode == TextBillingPerRequest {
|
||||
return RefundGenerationTask(tx, task, remark)
|
||||
}
|
||||
if task.CostRefunded {
|
||||
return false, nil
|
||||
}
|
||||
refunded, err := RefundTextCalls(tx, task.UserID, task.ID.String(), remark)
|
||||
if err == nil {
|
||||
task.CostRefunded = true
|
||||
}
|
||||
return refunded, err
|
||||
}
|
||||
|
||||
func LoadTextPricing(db *gorm.DB, modelID uuid.UUID) (TextPricing, error) {
|
||||
var pricing TextPricing
|
||||
err := db.Table("models m").
|
||||
Select(`coalesce(m.text_billing_mode,'') AS mode,
|
||||
coalesce(max(CASE WHEN lower(p.price_key)='default' THEN p.price::text END),'') AS per_request,
|
||||
coalesce(max(CASE WHEN lower(p.price_key)='input' THEN p.price::text END),'') AS input_per_m,
|
||||
coalesce(max(CASE WHEN lower(p.price_key)='output' THEN p.price::text END),'') AS output_per_m`).
|
||||
Joins("LEFT JOIN model_prices p ON p.model_id=m.id").
|
||||
Where("m.id=? AND m.model_type='text' AND m.enabled=true AND m.deleted_at IS NULL", modelID).
|
||||
Group("m.id,m.text_billing_mode").
|
||||
Take(&pricing).Error
|
||||
if err != nil {
|
||||
return pricing, err
|
||||
}
|
||||
if err := pricing.Validate(); err != nil {
|
||||
return pricing, err
|
||||
}
|
||||
return pricing, nil
|
||||
}
|
||||
|
||||
func (pricing TextPricing) Validate() error {
|
||||
switch pricing.Mode {
|
||||
case TextBillingPerRequest:
|
||||
if !validNonNegativeDecimal(pricing.PerRequest) {
|
||||
return errors.New("文本模型未配置按次价格")
|
||||
}
|
||||
case TextBillingPerToken:
|
||||
if !validNonNegativeDecimal(pricing.InputPerM) || !validNonNegativeDecimal(pricing.OutputPerM) {
|
||||
return errors.New("文本模型未完整配置输入、输出 Token 价格")
|
||||
}
|
||||
default:
|
||||
return errors.New("文本模型计费模式无效")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (pricing TextPricing) MarshalSnapshot() json.RawMessage {
|
||||
data, _ := json.Marshal(pricing)
|
||||
return data
|
||||
}
|
||||
|
||||
func ParseTextPricingSnapshot(value json.RawMessage) (TextPricing, error) {
|
||||
var pricing TextPricing
|
||||
if len(value) == 0 || string(value) == "{}" {
|
||||
return pricing, errors.New("文本模型计费快照缺失")
|
||||
}
|
||||
if err := json.Unmarshal(value, &pricing); err != nil {
|
||||
return pricing, errors.New("文本模型计费快照无效")
|
||||
}
|
||||
return pricing, pricing.Validate()
|
||||
}
|
||||
|
||||
func CalculateTextPoints(pricing TextPricing, usage TextUsage) (string, error) {
|
||||
if pricing.Mode != TextBillingPerToken || usage.Input < 0 || usage.Output < 0 {
|
||||
return "", errors.New("Token 计费参数无效")
|
||||
}
|
||||
inputRate, ok := new(big.Rat).SetString(pricing.InputPerM)
|
||||
if !ok || inputRate.Sign() < 0 {
|
||||
return "", errors.New("输入 Token 价格无效")
|
||||
}
|
||||
outputRate, ok := new(big.Rat).SetString(pricing.OutputPerM)
|
||||
if !ok || outputRate.Sign() < 0 {
|
||||
return "", errors.New("输出 Token 价格无效")
|
||||
}
|
||||
amount := new(big.Rat).Mul(inputRate, new(big.Rat).SetInt64(usage.Input))
|
||||
amount.Add(amount, new(big.Rat).Mul(outputRate, new(big.Rat).SetInt64(usage.Output)))
|
||||
amount.Quo(amount, big.NewRat(1000000, 1))
|
||||
return ceilPointCents(amount), nil
|
||||
}
|
||||
|
||||
func EstimateTextTokens(value string) int64 {
|
||||
var tokens int64
|
||||
ascii := 0
|
||||
flush := func() {
|
||||
if ascii > 0 {
|
||||
tokens += int64((ascii + 3) / 4)
|
||||
ascii = 0
|
||||
}
|
||||
}
|
||||
for _, r := range value {
|
||||
if r <= unicode.MaxASCII {
|
||||
ascii++
|
||||
} else {
|
||||
flush()
|
||||
tokens++
|
||||
}
|
||||
}
|
||||
flush()
|
||||
if tokens == 0 && value != "" {
|
||||
return 1
|
||||
}
|
||||
return tokens
|
||||
}
|
||||
|
||||
func ReserveTextCall(tx *gorm.DB, userID uuid.UUID, referenceID, callKey string, pricing TextPricing, estimatedInputTokens int64, remark string) (string, error) {
|
||||
if pricing.Mode != TextBillingPerToken || estimatedInputTokens < 0 {
|
||||
return "", errors.New("Token 预授权参数无效")
|
||||
}
|
||||
amount, err := CalculateTextReservePoints(pricing, estimatedInputTokens)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
businessID := textCallBusinessID(referenceID, callKey) + ":reserve"
|
||||
idempotencyKey := "text:reserve:" + businessID
|
||||
if existing, found, err := debitAmount(tx, idempotencyKey); err != nil {
|
||||
return "", err
|
||||
} else if found {
|
||||
return existing, nil
|
||||
}
|
||||
if amount == "0.00" {
|
||||
return amount, nil
|
||||
}
|
||||
_, err = DebitPoints(tx, userID, amount, "text_token_reserve", businessID, remark+"预授权", idempotencyKey)
|
||||
return amount, err
|
||||
}
|
||||
|
||||
func CalculateTextReservePoints(pricing TextPricing, estimatedInputTokens int64) (string, error) {
|
||||
if estimatedInputTokens < 0 || estimatedInputTokens > math.MaxInt64/2 {
|
||||
return "", errors.New("Token 预授权参数无效")
|
||||
}
|
||||
return CalculateTextPoints(pricing, TextUsage{Input: estimatedInputTokens, Output: estimatedInputTokens * 2})
|
||||
}
|
||||
|
||||
func SettleTextCall(tx *gorm.DB, userID uuid.UUID, referenceID, callKey string, pricing TextPricing, usage TextUsage, remark string) (string, error) {
|
||||
if usage.Input <= 0 && usage.Output <= 0 {
|
||||
return "", errors.New("上游未返回 Token 用量,无法完成计费")
|
||||
}
|
||||
actual, err := CalculateTextPoints(pricing, usage)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
baseID := textCallBusinessID(referenceID, callKey)
|
||||
reserveID := baseID + ":reserve"
|
||||
if _, found, err := debitAmount(tx, "text:reserve:"+reserveID); err != nil {
|
||||
return "", err
|
||||
} else if found {
|
||||
if _, err := RefundDebit(tx, userID, "text:reserve:"+reserveID, "text:reserve-refund:"+reserveID,
|
||||
"text_token_reserve_refund", reserveID, remark+"预授权释放"); err != nil {
|
||||
return "", err
|
||||
}
|
||||
}
|
||||
chargeKey := "generation:hold:" + baseID
|
||||
if existing, found, err := debitAmount(tx, chargeKey); err != nil {
|
||||
return "", err
|
||||
} else if found {
|
||||
return existing, nil
|
||||
}
|
||||
if actual == "0.00" {
|
||||
return actual, nil
|
||||
}
|
||||
_, err = DebitPoints(tx, userID, actual, "generation_hold", baseID, remark, chargeKey)
|
||||
return actual, err
|
||||
}
|
||||
|
||||
func RefundTextCalls(tx *gorm.DB, userID uuid.UUID, referenceID, remark string) (bool, error) {
|
||||
type debit struct {
|
||||
BusinessType string
|
||||
BusinessID string
|
||||
IdempotencyKey string
|
||||
}
|
||||
debits := make([]debit, 0)
|
||||
if err := tx.Table("point_ledger").
|
||||
Select("business_type,business_id,idempotency_key").
|
||||
Where("user_id=? AND change_amount<0 AND business_id LIKE ? AND business_type IN ?", userID, referenceID+":text:%", []string{"generation_hold", "text_token_reserve"}).
|
||||
Order("id").Find(&debits).Error; err != nil {
|
||||
return false, err
|
||||
}
|
||||
refunded := false
|
||||
for _, debit := range debits {
|
||||
refundKey := "generation:refund:" + debit.BusinessID
|
||||
refundType := "generation_refund"
|
||||
if debit.BusinessType == "text_token_reserve" {
|
||||
refundKey = "text:reserve-refund:" + debit.BusinessID
|
||||
refundType = "text_token_reserve_refund"
|
||||
}
|
||||
created, err := RefundDebit(tx, userID, debit.IdempotencyKey, refundKey, refundType, debit.BusinessID, remark)
|
||||
if err != nil {
|
||||
return refunded, err
|
||||
}
|
||||
refunded = refunded || created
|
||||
}
|
||||
return refunded, nil
|
||||
}
|
||||
|
||||
func CommitTextReserves(tx *gorm.DB, userID uuid.UUID, referenceID, remark string) (string, error) {
|
||||
type reserve struct {
|
||||
BusinessID string
|
||||
IdempotencyKey string
|
||||
}
|
||||
items := make([]reserve, 0)
|
||||
if err := tx.Table("point_ledger").Select("business_id,idempotency_key").
|
||||
Where("user_id=? AND business_type='text_token_reserve' AND business_id LIKE ? AND NOT EXISTS (SELECT 1 FROM point_ledger refund WHERE refund.idempotency_key='text:reserve-refund:'||point_ledger.business_id)", userID, referenceID+":text:%").
|
||||
Find(&items).Error; err != nil {
|
||||
return "", err
|
||||
}
|
||||
for _, item := range items {
|
||||
baseID := strings.TrimSuffix(item.BusinessID, ":reserve")
|
||||
if err := tx.Exec("UPDATE point_ledger SET business_type='generation_hold',business_id=?,idempotency_key=?,remark=? WHERE idempotency_key=?", baseID, "generation:hold:"+baseID, remark, item.IdempotencyKey).Error; err != nil {
|
||||
return "", err
|
||||
}
|
||||
}
|
||||
var amount string
|
||||
if err := tx.Table("point_ledger").Select("coalesce(sum(-change_amount),0)::text").
|
||||
Where("user_id=? AND business_type='generation_hold' AND business_id LIKE ?", userID, referenceID+":text:%").Scan(&amount).Error; err != nil {
|
||||
return "", err
|
||||
}
|
||||
return amount, nil
|
||||
}
|
||||
|
||||
func validNonNegativeDecimal(value string) bool {
|
||||
number, ok := new(big.Rat).SetString(strings.TrimSpace(value))
|
||||
return ok && number.Sign() >= 0
|
||||
}
|
||||
|
||||
func ceilPointCents(value *big.Rat) string {
|
||||
if value == nil || value.Sign() <= 0 {
|
||||
return "0.00"
|
||||
}
|
||||
cents := new(big.Rat).Mul(value, big.NewRat(100, 1))
|
||||
quotient := new(big.Int).Quo(cents.Num(), cents.Denom())
|
||||
if new(big.Int).Mod(cents.Num(), cents.Denom()).Sign() > 0 {
|
||||
quotient.Add(quotient, big.NewInt(1))
|
||||
}
|
||||
return fmt.Sprintf("%d.%02d", new(big.Int).Quo(quotient, big.NewInt(100)), new(big.Int).Mod(quotient, big.NewInt(100)))
|
||||
}
|
||||
|
||||
func textCallBusinessID(referenceID, callKey string) string {
|
||||
return referenceID + ":text:" + strings.TrimSpace(callKey)
|
||||
}
|
||||
|
||||
func debitAmount(tx *gorm.DB, idempotencyKey string) (string, bool, error) {
|
||||
var amount string
|
||||
result := tx.Table("point_ledger").Select("(-change_amount)::text").Where("idempotency_key=? AND change_amount<0", idempotencyKey).Limit(1).Scan(&amount)
|
||||
if result.Error != nil {
|
||||
return "", false, result.Error
|
||||
}
|
||||
return amount, strings.TrimSpace(amount) != "", nil
|
||||
}
|
||||
@@ -0,0 +1,70 @@
|
||||
package billing
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestCalculateTextPointsRoundsUpToPointCent(t *testing.T) {
|
||||
pricing := TextPricing{Mode: TextBillingPerToken, InputPerM: "100.0000", OutputPerM: "300.0000"}
|
||||
amount, err := CalculateTextPoints(pricing, TextUsage{Input: 120, Output: 30})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if amount != "0.03" {
|
||||
t.Fatalf("expected 0.03 points, got %s", amount)
|
||||
}
|
||||
amount, err = CalculateTextPoints(pricing, TextUsage{Input: 1})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if amount != "0.01" {
|
||||
t.Fatalf("expected minimum positive charge 0.01, got %s", amount)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTextPricingSnapshot(t *testing.T) {
|
||||
pricing := TextPricing{Mode: TextBillingPerToken, InputPerM: "1250.0000", OutputPerM: "2500.0000"}
|
||||
parsed, err := ParseTextPricingSnapshot(pricing.MarshalSnapshot())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if parsed != pricing {
|
||||
t.Fatalf("unexpected snapshot: %s", json.RawMessage(pricing.MarshalSnapshot()))
|
||||
}
|
||||
}
|
||||
|
||||
func TestTextPricingSnapshotRejectsLegacyThousandTokenRates(t *testing.T) {
|
||||
_, err := ParseTextPricingSnapshot(json.RawMessage(`{"mode":"per_token","input_per_k":"0.1000","output_per_k":"0.3000","max_output_tokens":1000}`))
|
||||
if err == nil {
|
||||
t.Fatal("expected legacy thousand-token snapshot to be rejected")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCalculateTextPointsUsesMillionTokens(t *testing.T) {
|
||||
pricing := TextPricing{Mode: TextBillingPerToken, InputPerM: "10", OutputPerM: "20"}
|
||||
amount, err := CalculateTextPoints(pricing, TextUsage{Input: 1_000_000})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if amount != "10.00" {
|
||||
t.Fatalf("expected 10.00 points for one million input tokens, got %s", amount)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCalculateTextReservePointsUsesDoubleInputAsEstimatedOutput(t *testing.T) {
|
||||
pricing := TextPricing{Mode: TextBillingPerToken, InputPerM: "10", OutputPerM: "20"}
|
||||
amount, err := CalculateTextReservePoints(pricing, 1_000_000)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if amount != "50.00" {
|
||||
t.Fatalf("expected 50.00 points for input plus double-input output estimate, got %s", amount)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEstimateTextTokens(t *testing.T) {
|
||||
if got := EstimateTextTokens("中文ABCD"); got != 3 {
|
||||
t.Fatalf("expected 3 estimated tokens, got %d", got)
|
||||
}
|
||||
}
|
||||
Vendored
+25
@@ -0,0 +1,25 @@
|
||||
package cache
|
||||
|
||||
import (
|
||||
"context"
|
||||
"time"
|
||||
|
||||
"juhe-factory/api/internal/config"
|
||||
|
||||
"github.com/redis/go-redis/v9"
|
||||
)
|
||||
|
||||
func Open(cfg config.Config) (*redis.Client, error) {
|
||||
client := redis.NewClient(&redis.Options{
|
||||
Addr: cfg.RedisAddress,
|
||||
Password: cfg.RedisPassword,
|
||||
DB: cfg.RedisDB,
|
||||
})
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
if err := client.Ping(ctx).Err(); err != nil {
|
||||
client.Close()
|
||||
return nil, err
|
||||
}
|
||||
return client, nil
|
||||
}
|
||||
@@ -0,0 +1,241 @@
|
||||
// Package config 读取并校验 API 服务运行所需的环境变量。
|
||||
package config
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net/url"
|
||||
"os"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"github.com/joho/godotenv"
|
||||
)
|
||||
|
||||
// Config 汇总数据库、缓存、认证、腾讯云 COS 和本地媒体处理配置。
|
||||
type Config struct {
|
||||
AppEnv string
|
||||
Debug bool
|
||||
ServerAddress string
|
||||
DatabaseURL string
|
||||
RedisAddress string
|
||||
RedisPassword string
|
||||
RedisDB int
|
||||
AsynqRedisDB int
|
||||
CORSOrigins []string
|
||||
StaticDirectory string
|
||||
JWTSecretKey string
|
||||
JWTAccessMinutes int
|
||||
JWTRefreshHours int
|
||||
BootstrapAdminUsername string
|
||||
BootstrapAdminPassword string
|
||||
Argon2Time uint32
|
||||
Argon2Memory uint32
|
||||
Argon2Parallelism uint8
|
||||
Argon2HashLength uint32
|
||||
Argon2SaltLength uint32
|
||||
EncryptionKeyVersion string
|
||||
EncryptionKey string
|
||||
COSSecretID string
|
||||
COSSecretKey string
|
||||
COSBucket string
|
||||
COSRegion string
|
||||
COSEndpoint string
|
||||
COSPublicBaseURL string
|
||||
COSMaxImageSizeMB int
|
||||
COSMaxVideoSizeMB int
|
||||
COSMaxAudioSizeMB int
|
||||
AIWorkerConcurrency int
|
||||
AIPollIntervalSeconds int
|
||||
AIHTTPMaxConnections int
|
||||
FFmpegPath string
|
||||
FFprobePath string
|
||||
PythonPath string
|
||||
ASRScriptPath string
|
||||
}
|
||||
|
||||
// Load 从本地环境文件和进程环境读取配置,并校验数值范围。
|
||||
func Load() (Config, error) {
|
||||
_ = godotenv.Load()
|
||||
|
||||
databasePort, err := envInt("DATABASE_PORT", 25432)
|
||||
if err != nil {
|
||||
return Config{}, err
|
||||
}
|
||||
redisPort, err := envInt("REDIS_PORT", 26379)
|
||||
if err != nil {
|
||||
return Config{}, err
|
||||
}
|
||||
redisDB, err := envInt("REDIS_DB", 0)
|
||||
if err != nil {
|
||||
return Config{}, err
|
||||
}
|
||||
asynqRedisDB, err := envInt("ASYNQ_REDIS_DB", 1)
|
||||
if err != nil {
|
||||
return Config{}, err
|
||||
}
|
||||
jwtAccessMinutes, err := envInt("JWT_ACCESS_TOKEN_EXPIRE_MINUTES", 60)
|
||||
if err != nil {
|
||||
return Config{}, err
|
||||
}
|
||||
jwtRefreshHours, err := envInt("JWT_REFRESH_TOKEN_EXPIRE_HOURS", 168)
|
||||
if err != nil {
|
||||
return Config{}, err
|
||||
}
|
||||
argon2Time, err := envInt("ARGON2_TIME_COST", 2)
|
||||
if err != nil {
|
||||
return Config{}, err
|
||||
}
|
||||
argon2Memory, err := envInt("ARGON2_MEMORY_COST", 19456)
|
||||
if err != nil {
|
||||
return Config{}, err
|
||||
}
|
||||
argon2Parallelism, err := envInt("ARGON2_PARALLELISM", 1)
|
||||
if err != nil {
|
||||
return Config{}, err
|
||||
}
|
||||
argon2HashLength, err := envInt("ARGON2_HASH_LENGTH", 32)
|
||||
if err != nil {
|
||||
return Config{}, err
|
||||
}
|
||||
argon2SaltLength, err := envInt("ARGON2_SALT_LENGTH", 16)
|
||||
if err != nil {
|
||||
return Config{}, err
|
||||
}
|
||||
cosMaxImageSizeMB, err := envInt("COS_MAX_IMAGE_SIZE_MB", 20)
|
||||
if err != nil {
|
||||
return Config{}, err
|
||||
}
|
||||
cosMaxVideoSizeMB, err := envInt("COS_MAX_VIDEO_SIZE_MB", 100)
|
||||
if err != nil {
|
||||
return Config{}, err
|
||||
}
|
||||
cosMaxAudioSizeMB, err := envInt("COS_MAX_AUDIO_SIZE_MB", 30)
|
||||
if err != nil {
|
||||
return Config{}, err
|
||||
}
|
||||
aiWorkerConcurrency, err := envInt("AI_WORKER_CONCURRENCY", 100)
|
||||
if err != nil {
|
||||
return Config{}, err
|
||||
}
|
||||
aiPollIntervalSeconds, err := envInt("AI_POLL_INTERVAL_SECONDS", 10)
|
||||
if err != nil {
|
||||
return Config{}, err
|
||||
}
|
||||
aiHTTPMaxConnections, err := envInt("AI_HTTP_MAX_CONNECTIONS", 200)
|
||||
if err != nil {
|
||||
return Config{}, err
|
||||
}
|
||||
if cosMaxImageSizeMB < 1 || cosMaxVideoSizeMB < 1 || cosMaxAudioSizeMB < 1 {
|
||||
return Config{}, fmt.Errorf("COS 文件大小限制必须为正整数")
|
||||
}
|
||||
if aiWorkerConcurrency < 1 || aiWorkerConcurrency > 1000 {
|
||||
return Config{}, fmt.Errorf("AI_WORKER_CONCURRENCY 必须在 1 至 1000 之间")
|
||||
}
|
||||
if aiPollIntervalSeconds < 2 || aiPollIntervalSeconds > 300 {
|
||||
return Config{}, fmt.Errorf("AI_POLL_INTERVAL_SECONDS 必须在 2 至 300 之间")
|
||||
}
|
||||
if aiHTTPMaxConnections < 10 || aiHTTPMaxConnections > 5000 {
|
||||
return Config{}, fmt.Errorf("AI_HTTP_MAX_CONNECTIONS 必须在 10 至 5000 之间")
|
||||
}
|
||||
debug, err := strconv.ParseBool(env("DEBUG", "true"))
|
||||
if err != nil {
|
||||
return Config{}, fmt.Errorf("DEBUG 必须是布尔值: %w", err)
|
||||
}
|
||||
|
||||
databaseURL := strings.TrimSpace(os.Getenv("DATABASE_URL"))
|
||||
if databaseURL == "" {
|
||||
databaseURL = buildDatabaseURL(databasePort)
|
||||
} else if parsed, parseErr := url.Parse(databaseURL); parseErr == nil {
|
||||
query := parsed.Query()
|
||||
if query.Get("TimeZone") == "" {
|
||||
query.Set("TimeZone", "Asia/Shanghai")
|
||||
}
|
||||
parsed.RawQuery = query.Encode()
|
||||
databaseURL = parsed.String()
|
||||
}
|
||||
|
||||
return Config{
|
||||
AppEnv: env("APP_ENV", "development"),
|
||||
Debug: debug,
|
||||
ServerAddress: env("SERVER_ADDRESS", ":8900"),
|
||||
DatabaseURL: databaseURL,
|
||||
RedisAddress: fmt.Sprintf("%s:%d", env("REDIS_HOST", "127.0.0.1"), redisPort),
|
||||
RedisPassword: os.Getenv("REDIS_PASSWORD"),
|
||||
RedisDB: redisDB,
|
||||
AsynqRedisDB: asynqRedisDB,
|
||||
CORSOrigins: splitCSV(env("CORS_ORIGINS", "http://localhost:5500,http://localhost:5501")),
|
||||
StaticDirectory: env("STATIC_DIRECTORY", "static"),
|
||||
JWTSecretKey: env("JWT_SECRET_KEY", "replace-with-a-strong-random-secret"),
|
||||
JWTAccessMinutes: jwtAccessMinutes,
|
||||
JWTRefreshHours: jwtRefreshHours,
|
||||
BootstrapAdminUsername: env("ADMIN_BOOTSTRAP_USERNAME", "admin"),
|
||||
BootstrapAdminPassword: os.Getenv("ADMIN_BOOTSTRAP_PASSWORD"),
|
||||
Argon2Time: uint32(argon2Time),
|
||||
Argon2Memory: uint32(argon2Memory),
|
||||
Argon2Parallelism: uint8(argon2Parallelism),
|
||||
Argon2HashLength: uint32(argon2HashLength),
|
||||
Argon2SaltLength: uint32(argon2SaltLength),
|
||||
EncryptionKeyVersion: env("CONFIG_ENCRYPTION_KEY_VERSION", "v1"),
|
||||
EncryptionKey: os.Getenv("CONFIG_ENCRYPTION_KEY"),
|
||||
COSSecretID: os.Getenv("COS_SECRET_ID"),
|
||||
COSSecretKey: os.Getenv("COS_SECRET_KEY"),
|
||||
COSBucket: os.Getenv("COS_BUCKET"),
|
||||
COSRegion: env("COS_REGION", "ap-chengdu"),
|
||||
COSEndpoint: os.Getenv("COS_ENDPOINT"),
|
||||
COSPublicBaseURL: strings.TrimRight(os.Getenv("COS_PUBLIC_BASE_URL"), "/"),
|
||||
COSMaxImageSizeMB: cosMaxImageSizeMB,
|
||||
COSMaxVideoSizeMB: cosMaxVideoSizeMB,
|
||||
COSMaxAudioSizeMB: cosMaxAudioSizeMB,
|
||||
AIWorkerConcurrency: aiWorkerConcurrency,
|
||||
AIPollIntervalSeconds: aiPollIntervalSeconds,
|
||||
AIHTTPMaxConnections: aiHTTPMaxConnections,
|
||||
FFmpegPath: env("FFMPEG_PATH", "ffmpeg"),
|
||||
FFprobePath: env("FFPROBE_PATH", "ffprobe"),
|
||||
PythonPath: env("PYTHON_PATH", "python3"),
|
||||
ASRScriptPath: env("ASR_SCRIPT_PATH", "workers/transcribe.py"),
|
||||
}, nil
|
||||
}
|
||||
|
||||
func buildDatabaseURL(port int) string {
|
||||
credentials := url.UserPassword(
|
||||
env("DATABASE_USER", "postgres"),
|
||||
env("DATABASE_PASSWORD", "juchuang_dev"),
|
||||
).String()
|
||||
return fmt.Sprintf(
|
||||
"postgres://%s@%s:%d/%s?sslmode=disable&TimeZone=Asia/Shanghai",
|
||||
credentials,
|
||||
env("DATABASE_HOST", "127.0.0.1"),
|
||||
port,
|
||||
url.PathEscape(env("DATABASE_NAME", "juchuang_factory")),
|
||||
)
|
||||
}
|
||||
|
||||
func env(key, fallback string) string {
|
||||
if value := strings.TrimSpace(os.Getenv(key)); value != "" {
|
||||
return value
|
||||
}
|
||||
return fallback
|
||||
}
|
||||
|
||||
func envInt(key string, fallback int) (int, error) {
|
||||
value := strings.TrimSpace(os.Getenv(key))
|
||||
if value == "" {
|
||||
return fallback, nil
|
||||
}
|
||||
parsed, err := strconv.Atoi(value)
|
||||
if err != nil {
|
||||
return 0, fmt.Errorf("%s 必须是整数: %w", key, err)
|
||||
}
|
||||
return parsed, nil
|
||||
}
|
||||
|
||||
func splitCSV(value string) []string {
|
||||
items := strings.Split(value, ",")
|
||||
result := make([]string, 0, len(items))
|
||||
for _, item := range items {
|
||||
if item = strings.TrimSpace(item); item != "" {
|
||||
result = append(result, item)
|
||||
}
|
||||
}
|
||||
return result
|
||||
}
|
||||
@@ -0,0 +1,76 @@
|
||||
package database
|
||||
|
||||
import (
|
||||
_ "embed"
|
||||
"fmt"
|
||||
|
||||
"juhe-factory/api/internal/model"
|
||||
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/logger"
|
||||
)
|
||||
|
||||
//go:embed schema.sql
|
||||
var baselineSchema string
|
||||
|
||||
func bootstrapEmptyDatabase(db *gorm.DB) error {
|
||||
var businessTableCount int64
|
||||
if err := db.Raw(`SELECT count(*) FROM pg_tables
|
||||
WHERE schemaname = 'public' AND tablename NOT IN ('jcf_schema_migrations', 'schema_migrations')`).
|
||||
Scan(&businessTableCount).Error; err != nil {
|
||||
return fmt.Errorf("inspect database schema: %w", err)
|
||||
}
|
||||
if businessTableCount == 0 {
|
||||
baselineDB := db.Session(&gorm.Session{Logger: logger.Default.LogMode(logger.Silent)})
|
||||
if err := baselineDB.Transaction(func(tx *gorm.DB) error {
|
||||
return tx.Exec(baselineSchema).Error
|
||||
}); err != nil {
|
||||
return fmt.Errorf("initialize database baseline: %w", err)
|
||||
}
|
||||
if err := db.Exec(`SET search_path TO public`).Error; err != nil {
|
||||
return fmt.Errorf("restore database search path: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
if err := autoMigrateMissingTables(db, autoMigrateModels()...); err != nil {
|
||||
return fmt.Errorf("auto migrate database models: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func autoMigrateMissingTables(db *gorm.DB, tableModels ...any) error {
|
||||
for _, tableModel := range tableModels {
|
||||
if db.Migrator().HasTable(tableModel) {
|
||||
continue
|
||||
}
|
||||
if err := db.AutoMigrate(tableModel); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func autoMigrateModels() []any {
|
||||
return []any{
|
||||
&model.AdminUser{},
|
||||
&model.AdminRefreshToken{},
|
||||
&model.AdminAuditLog{},
|
||||
&model.WebUser{},
|
||||
&model.WebRefreshToken{},
|
||||
&model.MediaAsset{},
|
||||
&model.CreativeProject{},
|
||||
&model.ProjectModelConfig{},
|
||||
&model.UserModelConfig{},
|
||||
&model.ScriptAnalysis{},
|
||||
&model.ScriptAnalysisCharacter{},
|
||||
&model.ProjectEpisode{},
|
||||
&model.ProjectAsset{},
|
||||
&model.EpisodeStoryboard{},
|
||||
&model.EpisodeSource{},
|
||||
&model.DramaImportSession{},
|
||||
&model.DramaParseBatch{},
|
||||
&model.DramaParseTask{},
|
||||
&model.GenerationTask{},
|
||||
&model.GenerationOutput{},
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,59 @@
|
||||
package database
|
||||
|
||||
import (
|
||||
"os"
|
||||
"testing"
|
||||
|
||||
"juhe-factory/api/internal/config"
|
||||
)
|
||||
|
||||
type autoMigrateProbe struct {
|
||||
ID uint `gorm:"primaryKey"`
|
||||
}
|
||||
|
||||
func (autoMigrateProbe) TableName() string { return "auto_migrate_probe" }
|
||||
|
||||
func TestBootstrapEmptyDatabase(t *testing.T) {
|
||||
databaseURL := os.Getenv("JCF_BOOTSTRAP_TEST_DATABASE_URL")
|
||||
if databaseURL == "" {
|
||||
t.Skip("JCF_BOOTSTRAP_TEST_DATABASE_URL is not configured")
|
||||
}
|
||||
|
||||
db, sqlDB, err := Open(config.Config{DatabaseURL: databaseURL})
|
||||
if err != nil {
|
||||
t.Fatalf("open empty database: %v", err)
|
||||
}
|
||||
defer sqlDB.Close()
|
||||
|
||||
var tableCount int64
|
||||
if err := db.Raw(`SELECT count(*) FROM pg_tables WHERE schemaname = 'public'`).Scan(&tableCount).Error; err != nil {
|
||||
t.Fatalf("count initialized tables: %v", err)
|
||||
}
|
||||
if tableCount != 33 {
|
||||
t.Fatalf("expected 33 initialized tables, got %d", tableCount)
|
||||
}
|
||||
|
||||
var migrationCount int64
|
||||
if err := db.Table("jcf_schema_migrations").Count(&migrationCount).Error; err != nil {
|
||||
t.Fatalf("count baseline migrations: %v", err)
|
||||
}
|
||||
if migrationCount != 0 {
|
||||
t.Fatalf("expected no baseline migrations, got %d", migrationCount)
|
||||
}
|
||||
|
||||
if err := autoMigrateMissingTables(db, &autoMigrateProbe{}); err != nil {
|
||||
t.Fatalf("auto migrate missing table: %v", err)
|
||||
}
|
||||
if !db.Migrator().HasTable(&autoMigrateProbe{}) {
|
||||
t.Fatal("expected AutoMigrate to create missing model table")
|
||||
}
|
||||
if err := db.Migrator().DropTable(&autoMigrateProbe{}); err != nil {
|
||||
t.Fatalf("drop AutoMigrate probe table: %v", err)
|
||||
}
|
||||
|
||||
_, reopenedSQLDB, err := Open(config.Config{DatabaseURL: databaseURL})
|
||||
if err != nil {
|
||||
t.Fatalf("reopen initialized database: %v", err)
|
||||
}
|
||||
defer reopenedSQLDB.Close()
|
||||
}
|
||||
@@ -0,0 +1,133 @@
|
||||
package database
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"juhe-factory/api/internal/config"
|
||||
|
||||
"gorm.io/driver/postgres"
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/logger"
|
||||
)
|
||||
|
||||
func Open(cfg config.Config) (*gorm.DB, *sql.DB, error) {
|
||||
logLevel := logger.Warn
|
||||
if cfg.Debug {
|
||||
logLevel = logger.Info
|
||||
}
|
||||
|
||||
databaseURL := cfg.DatabaseURL
|
||||
if !strings.Contains(strings.ToLower(databaseURL), "timezone=") {
|
||||
separator := "?"
|
||||
if strings.Contains(databaseURL, " ") {
|
||||
separator = " "
|
||||
}
|
||||
if strings.Contains(databaseURL, "?") {
|
||||
separator = "&"
|
||||
}
|
||||
databaseURL += separator + "TimeZone=Asia%2FShanghai"
|
||||
}
|
||||
db, err := gorm.Open(postgres.Open(databaseURL), &gorm.Config{
|
||||
Logger: logger.Default.LogMode(logLevel),
|
||||
})
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
|
||||
sqlDB, err := db.DB()
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
sqlDB.SetMaxOpenConns(20)
|
||||
sqlDB.SetMaxIdleConns(5)
|
||||
sqlDB.SetConnMaxLifetime(30 * time.Minute)
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
if err := sqlDB.PingContext(ctx); err != nil {
|
||||
sqlDB.Close()
|
||||
return nil, nil, err
|
||||
}
|
||||
if err := db.Exec(`SET TIME ZONE 'Asia/Shanghai'`).Error; err != nil {
|
||||
sqlDB.Close()
|
||||
return nil, nil, err
|
||||
}
|
||||
if err := bootstrapEmptyDatabase(db); err != nil {
|
||||
sqlDB.Close()
|
||||
return nil, nil, err
|
||||
}
|
||||
// Keep older installations compatible when migrations are applied manually.
|
||||
if err := db.Exec(`ALTER TABLE project_model_configs ADD COLUMN IF NOT EXISTS settings jsonb NOT NULL DEFAULT '{}'::jsonb`).Error; err != nil {
|
||||
sqlDB.Close()
|
||||
return nil, nil, err
|
||||
}
|
||||
if err := db.Exec(`ALTER TABLE creative_projects ADD COLUMN IF NOT EXISTS video_resolution varchar(8) NOT NULL DEFAULT '720p'`).Error; err != nil {
|
||||
sqlDB.Close()
|
||||
return nil, nil, err
|
||||
}
|
||||
if err := db.Exec(`ALTER TABLE episode_storyboards ADD COLUMN IF NOT EXISTS image_prompt text`).Error; err != nil {
|
||||
sqlDB.Close()
|
||||
return nil, nil, err
|
||||
}
|
||||
if err := db.Exec(`ALTER TABLE prompts ADD COLUMN IF NOT EXISTS type varchar(50); ALTER TABLE prompts ADD COLUMN IF NOT EXISTS content text; ALTER TABLE prompts ADD COLUMN IF NOT EXISTS scope varchar(20) NOT NULL DEFAULT 'system'; ALTER TABLE prompts ADD COLUMN IF NOT EXISTS owner_user_id uuid REFERENCES web_users(id);`).Error; err != nil {
|
||||
sqlDB.Close()
|
||||
return nil, nil, err
|
||||
}
|
||||
if err := db.Exec(`ALTER TABLE web_users ADD COLUMN IF NOT EXISTS account citext; UPDATE web_users SET account=username WHERE account IS NULL; CREATE UNIQUE INDEX IF NOT EXISTS idx_web_users_account ON web_users(account);`).Error; err != nil {
|
||||
sqlDB.Close()
|
||||
return nil, nil, err
|
||||
}
|
||||
if err := db.Exec(`ALTER TABLE redemption_codes ADD COLUMN IF NOT EXISTS code character varying(40);`).Error; err != nil {
|
||||
sqlDB.Close()
|
||||
return nil, nil, err
|
||||
}
|
||||
if err := db.Exec(`CREATE TABLE IF NOT EXISTS user_model_configs (
|
||||
id uuid PRIMARY KEY DEFAULT gen_random_uuid(),
|
||||
user_id uuid NOT NULL REFERENCES web_users(id),
|
||||
purpose varchar(32) NOT NULL CHECK (purpose IN ('prompt_reverse','image_generation','video_generation')),
|
||||
model_id uuid NOT NULL REFERENCES models(id),
|
||||
prompt_id uuid REFERENCES prompts(id),
|
||||
settings jsonb NOT NULL DEFAULT '{}'::jsonb,
|
||||
created_at timestamptz NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
updated_at timestamptz NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
UNIQUE(user_id, purpose)
|
||||
); CREATE INDEX IF NOT EXISTS idx_user_model_configs_user ON user_model_configs(user_id);
|
||||
DO $$
|
||||
BEGIN
|
||||
IF EXISTS (SELECT 1 FROM information_schema.columns WHERE table_name='user_model_configs' AND column_name='model_type') THEN
|
||||
INSERT INTO user_model_configs(id,user_id,project_type,model_type,purpose,model_id,prompt_id,settings)
|
||||
SELECT DISTINCT ON (project.user_id, config.model_type)
|
||||
gen_random_uuid(), project.user_id, 'video_redraw', config.model_type, config.purpose, config.model_id, config.prompt_id, config.settings
|
||||
FROM project_model_configs config
|
||||
JOIN creative_projects project ON project.id=config.project_id
|
||||
WHERE project.deleted_at IS NULL AND project.project_type='video_redraw'
|
||||
ORDER BY project.user_id, config.model_type, project.updated_at DESC, config.updated_at DESC
|
||||
ON CONFLICT DO NOTHING;
|
||||
ELSE
|
||||
INSERT INTO user_model_configs(id,user_id,purpose,model_id,prompt_id,settings)
|
||||
SELECT DISTINCT ON (project.user_id, config.purpose)
|
||||
gen_random_uuid(), project.user_id, config.purpose, config.model_id, config.prompt_id, config.settings
|
||||
FROM project_model_configs config
|
||||
JOIN creative_projects project ON project.id=config.project_id
|
||||
WHERE project.deleted_at IS NULL
|
||||
ORDER BY project.user_id, config.purpose, project.updated_at DESC, config.updated_at DESC
|
||||
ON CONFLICT DO NOTHING;
|
||||
END IF;
|
||||
END $$;`).Error; err != nil {
|
||||
sqlDB.Close()
|
||||
return nil, nil, err
|
||||
}
|
||||
if err := db.Exec(`CREATE TABLE IF NOT EXISTS user_prompt_preferences (
|
||||
id uuid PRIMARY KEY DEFAULT gen_random_uuid(), user_id uuid NOT NULL REFERENCES web_users(id) ON DELETE CASCADE,
|
||||
prompt_id uuid NOT NULL REFERENCES prompts(id) ON DELETE CASCADE, prompt_type varchar(64) NOT NULL,
|
||||
created_at timestamptz NOT NULL DEFAULT CURRENT_TIMESTAMP, updated_at timestamptz NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
UNIQUE(user_id,prompt_type)
|
||||
); CREATE INDEX IF NOT EXISTS idx_user_prompt_preferences_user ON user_prompt_preferences(user_id);`).Error; err != nil {
|
||||
sqlDB.Close()
|
||||
return nil, nil, err
|
||||
}
|
||||
return db, sqlDB, nil
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,30 @@
|
||||
// Package datetime centralizes the application's Beijing time policy.
|
||||
package datetime
|
||||
|
||||
import "time"
|
||||
|
||||
const Zone = "Asia/Shanghai"
|
||||
|
||||
var Location = func() *time.Location {
|
||||
location, err := time.LoadLocation(Zone)
|
||||
if err != nil {
|
||||
return time.FixedZone("CST", 8*60*60)
|
||||
}
|
||||
return location
|
||||
}()
|
||||
|
||||
func Now() time.Time { return time.Now().In(Location) }
|
||||
|
||||
// Normalize converts any source instant (including COS metadata) to Beijing time.
|
||||
func Normalize(value time.Time) time.Time {
|
||||
if value.IsZero() { return value }
|
||||
return value.In(Location)
|
||||
}
|
||||
|
||||
func Parse(value string) (time.Time, error) {
|
||||
parsed, err := time.Parse(time.RFC3339Nano, value)
|
||||
if err != nil { return time.Time{}, err }
|
||||
return Normalize(parsed), nil
|
||||
}
|
||||
|
||||
func Format(value time.Time) string { return Normalize(value).Format(time.RFC3339Nano) }
|
||||
@@ -0,0 +1,20 @@
|
||||
package drama
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
"unicode/utf8"
|
||||
)
|
||||
|
||||
const MaxEpisodeCharacters = 5000
|
||||
|
||||
func ValidateEpisodeContent(content string) error {
|
||||
count := utf8.RuneCountInString(strings.TrimSpace(content))
|
||||
if count == 0 {
|
||||
return fmt.Errorf("剧集原文不能为空")
|
||||
}
|
||||
if count > MaxEpisodeCharacters {
|
||||
return fmt.Errorf("每集原文不能超过%d字,当前%d字", MaxEpisodeCharacters, count)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,18 @@
|
||||
package drama
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestValidateEpisodeContent(t *testing.T) {
|
||||
if err := ValidateEpisodeContent(strings.Repeat("字", MaxEpisodeCharacters)); err != nil {
|
||||
t.Fatalf("expected content at limit to pass: %v", err)
|
||||
}
|
||||
if err := ValidateEpisodeContent(strings.Repeat("字", MaxEpisodeCharacters+1)); err == nil {
|
||||
t.Fatal("expected over-limit content to fail")
|
||||
}
|
||||
if err := ValidateEpisodeContent(" \n\t"); err == nil {
|
||||
t.Fatal("expected blank content to fail")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,111 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"juhe-factory/api/internal/model"
|
||||
"juhe-factory/api/internal/service"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
const adminContextKey = "admin_user"
|
||||
|
||||
type AdminAuth struct{ service *service.Auth }
|
||||
|
||||
func NewAdminAuth(service *service.Auth) *AdminAuth { return &AdminAuth{service: service} }
|
||||
|
||||
func (h *AdminAuth) Login(c *gin.Context) {
|
||||
var body struct {
|
||||
Username string `json:"username" binding:"required"`
|
||||
Password string `json:"password" binding:"required"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&body); err != nil {
|
||||
fail(c, http.StatusBadRequest, "invalid_request", "请输入管理员账号和密码")
|
||||
return
|
||||
}
|
||||
pair, err := h.service.Login(body.Username, body.Password)
|
||||
if err != nil {
|
||||
fail(c, http.StatusUnauthorized, "invalid_credentials", err.Error())
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"data": pair})
|
||||
}
|
||||
|
||||
func (h *AdminAuth) Refresh(c *gin.Context) {
|
||||
var body struct {
|
||||
RefreshToken string `json:"refresh_token" binding:"required"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&body); err != nil {
|
||||
fail(c, http.StatusBadRequest, "invalid_request", "缺少刷新令牌")
|
||||
return
|
||||
}
|
||||
pair, err := h.service.Refresh(body.RefreshToken)
|
||||
if err != nil {
|
||||
fail(c, http.StatusUnauthorized, "invalid_refresh_token", err.Error())
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"data": pair})
|
||||
}
|
||||
|
||||
func (h *AdminAuth) Logout(c *gin.Context) {
|
||||
var body struct {
|
||||
RefreshToken string `json:"refresh_token"`
|
||||
}
|
||||
_ = c.ShouldBindJSON(&body)
|
||||
if err := h.service.Logout(body.RefreshToken); err != nil {
|
||||
fail(c, http.StatusInternalServerError, "logout_failed", "退出登录失败")
|
||||
return
|
||||
}
|
||||
c.Status(http.StatusNoContent)
|
||||
}
|
||||
|
||||
func (h *AdminAuth) ChangePassword(c *gin.Context) {
|
||||
var body struct {
|
||||
CurrentPassword string `json:"current_password" binding:"required"`
|
||||
NewPassword string `json:"new_password" binding:"required"`
|
||||
ConfirmPassword string `json:"confirm_password" binding:"required"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&body); err != nil {
|
||||
fail(c, http.StatusBadRequest, "invalid_request", "请完整填写密码")
|
||||
return
|
||||
}
|
||||
admin := currentAdmin(c)
|
||||
if err := h.service.ChangePassword(admin.ID, body.CurrentPassword, body.NewPassword, body.ConfirmPassword); err != nil {
|
||||
if errors.Is(err, service.ErrReauthFailed) || errors.Is(err, service.ErrAdminPasswordConfirmation) || errors.Is(err, service.ErrAdminPasswordLength) {
|
||||
_ = h.service.Audit(admin, "password_change_failed", "admin-auth", admin.ID.String(), err.Error(), c.ClientIP(), c.GetString("trace_id"), nil)
|
||||
fail(c, http.StatusBadRequest, "password_invalid", err.Error())
|
||||
return
|
||||
}
|
||||
fail(c, http.StatusInternalServerError, "password_update_failed", "密码修改失败,请稍后重试")
|
||||
return
|
||||
}
|
||||
_ = h.service.Audit(admin, "password_changed", "admin-auth", admin.ID.String(), "", c.ClientIP(), c.GetString("trace_id"), nil)
|
||||
c.Status(http.StatusNoContent)
|
||||
}
|
||||
|
||||
func (h *AdminAuth) Middleware() gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
raw := strings.TrimSpace(strings.TrimPrefix(c.GetHeader("Authorization"), "Bearer "))
|
||||
admin, err := h.service.Authenticate(raw)
|
||||
if err != nil {
|
||||
fail(c, http.StatusUnauthorized, "unauthorized", err.Error())
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
c.Set(adminContextKey, admin)
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
|
||||
func currentAdmin(c *gin.Context) *model.AdminUser {
|
||||
value, _ := c.Get(adminContextKey)
|
||||
admin, _ := value.(*model.AdminUser)
|
||||
return admin
|
||||
}
|
||||
|
||||
func fail(c *gin.Context, status int, code, message string) {
|
||||
c.JSON(status, gin.H{"code": code, "message": message, "trace_id": c.GetString("trace_id")})
|
||||
}
|
||||
@@ -0,0 +1,544 @@
|
||||
// 管理后台接口处理器,负责解析管理端请求并调用对应业务模块组织响应。
|
||||
package handler
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"image"
|
||||
_ "image/gif"
|
||||
_ "image/jpeg"
|
||||
_ "image/png"
|
||||
"log/slog"
|
||||
"mime/multipart"
|
||||
"net/http"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
mediakey "juhe-factory/api/internal/media"
|
||||
adminmodule "juhe-factory/api/internal/modules/admin"
|
||||
"juhe-factory/api/internal/modules/prompt"
|
||||
"juhe-factory/api/internal/service"
|
||||
"juhe-factory/api/internal/storage"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/google/uuid"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
type AdminData struct {
|
||||
data *service.AdminData
|
||||
admin *adminmodule.Service
|
||||
prompts *prompt.Service
|
||||
cos *storage.COS
|
||||
}
|
||||
|
||||
type optionalTime struct{ Time *time.Time }
|
||||
|
||||
func (t *optionalTime) UnmarshalJSON(data []byte) error {
|
||||
if string(data) == "null" || strings.TrimSpace(string(data)) == `""` {
|
||||
t.Time = nil
|
||||
return nil
|
||||
}
|
||||
var value time.Time
|
||||
if err := json.Unmarshal(data, &value); err != nil {
|
||||
return err
|
||||
}
|
||||
t.Time = &value
|
||||
return nil
|
||||
}
|
||||
|
||||
func NewAdminData(data *service.AdminData, admin *adminmodule.Service, prompts *prompt.Service, cos *storage.COS) *AdminData {
|
||||
return &AdminData{data: data, admin: admin, prompts: prompts, cos: cos}
|
||||
}
|
||||
|
||||
func (h *AdminData) ListResource(c *gin.Context) {
|
||||
filters := map[string]string{"enabled": c.Query("enabled"), "payment_type": c.Query("payment_type"), "channel_type": c.Query("channel_type"), "status": c.Query("status"), "pinned": c.Query("pinned")}
|
||||
page, err := h.data.ListResource(c.Request.Context(), c.Param("resource"), c.Query("keyword"), filters, c.Query("page"), c.Query("page_size"))
|
||||
if err != nil {
|
||||
h.bad(c, err)
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"data": page})
|
||||
}
|
||||
|
||||
func (h *AdminData) CreateResource(c *gin.Context) {
|
||||
var body map[string]any
|
||||
if err := c.ShouldBindJSON(&body); err != nil {
|
||||
h.bad(c, err)
|
||||
return
|
||||
}
|
||||
id, err := h.data.SaveResource(c.Param("resource"), "", body)
|
||||
if err != nil {
|
||||
h.bad(c, err)
|
||||
return
|
||||
}
|
||||
h.audit(c, "create", c.Param("resource"), id, "", body)
|
||||
c.JSON(http.StatusCreated, gin.H{"data": gin.H{"id": id}})
|
||||
}
|
||||
func (h *AdminData) UpdateResource(c *gin.Context) {
|
||||
var body map[string]any
|
||||
if err := c.ShouldBindJSON(&body); err != nil {
|
||||
h.bad(c, err)
|
||||
return
|
||||
}
|
||||
id, err := h.data.SaveResource(c.Param("resource"), c.Param("id"), body)
|
||||
if err != nil {
|
||||
h.bad(c, err)
|
||||
return
|
||||
}
|
||||
h.audit(c, "update", c.Param("resource"), id, stringValue(body["reason"]), body)
|
||||
c.JSON(http.StatusOK, gin.H{"data": gin.H{"id": id}})
|
||||
}
|
||||
func (h *AdminData) ToggleResource(c *gin.Context) {
|
||||
var body struct {
|
||||
Enabled bool `json:"enabled"`
|
||||
Reason string `json:"reason"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&body); err != nil {
|
||||
h.bad(c, err)
|
||||
return
|
||||
}
|
||||
if err := h.data.ToggleResource(c.Param("resource"), c.Param("id"), body.Enabled); err != nil {
|
||||
h.bad(c, err)
|
||||
return
|
||||
}
|
||||
h.audit(c, "toggle", c.Param("resource"), c.Param("id"), body.Reason, body)
|
||||
c.Status(http.StatusNoContent)
|
||||
}
|
||||
func (h *AdminData) DeleteResource(c *gin.Context) {
|
||||
var body struct {
|
||||
Reason string `json:"reason"`
|
||||
}
|
||||
_ = c.ShouldBindJSON(&body)
|
||||
if err := h.data.DeleteResource(c.Param("resource"), c.Param("id")); err != nil {
|
||||
h.bad(c, err)
|
||||
return
|
||||
}
|
||||
h.audit(c, "delete", c.Param("resource"), c.Param("id"), body.Reason, nil)
|
||||
c.Status(http.StatusNoContent)
|
||||
}
|
||||
|
||||
func (h *AdminData) ReorderStyles(c *gin.Context) {
|
||||
var body struct {
|
||||
IDs []string `json:"ids"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&body); err != nil {
|
||||
h.bad(c, err)
|
||||
return
|
||||
}
|
||||
if err := h.data.ReorderStyles(body.IDs); err != nil {
|
||||
h.bad(c, err)
|
||||
return
|
||||
}
|
||||
h.audit(c, "reorder", "styles", strings.Join(body.IDs, ","), "", map[string]any{"count": len(body.IDs)})
|
||||
c.Status(http.StatusNoContent)
|
||||
}
|
||||
|
||||
func (h *AdminData) ListUsers(c *gin.Context) {
|
||||
page, err := h.data.ListUsers(c.Query("keyword"), c.Query("enabled"), c.Query("page"), c.Query("page_size"))
|
||||
if err != nil {
|
||||
h.internal(c, err)
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"data": page})
|
||||
}
|
||||
|
||||
// ListAuditLogs 返回管理端操作审计记录,不提供修改或删除能力。
|
||||
func (h *AdminData) ListAuditLogs(c *gin.Context) {
|
||||
page, size := parsePage(c)
|
||||
result, err := h.admin.ListAuditLogs(c.Query("keyword"), page, size)
|
||||
if err != nil {
|
||||
h.internal(c, err)
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"data": result})
|
||||
}
|
||||
|
||||
func (h *AdminData) CreateUser(c *gin.Context) {
|
||||
var body struct {
|
||||
Account string `json:"account"`
|
||||
Password string `json:"password"`
|
||||
DailyLimit any `json:"daily_limit"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&body); err != nil {
|
||||
h.bad(c, err)
|
||||
return
|
||||
}
|
||||
item, err := h.data.CreateUser(body.Account, body.Password, body.DailyLimit)
|
||||
if err != nil {
|
||||
h.bad(c, err)
|
||||
return
|
||||
}
|
||||
h.audit(c, "create", "users", stringValue(item["id"]), "", map[string]any{"account": item["account"], "username": item["username"], "daily_limit": body.DailyLimit})
|
||||
c.JSON(http.StatusCreated, gin.H{"data": item})
|
||||
}
|
||||
|
||||
func (h *AdminData) BatchCreateUsers(c *gin.Context) {
|
||||
var body struct {
|
||||
Prefix string `json:"prefix"`
|
||||
StartSequence string `json:"start_sequence"`
|
||||
EndSequence string `json:"end_sequence"`
|
||||
Password string `json:"password"`
|
||||
DailyLimit any `json:"daily_limit"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&body); err != nil {
|
||||
h.bad(c, err)
|
||||
return
|
||||
}
|
||||
items, err := h.data.BatchCreateUsers(body.Prefix, body.StartSequence, body.EndSequence, body.Password, body.DailyLimit)
|
||||
if err != nil {
|
||||
h.bad(c, err)
|
||||
return
|
||||
}
|
||||
h.audit(c, "batch_create", "users", "", "", map[string]any{"prefix": body.Prefix, "start_sequence": body.StartSequence, "end_sequence": body.EndSequence, "count": len(items)})
|
||||
c.JSON(http.StatusCreated, gin.H{"data": items})
|
||||
}
|
||||
|
||||
func (h *AdminData) BatchUpdateUsers(c *gin.Context) {
|
||||
var body struct {
|
||||
IDs []string `json:"ids"`
|
||||
DailyLimit any `json:"daily_limit"`
|
||||
Enabled *bool `json:"enabled"`
|
||||
Reason string `json:"reason"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&body); err != nil {
|
||||
h.bad(c, err)
|
||||
return
|
||||
}
|
||||
updates := map[string]any{}
|
||||
if body.DailyLimit != nil {
|
||||
updates["daily_limit"] = body.DailyLimit
|
||||
}
|
||||
if body.Enabled != nil {
|
||||
updates["enabled"] = *body.Enabled
|
||||
}
|
||||
if err := h.data.UpdateUsers(body.IDs, updates); err != nil {
|
||||
h.bad(c, err)
|
||||
return
|
||||
}
|
||||
h.audit(c, "batch_update", "users", strings.Join(body.IDs, ","), body.Reason, updates)
|
||||
c.Status(http.StatusNoContent)
|
||||
}
|
||||
|
||||
func (h *AdminData) BatchDeleteUsers(c *gin.Context) {
|
||||
var body struct {
|
||||
IDs []string `json:"ids"`
|
||||
Reason string `json:"reason"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&body); err != nil {
|
||||
h.bad(c, err)
|
||||
return
|
||||
}
|
||||
if err := h.data.DeleteUsers(body.IDs); err != nil {
|
||||
h.bad(c, err)
|
||||
return
|
||||
}
|
||||
h.audit(c, "batch_delete", "users", strings.Join(body.IDs, ","), body.Reason, nil)
|
||||
c.Status(http.StatusNoContent)
|
||||
}
|
||||
|
||||
// GrantUserPoints 校验管理端积分发放请求,完成入账后记录操作审计。
|
||||
func (h *AdminData) GrantUserPoints(c *gin.Context) {
|
||||
userID, err := uuid.Parse(c.Param("id"))
|
||||
if err != nil {
|
||||
h.bad(c, errors.New("用户标识无效"))
|
||||
return
|
||||
}
|
||||
var body struct {
|
||||
RequestID string `json:"request_id"`
|
||||
Points any `json:"points"`
|
||||
Reason string `json:"reason"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&body); err != nil {
|
||||
h.bad(c, err)
|
||||
return
|
||||
}
|
||||
body.Reason = strings.TrimSpace(body.Reason)
|
||||
if body.Reason == "" {
|
||||
h.bad(c, errors.New("请填写积分发放原因"))
|
||||
return
|
||||
}
|
||||
if len([]rune(body.Reason)) > 200 {
|
||||
h.bad(c, errors.New("积分发放原因不能超过 200 个字符"))
|
||||
return
|
||||
}
|
||||
balance, credited, err := h.admin.GrantUserPoints(userID, body.RequestID, body.Points)
|
||||
if err != nil {
|
||||
h.bad(c, err)
|
||||
return
|
||||
}
|
||||
if credited {
|
||||
h.audit(c, "grant_points", "users", userID.String(), body.Reason, map[string]any{
|
||||
"points": body.Points, "balance": balance, "request_id": body.RequestID,
|
||||
})
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"data": gin.H{"balance": balance, "credited": credited}})
|
||||
}
|
||||
|
||||
func (h *AdminData) ListModels(c *gin.Context) {
|
||||
page, size := parsePage(c)
|
||||
result, err := h.admin.ListModels(c.Query("keyword"), c.Query("model_type"), c.Query("channel_id"), page, size)
|
||||
if err != nil {
|
||||
h.internal(c, err)
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"data": result})
|
||||
}
|
||||
|
||||
func (h *AdminData) SaveModel(c *gin.Context) {
|
||||
var body adminmodule.ModelInput
|
||||
if err := c.ShouldBindJSON(&body); err != nil {
|
||||
h.bad(c, err)
|
||||
return
|
||||
}
|
||||
id, created, normalizedBody, err := h.admin.SaveModel(c.Param("id"), body)
|
||||
if err != nil {
|
||||
h.bad(c, err)
|
||||
return
|
||||
}
|
||||
action := "update"
|
||||
status := http.StatusOK
|
||||
if created {
|
||||
action = "create"
|
||||
status = http.StatusCreated
|
||||
}
|
||||
h.audit(c, action, "models", id, "", normalizedBody)
|
||||
c.JSON(status, gin.H{"data": gin.H{"id": id}})
|
||||
}
|
||||
|
||||
func (h *AdminData) ListPrompts(c *gin.Context) {
|
||||
page, size := parsePage(c)
|
||||
items, total, err := h.prompts.ListSystemPrompts(c.Query("keyword"), c.Query("type"), page, size)
|
||||
if err != nil {
|
||||
h.internal(c, err)
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"data": service.Page{Items: items, Total: total, Page: page, PageSize: size}})
|
||||
}
|
||||
|
||||
func (h *AdminData) SavePrompt(c *gin.Context) {
|
||||
var body struct {
|
||||
Code string `json:"code"`
|
||||
Name string `json:"name"`
|
||||
Category string `json:"category"`
|
||||
Type string `json:"type"`
|
||||
Content string `json:"content"`
|
||||
Variables []string `json:"variables"`
|
||||
VersionNote string `json:"version_note"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&body); err != nil {
|
||||
h.bad(c, err)
|
||||
return
|
||||
}
|
||||
if strings.TrimSpace(body.Category) == "" {
|
||||
body.Category = body.Type
|
||||
}
|
||||
if strings.TrimSpace(body.Type) == "" {
|
||||
body.Type = body.Category
|
||||
}
|
||||
if err := service.ValidatePromptVariables(body.Content, body.Variables); err != nil {
|
||||
h.bad(c, err)
|
||||
return
|
||||
}
|
||||
id, err := h.prompts.SaveVersionedPrompt(currentAdmin(c).ID, c.Param("id"), prompt.VersionedPromptInput{
|
||||
Code: body.Code, Name: body.Name, Category: body.Category, Type: body.Type,
|
||||
Content: body.Content, Variables: body.Variables, VersionNote: body.VersionNote,
|
||||
})
|
||||
if err != nil {
|
||||
h.bad(c, err)
|
||||
return
|
||||
}
|
||||
h.audit(c, "save_draft", "prompts", id, body.VersionNote, nil)
|
||||
c.JSON(http.StatusOK, gin.H{"data": gin.H{"id": id}})
|
||||
}
|
||||
|
||||
// SavePromptSimple 保存系统提示词并立即应用,不启用历史版本工作流。
|
||||
func (h *AdminData) SavePromptSimple(c *gin.Context) {
|
||||
var body prompt.SystemPromptInput
|
||||
if err := c.ShouldBindJSON(&body); err != nil {
|
||||
h.bad(c, err)
|
||||
return
|
||||
}
|
||||
item, err := h.prompts.SaveSystemPrompt(c.Param("id"), body)
|
||||
if err != nil {
|
||||
h.bad(c, err)
|
||||
return
|
||||
}
|
||||
h.audit(c, "save", "prompts", item.ID, "", body)
|
||||
c.JSON(http.StatusOK, gin.H{"data": item})
|
||||
}
|
||||
|
||||
// DeletePromptSimple 删除管理端指定的系统提示词。
|
||||
func (h *AdminData) DeletePromptSimple(c *gin.Context) {
|
||||
if err := h.prompts.DeleteSystemPrompt(c.Param("id")); err != nil {
|
||||
h.bad(c, err)
|
||||
return
|
||||
}
|
||||
c.Status(http.StatusNoContent)
|
||||
}
|
||||
|
||||
func (h *AdminData) PromptAction(c *gin.Context) {
|
||||
var body struct {
|
||||
Action string `json:"action"`
|
||||
VersionID string `json:"version_id"`
|
||||
VersionNote string `json:"version_note"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&body); err != nil {
|
||||
h.bad(c, err)
|
||||
return
|
||||
}
|
||||
id := c.Param("id")
|
||||
if err := h.prompts.ApplyPromptAction(currentAdmin(c).ID, id, body.Action, body.VersionID, body.VersionNote); err != nil {
|
||||
h.bad(c, err)
|
||||
return
|
||||
}
|
||||
h.audit(c, body.Action, "prompts", id, body.VersionNote, body)
|
||||
c.Status(http.StatusNoContent)
|
||||
}
|
||||
|
||||
func (h *AdminData) PromptHistory(c *gin.Context) {
|
||||
items, err := h.prompts.ListPromptHistory(c.Param("id"))
|
||||
if err != nil {
|
||||
h.internal(c, err)
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"data": items})
|
||||
}
|
||||
|
||||
func (h *AdminData) PromptActionSimple(c *gin.Context) { c.Status(http.StatusNoContent) }
|
||||
func (h *AdminData) PromptHistorySimple(c *gin.Context) {
|
||||
c.JSON(http.StatusOK, gin.H{"data": []any{}})
|
||||
}
|
||||
|
||||
func (h *AdminData) ListRedemptions(c *gin.Context) {
|
||||
page, size := parsePage(c)
|
||||
result, err := h.admin.ListRedemptions(c.Query("keyword"), c.Query("status"), page, size)
|
||||
if err != nil {
|
||||
h.internal(c, err)
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"data": result})
|
||||
}
|
||||
|
||||
func (h *AdminData) CreateRedemptionBatch(c *gin.Context) {
|
||||
var body struct {
|
||||
Name string `json:"name"`
|
||||
Points any `json:"points"`
|
||||
Quantity int `json:"quantity"`
|
||||
ExpiresAt optionalTime `json:"expires_at"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&body); err != nil {
|
||||
h.bad(c, err)
|
||||
return
|
||||
}
|
||||
if body.ExpiresAt.Time == nil {
|
||||
expiresAt := time.Now().AddDate(0, 0, 7)
|
||||
body.ExpiresAt.Time = &expiresAt
|
||||
}
|
||||
admin := currentAdmin(c)
|
||||
batchID, codes, points, err := h.admin.CreateRedemptionBatch(admin.ID, body.Name, body.Points, body.Quantity, *body.ExpiresAt.Time)
|
||||
if err != nil {
|
||||
h.bad(c, err)
|
||||
return
|
||||
}
|
||||
h.audit(c, "generate", "redemption-codes", batchID, "", map[string]any{"quantity": body.Quantity, "points": points})
|
||||
c.JSON(http.StatusCreated, gin.H{"data": gin.H{"batch_id": batchID, "codes": codes}})
|
||||
}
|
||||
|
||||
// validateRedemptionBatchLimits 保留 Handler 包内既有测试入口,实际规则由管理后台模块统一实现。
|
||||
func validateRedemptionBatchLimits(value any, quantity int) (string, error) {
|
||||
return adminmodule.ValidateRedemptionBatchLimits(value, quantity)
|
||||
}
|
||||
|
||||
func (h *AdminData) UploadStyleImage(c *gin.Context) {
|
||||
if h.cos == nil {
|
||||
fail(c, http.StatusServiceUnavailable, "cos_not_configured", "COS 对象存储未配置")
|
||||
return
|
||||
}
|
||||
file, header, err := c.Request.FormFile("file")
|
||||
if err != nil {
|
||||
h.bad(c, errors.New("请选择图片文件"))
|
||||
return
|
||||
}
|
||||
defer file.Close()
|
||||
if header.Size <= 0 || header.Size > h.cos.MaxImageBytes() {
|
||||
h.bad(c, errors.New("图片大小超出限制"))
|
||||
return
|
||||
}
|
||||
contentType := header.Header.Get("Content-Type")
|
||||
if contentType != "image/jpeg" && contentType != "image/png" && contentType != "image/gif" {
|
||||
h.bad(c, errors.New("仅支持 JPEG、PNG 或 GIF 图片"))
|
||||
return
|
||||
}
|
||||
width, height, err := imageDimensions(file, contentType)
|
||||
if err != nil {
|
||||
h.bad(c, errors.New("无法识别图片内容"))
|
||||
return
|
||||
}
|
||||
if width < 64 || height < 64 || width > 8192 || height > 8192 {
|
||||
h.bad(c, errors.New("图片尺寸必须在 64×64 至 8192×8192 之间"))
|
||||
return
|
||||
}
|
||||
if _, err = file.Seek(0, 0); err != nil {
|
||||
h.internal(c, err)
|
||||
return
|
||||
}
|
||||
ext := strings.ToLower(filepath.Ext(header.Filename))
|
||||
if ext == "" {
|
||||
ext = ".img"
|
||||
}
|
||||
key := mediakey.StyleImage(uuid.New(), ext)
|
||||
url, err := h.cos.Put(c, key, contentType, file, header.Size)
|
||||
if err != nil {
|
||||
h.internal(c, err)
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusCreated, gin.H{"data": gin.H{"image_key": key, "image_url": url, "image_mime": contentType, "image_size": header.Size, "image_width": width, "image_height": height}})
|
||||
}
|
||||
|
||||
func imageDimensions(file multipart.File, contentType string) (int, int, error) {
|
||||
if contentType == "image/webp" {
|
||||
return 0, 0, errors.New("WebP 尺寸解析暂不可用")
|
||||
}
|
||||
cfg, _, err := image.DecodeConfig(file)
|
||||
return cfg.Width, cfg.Height, err
|
||||
}
|
||||
|
||||
func parsePage(c *gin.Context) (int, int) {
|
||||
page, _ := strconv.Atoi(c.Query("page"))
|
||||
if page < 1 {
|
||||
page = 1
|
||||
}
|
||||
size, _ := strconv.Atoi(c.Query("page_size"))
|
||||
if size < 1 {
|
||||
size = 20
|
||||
}
|
||||
if size > 100 {
|
||||
size = 100
|
||||
}
|
||||
return page, size
|
||||
}
|
||||
func stringValue(v any) string {
|
||||
if v == nil {
|
||||
return ""
|
||||
}
|
||||
return fmt.Sprint(v)
|
||||
}
|
||||
func (h *AdminData) audit(c *gin.Context, action, resource, id, reason string, detail any) {
|
||||
_ = h.admin.WriteAudit(currentAdmin(c), action, resource, id, reason, c.ClientIP(), c.GetString("trace_id"), detail)
|
||||
}
|
||||
func (h *AdminData) bad(c *gin.Context, err error) {
|
||||
message := err.Error()
|
||||
code := "invalid_request"
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
code = "not_found"
|
||||
}
|
||||
fail(c, http.StatusBadRequest, code, message)
|
||||
}
|
||||
func (h *AdminData) internal(c *gin.Context, err error) {
|
||||
slog.ErrorContext(c.Request.Context(), "admin request failed", "trace_id", c.GetString("trace_id"), "error", err)
|
||||
fail(c, http.StatusInternalServerError, "internal_error", "操作失败,请稍后重试")
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,47 @@
|
||||
// 本文件负责剧集视频 ZIP 下载请求的参数校验、响应头设置和对象流写入。
|
||||
package handler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"io"
|
||||
"mime"
|
||||
"net/http"
|
||||
|
||||
"juhe-factory/api/internal/service"
|
||||
"juhe-factory/api/internal/videoarchive"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// DownloadEpisodeVideos 下载当前剧集按项目类型筛选后的视频 ZIP。
|
||||
func (h *Creative) DownloadEpisodeVideos(c *gin.Context) {
|
||||
if !h.objectStorageAvailable(c) {
|
||||
return
|
||||
}
|
||||
projectID, ok := h.projectID(c)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
episodeID, err := service.ParseUUID(c.Param("episode_id"), "剧集")
|
||||
if err != nil {
|
||||
h.bad(c, err)
|
||||
return
|
||||
}
|
||||
manifest, err := h.service.EpisodeVideoArchive(currentWebUser(c).ID, projectID, episodeID)
|
||||
if err != nil {
|
||||
h.bad(c, err)
|
||||
return
|
||||
}
|
||||
c.Header("Content-Type", "application/zip")
|
||||
c.Header("Content-Disposition", mime.FormatMediaType("attachment", map[string]string{"filename": manifest.Name}))
|
||||
c.Header("X-Content-Type-Options", "nosniff")
|
||||
c.Status(http.StatusOK)
|
||||
err = videoarchive.Write(c.Request.Context(), c.Writer, manifest.Entries, func(ctx context.Context, objectKey string) (io.ReadCloser, error) {
|
||||
body, _, _, openErr := h.cos.Open(ctx, objectKey)
|
||||
return body, openErr
|
||||
})
|
||||
if err != nil {
|
||||
_ = c.Error(err)
|
||||
c.Abort()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,210 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"juhe-factory/api/internal/service"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
func (h *Creative) PreviewDramaImport(c *gin.Context) {
|
||||
projectID, ok := h.projectID(c)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
c.Request.Body = http.MaxBytesReader(c.Writer, c.Request.Body, 4<<20)
|
||||
if err := c.Request.ParseMultipartForm(4 << 20); err != nil {
|
||||
h.bad(c, errors.New("导入内容无效"))
|
||||
return
|
||||
}
|
||||
if c.Request.MultipartForm != nil {
|
||||
defer c.Request.MultipartForm.RemoveAll()
|
||||
}
|
||||
rawText := strings.TrimSpace(c.PostForm("raw_text"))
|
||||
var filename string
|
||||
var content []byte
|
||||
file, header, err := c.Request.FormFile("file")
|
||||
if err == nil {
|
||||
defer file.Close()
|
||||
filename = header.Filename
|
||||
content, err = io.ReadAll(io.LimitReader(file, 3*1024*1024+1))
|
||||
if err != nil {
|
||||
h.bad(c, errors.New("读取小说文件失败"))
|
||||
return
|
||||
}
|
||||
} else if rawText == "" {
|
||||
h.bad(c, errors.New("请选择小说文件或粘贴文本"))
|
||||
return
|
||||
}
|
||||
preview, err := h.service.PreviewDramaImport(currentWebUser(c).ID, projectID, filename, content, rawText)
|
||||
if err != nil {
|
||||
h.respondError(c, err)
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"data": preview})
|
||||
}
|
||||
|
||||
func (h *Creative) ConfirmDramaImport(c *gin.Context) {
|
||||
projectID, ok := h.projectID(c)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
var body struct {
|
||||
ImportToken uuid.UUID `json:"import_token"`
|
||||
ChaptersPerEpisode int `json:"chapters_per_episode"`
|
||||
StartEpisodeNo int `json:"start_episode_no"`
|
||||
TreatAsSingleEpisode bool `json:"treat_as_single_episode"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&body); err != nil || body.ImportToken == uuid.Nil {
|
||||
h.bad(c, errors.New("导入确认参数无效"))
|
||||
return
|
||||
}
|
||||
items, err := h.service.ConfirmDramaImport(currentWebUser(c).ID, projectID, body.ImportToken, body.ChaptersPerEpisode, body.StartEpisodeNo, body.TreatAsSingleEpisode)
|
||||
if err != nil {
|
||||
h.respondError(c, err)
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusCreated, gin.H{"data": items})
|
||||
}
|
||||
|
||||
func (h *Creative) GetDramaEpisodeSource(c *gin.Context) {
|
||||
projectID, episodeID, ok := h.dramaEpisodeIDs(c)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
source, err := h.service.EpisodeSource(currentWebUser(c).ID, projectID, episodeID)
|
||||
if err != nil {
|
||||
h.respondError(c, err)
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"data": source})
|
||||
}
|
||||
|
||||
func (h *Creative) SaveDramaEpisodeSource(c *gin.Context) {
|
||||
projectID, episodeID, ok := h.dramaEpisodeIDs(c)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
var body struct {
|
||||
Title string `json:"title"`
|
||||
RawContent string `json:"raw_content"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&body); err != nil {
|
||||
h.bad(c, err)
|
||||
return
|
||||
}
|
||||
source, err := h.service.SaveEpisodeSource(currentWebUser(c).ID, projectID, episodeID, body.RawContent, body.Title)
|
||||
if err != nil {
|
||||
h.respondError(c, err)
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"data": source})
|
||||
}
|
||||
|
||||
func (h *Creative) QueueDramaParse(c *gin.Context) {
|
||||
projectID, episodeID, ok := h.dramaEpisodeIDs(c)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
task, err := h.service.QueueDramaParse(currentWebUser(c).ID, projectID, episodeID)
|
||||
if err != nil {
|
||||
h.respondError(c, err)
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusAccepted, gin.H{"data": task})
|
||||
}
|
||||
|
||||
func (h *Creative) ReparseDrama(c *gin.Context) {
|
||||
if !h.objectStorageAvailable(c) {
|
||||
return
|
||||
}
|
||||
projectID, episodeID, ok := h.dramaEpisodeIDs(c)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
userID := currentWebUser(c).ID
|
||||
objectKeys, err := h.service.ResetEpisodeAnalysis(userID, projectID, episodeID)
|
||||
if err != nil {
|
||||
h.respondError(c, err)
|
||||
return
|
||||
}
|
||||
if !h.deleteStoredObjects(c, objectKeys) {
|
||||
return
|
||||
}
|
||||
task, err := h.service.QueueDramaParse(userID, projectID, episodeID)
|
||||
if err != nil {
|
||||
h.respondError(c, err)
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusAccepted, gin.H{"data": task})
|
||||
}
|
||||
|
||||
func (h *Creative) ListDramaParseTasks(c *gin.Context) {
|
||||
projectID, ok := h.projectID(c)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
var episodeID *uuid.UUID
|
||||
if raw := strings.TrimSpace(c.Query("episode_id")); raw != "" {
|
||||
parsed, err := service.ParseUUID(raw, "剧集")
|
||||
if err != nil {
|
||||
h.bad(c, err)
|
||||
return
|
||||
}
|
||||
episodeID = &parsed
|
||||
}
|
||||
items, err := h.service.ListDramaParseTasks(currentWebUser(c).ID, projectID, episodeID)
|
||||
if err != nil {
|
||||
h.respondError(c, err)
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"data": items})
|
||||
}
|
||||
|
||||
func (h *Creative) CancelDramaParse(c *gin.Context) {
|
||||
projectID, ok := h.projectID(c)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
taskID, err := service.ParseUUID(c.Param("task_id"), "解析任务")
|
||||
if err != nil {
|
||||
h.bad(c, err)
|
||||
return
|
||||
}
|
||||
if err = h.service.CancelDramaParse(currentWebUser(c).ID, projectID, taskID); err != nil {
|
||||
h.respondError(c, err)
|
||||
return
|
||||
}
|
||||
c.Status(http.StatusNoContent)
|
||||
}
|
||||
|
||||
func (h *Creative) CreateDramaStoryboard(c *gin.Context) {
|
||||
projectID, episodeID, ok := h.dramaEpisodeIDs(c)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
storyboard, err := h.service.CreateStoryboard(currentWebUser(c).ID, projectID, episodeID)
|
||||
if err != nil {
|
||||
h.respondError(c, err)
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusCreated, gin.H{"data": storyboard})
|
||||
}
|
||||
|
||||
func (h *Creative) dramaEpisodeIDs(c *gin.Context) (uuid.UUID, uuid.UUID, bool) {
|
||||
projectID, ok := h.projectID(c)
|
||||
if !ok {
|
||||
return uuid.Nil, uuid.Nil, false
|
||||
}
|
||||
episodeID, err := service.ParseUUID(c.Param("episode_id"), "剧集")
|
||||
if err != nil {
|
||||
h.bad(c, err)
|
||||
return uuid.Nil, uuid.Nil, false
|
||||
}
|
||||
return projectID, episodeID, true
|
||||
}
|
||||
@@ -0,0 +1,48 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"net/http"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/redis/go-redis/v9"
|
||||
)
|
||||
|
||||
type Health struct {
|
||||
db *sql.DB
|
||||
redis *redis.Client
|
||||
}
|
||||
|
||||
func NewHealth(db *sql.DB, redisClient *redis.Client) *Health {
|
||||
return &Health{db: db, redis: redisClient}
|
||||
}
|
||||
|
||||
func (h *Health) Check(c *gin.Context) {
|
||||
ctx, cancel := context.WithTimeout(c.Request.Context(), 2*time.Second)
|
||||
defer cancel()
|
||||
|
||||
services := map[string]string{
|
||||
"postgresql": "ok",
|
||||
"redis": "ok",
|
||||
}
|
||||
status := "ok"
|
||||
statusCode := http.StatusOK
|
||||
|
||||
if err := h.db.PingContext(ctx); err != nil {
|
||||
services["postgresql"] = "unavailable"
|
||||
status = "degraded"
|
||||
statusCode = http.StatusServiceUnavailable
|
||||
}
|
||||
if err := h.redis.Ping(ctx).Err(); err != nil {
|
||||
services["redis"] = "unavailable"
|
||||
status = "degraded"
|
||||
statusCode = http.StatusServiceUnavailable
|
||||
}
|
||||
|
||||
c.JSON(statusCode, gin.H{
|
||||
"status": status,
|
||||
"services": services,
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,236 @@
|
||||
// 图片生成接口处理器,负责生图参数校验、参考图上传、历史查询和结果删除。
|
||||
package handler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"errors"
|
||||
"image"
|
||||
"io"
|
||||
"mime"
|
||||
"net/http"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
|
||||
mediakey "juhe-factory/api/internal/media"
|
||||
"juhe-factory/api/internal/model"
|
||||
"juhe-factory/api/internal/modules/productimage"
|
||||
"juhe-factory/api/internal/service"
|
||||
"juhe-factory/api/internal/storage"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/google/uuid"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
// ProductImage 组合图片生成生图服务与对象存储依赖。
|
||||
type ProductImage struct {
|
||||
service *productimage.Service
|
||||
cos *storage.COS
|
||||
}
|
||||
|
||||
// NewProductImage 创建图片生成 HTTP 处理器。
|
||||
func NewProductImage(service *productimage.Service, cos *storage.COS) *ProductImage {
|
||||
return &ProductImage{service: service, cos: cos}
|
||||
}
|
||||
|
||||
// Generate 接收当前页面的临时配置和参考图,并创建图片生成任务。
|
||||
func (h *ProductImage) Generate(c *gin.Context) {
|
||||
if h.cos == nil {
|
||||
fail(c, http.StatusServiceUnavailable, "cos_not_configured", "对象存储未配置")
|
||||
return
|
||||
}
|
||||
c.Request.Body = http.MaxBytesReader(c.Writer, c.Request.Body, 48<<20)
|
||||
if err := c.Request.ParseMultipartForm(48 << 20); err != nil {
|
||||
fail(c, http.StatusBadRequest, "invalid_request", "生图请求无效")
|
||||
return
|
||||
}
|
||||
if c.Request.MultipartForm != nil {
|
||||
defer c.Request.MultipartForm.RemoveAll()
|
||||
}
|
||||
modelID, err := service.ParseUUID(c.PostForm("model_id"), "图片模型")
|
||||
if err != nil {
|
||||
fail(c, http.StatusBadRequest, "invalid_request", err.Error())
|
||||
return
|
||||
}
|
||||
ratio := strings.TrimSpace(c.PostForm("aspect_ratio"))
|
||||
if ratio == "" {
|
||||
ratio = "9:16"
|
||||
}
|
||||
resolution := strings.TrimSpace(c.PostForm("resolution"))
|
||||
if resolution == "" {
|
||||
resolution = "1k"
|
||||
}
|
||||
userID := currentWebUser(c).ID
|
||||
files := c.Request.MultipartForm.File["references"]
|
||||
if len(files) > 4 {
|
||||
fail(c, http.StatusBadRequest, "invalid_request", "最多上传 4 张参考图")
|
||||
return
|
||||
}
|
||||
references := make([]productimage.ReferenceUpload, 0, len(files))
|
||||
objectKeys := make([]string, 0, len(files))
|
||||
for index, header := range files {
|
||||
file, openErr := header.Open()
|
||||
if openErr != nil {
|
||||
if !h.cleanupUploadedObjects(c, objectKeys) {
|
||||
return
|
||||
}
|
||||
fail(c, http.StatusBadRequest, "invalid_request", "参考图无法读取")
|
||||
return
|
||||
}
|
||||
contentType := cleanProductImageContentType(header.Header.Get("Content-Type"))
|
||||
if contentType == "" {
|
||||
contentType = cleanProductImageContentType(mime.TypeByExtension(strings.ToLower(filepath.Ext(header.Filename))))
|
||||
}
|
||||
allowed := map[string]bool{"image/jpeg": true, "image/png": true, "image/webp": true}
|
||||
if !allowed[contentType] || header.Size <= 0 || header.Size > h.cos.MaxImageBytes() {
|
||||
file.Close()
|
||||
if !h.cleanupUploadedObjects(c, objectKeys) {
|
||||
return
|
||||
}
|
||||
fail(c, http.StatusBadRequest, "invalid_request", "参考图格式或大小不符合要求")
|
||||
return
|
||||
}
|
||||
config, _, decodeErr := image.DecodeConfig(file)
|
||||
if decodeErr != nil || config.Width <= 0 || config.Height <= 0 {
|
||||
file.Close()
|
||||
if !h.cleanupUploadedObjects(c, objectKeys) {
|
||||
return
|
||||
}
|
||||
fail(c, http.StatusBadRequest, "invalid_request", "参考图内容无法解析")
|
||||
return
|
||||
}
|
||||
if _, err = file.Seek(0, io.SeekStart); err != nil {
|
||||
file.Close()
|
||||
if !h.cleanupUploadedObjects(c, objectKeys) {
|
||||
return
|
||||
}
|
||||
h.respondError(c, err)
|
||||
return
|
||||
}
|
||||
mediaID := uuid.New()
|
||||
key := mediakey.ProductImageReference(userID, mediaID, header.Filename, contentType)
|
||||
hash := sha256.New()
|
||||
if _, err = io.Copy(hash, file); err != nil {
|
||||
file.Close()
|
||||
if !h.cleanupUploadedObjects(c, objectKeys) {
|
||||
return
|
||||
}
|
||||
h.respondError(c, err)
|
||||
return
|
||||
}
|
||||
if _, err = file.Seek(0, io.SeekStart); err != nil {
|
||||
file.Close()
|
||||
if !h.cleanupUploadedObjects(c, objectKeys) {
|
||||
return
|
||||
}
|
||||
h.respondError(c, err)
|
||||
return
|
||||
}
|
||||
url, putErr := h.cos.Put(c, key, contentType, file, header.Size)
|
||||
file.Close()
|
||||
if putErr != nil {
|
||||
if !h.cleanupUploadedObjects(c, objectKeys) {
|
||||
return
|
||||
}
|
||||
h.respondError(c, putErr)
|
||||
return
|
||||
}
|
||||
objectKeys = append(objectKeys, key)
|
||||
name := strings.TrimSpace(filepath.Base(header.Filename))
|
||||
if name == "" {
|
||||
name = "参考图" + string(rune('1'+index))
|
||||
}
|
||||
owner := userID
|
||||
references = append(references, productimage.ReferenceUpload{Asset: &model.MediaAsset{ID: mediaID, OwnerUserID: &owner, StorageProvider: "cos", ObjectKey: key, PublicURL: url, OriginalName: name, DisplayName: "图片生成参考图", MimeType: contentType, SizeBytes: header.Size, SHA256: hex.EncodeToString(hash.Sum(nil)), Width: &config.Width, Height: &config.Height}, Name: name})
|
||||
}
|
||||
task, err := h.service.QueueGeneration(userID, productimage.GenerateInput{Prompt: c.PostForm("prompt"), ModelID: modelID, AspectRatio: ratio, Resolution: resolution, References: references})
|
||||
if err != nil {
|
||||
if !h.cleanupUploadedObjects(c, objectKeys) {
|
||||
return
|
||||
}
|
||||
h.respondError(c, err)
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusAccepted, gin.H{"data": gin.H{"task_id": task.ID}})
|
||||
}
|
||||
|
||||
// ListGenerations 返回当前用户最近的图片生成历史。
|
||||
func (h *ProductImage) ListGenerations(c *gin.Context) {
|
||||
items, err := h.service.Generations(currentWebUser(c).ID)
|
||||
if err != nil {
|
||||
h.respondError(c, err)
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"data": items})
|
||||
}
|
||||
|
||||
// DeleteGeneration 删除当前用户的一条图片生成记录及其 COS 对象。
|
||||
func (h *ProductImage) DeleteGeneration(c *gin.Context) {
|
||||
taskID, err := service.ParseUUID(c.Param("task_id"), "图片生成任务")
|
||||
if err != nil {
|
||||
fail(c, http.StatusBadRequest, "invalid_request", err.Error())
|
||||
return
|
||||
}
|
||||
err = h.service.DeleteGeneration(currentWebUser(c).ID, taskID, func(keys []string) error {
|
||||
return h.deleteObjects(c.Request.Context(), keys)
|
||||
})
|
||||
if err != nil {
|
||||
h.respondError(c, err)
|
||||
return
|
||||
}
|
||||
c.Status(http.StatusNoContent)
|
||||
}
|
||||
|
||||
// deleteObjects 删除一组已经精确解析的 COS 对象键。
|
||||
func (h *ProductImage) deleteObjects(ctx context.Context, keys []string) error {
|
||||
if h.cos == nil {
|
||||
return nil
|
||||
}
|
||||
for _, key := range keys {
|
||||
if err := h.cos.Delete(ctx, key); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// cleanupUploadedObjects 清理提交失败前已经上传的参考图,并统一响应清理错误。
|
||||
func (h *ProductImage) cleanupUploadedObjects(c *gin.Context, keys []string) bool {
|
||||
if err := h.deleteObjects(c.Request.Context(), keys); err != nil {
|
||||
h.respondError(c, err)
|
||||
return false
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
// cleanProductImageContentType 去除图片生成上传类型的参数部分。
|
||||
func cleanProductImageContentType(value string) string {
|
||||
return strings.TrimSpace(strings.Split(value, ";")[0])
|
||||
}
|
||||
|
||||
// respondError 将图片生成模块错误转换为一致的 Web API 响应。
|
||||
func (h *ProductImage) respondError(c *gin.Context, err error) {
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
fail(c, http.StatusNotFound, "not_found", "图片生成记录不存在")
|
||||
return
|
||||
}
|
||||
message := strings.TrimSpace(err.Error())
|
||||
if message == "" {
|
||||
message = "请求处理失败"
|
||||
}
|
||||
if strings.Contains(message, "不能为空") || strings.Contains(message, "不能") || strings.Contains(message, "无效") || strings.Contains(message, "请选择") || strings.Contains(message, "最多") || strings.Contains(message, "不可用") || strings.Contains(message, "未配置") {
|
||||
fail(c, http.StatusBadRequest, "invalid_request", message)
|
||||
return
|
||||
}
|
||||
if strings.Contains(message, "积分") {
|
||||
fail(c, http.StatusConflict, "insufficient_points", message)
|
||||
return
|
||||
}
|
||||
if strings.Contains(message, "生成中") {
|
||||
fail(c, http.StatusConflict, "generation_active", message)
|
||||
return
|
||||
}
|
||||
fail(c, http.StatusInternalServerError, "internal_error", message)
|
||||
}
|
||||
@@ -0,0 +1,29 @@
|
||||
package handler
|
||||
|
||||
import "testing"
|
||||
|
||||
func TestValidateRedemptionBatchLimits(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
points any
|
||||
quantity int
|
||||
wantErr bool
|
||||
}{
|
||||
{name: "minimum", points: 1, quantity: 1},
|
||||
{name: "maximum", points: 100, quantity: 10},
|
||||
{name: "decimal points", points: 1.5, quantity: 2},
|
||||
{name: "points below minimum", points: 0.99, quantity: 1, wantErr: true},
|
||||
{name: "points above maximum", points: 100.01, quantity: 1, wantErr: true},
|
||||
{name: "quantity below minimum", points: 1, quantity: 0, wantErr: true},
|
||||
{name: "quantity above maximum", points: 1, quantity: 11, wantErr: true},
|
||||
{name: "invalid points", points: "invalid", quantity: 1, wantErr: true},
|
||||
}
|
||||
for _, test := range tests {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
_, err := validateRedemptionBatchLimits(test.points, test.quantity)
|
||||
if (err != nil) != test.wantErr {
|
||||
t.Fatalf("validateRedemptionBatchLimits(%v, %d) error = %v", test.points, test.quantity, err)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,257 @@
|
||||
// 视频重绘接口处理器,负责重绘工作台、分析任务和源媒体上传请求。
|
||||
package handler
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
|
||||
mediakey "juhe-factory/api/internal/media"
|
||||
"juhe-factory/api/internal/service"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
func (h *Creative) GetRedrawWorkbench(c *gin.Context) {
|
||||
projectID, ok := h.projectID(c)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
data, err := h.service.GetRedrawWorkbench(currentWebUser(c).ID, projectID)
|
||||
if err != nil {
|
||||
h.respondError(c, err)
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"data": data})
|
||||
}
|
||||
|
||||
func (h *Creative) GetRedrawEpisodeWorkbench(c *gin.Context) {
|
||||
projectID, ok := h.projectID(c)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
episodeID, err := service.ParseUUID(c.Param("episode_id"), "剧集")
|
||||
if err != nil {
|
||||
h.bad(c, err)
|
||||
return
|
||||
}
|
||||
data, err := h.service.GetRedrawEpisodeWorkbench(currentWebUser(c).ID, projectID, episodeID)
|
||||
if err != nil {
|
||||
h.respondError(c, err)
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"data": data})
|
||||
}
|
||||
|
||||
func (h *Creative) QueueRedrawEpisodeScript(c *gin.Context) {
|
||||
projectID, ok := h.projectID(c)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
episodeID, err := service.ParseUUID(c.Param("episode_id"), "剧集")
|
||||
if err != nil {
|
||||
h.bad(c, err)
|
||||
return
|
||||
}
|
||||
task, err := h.service.QueueRedrawEpisodeScript(currentWebUser(c).ID, projectID, episodeID)
|
||||
if err != nil {
|
||||
h.respondError(c, err)
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusAccepted, gin.H{"data": task})
|
||||
}
|
||||
|
||||
// SaveRedrawEpisodeScript 保存用户修改后的剧集反推剧本正文。
|
||||
func (h *Creative) SaveRedrawEpisodeScript(c *gin.Context) {
|
||||
projectID, ok := h.projectID(c)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
episodeID, err := service.ParseUUID(c.Param("episode_id"), "剧集")
|
||||
if err != nil {
|
||||
h.bad(c, err)
|
||||
return
|
||||
}
|
||||
var body struct {
|
||||
Content string `json:"content"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&body); err != nil {
|
||||
h.bad(c, err)
|
||||
return
|
||||
}
|
||||
if err := h.service.SaveRedrawEpisodeScript(currentWebUser(c).ID, projectID, episodeID, body.Content); err != nil {
|
||||
h.respondError(c, err)
|
||||
return
|
||||
}
|
||||
c.Status(http.StatusNoContent)
|
||||
}
|
||||
|
||||
func (h *Creative) UpdateRedrawAnalysisSettings(c *gin.Context) {
|
||||
projectID, ok := h.projectID(c)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
var body struct {
|
||||
AudioSource string `json:"audio_source"`
|
||||
SourceLanguage string `json:"source_language"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&body); err != nil {
|
||||
h.bad(c, err)
|
||||
return
|
||||
}
|
||||
if err := h.service.UpdateRedrawAnalysisSettings(currentWebUser(c).ID, projectID, body.AudioSource, body.SourceLanguage); err != nil {
|
||||
h.respondError(c, err)
|
||||
return
|
||||
}
|
||||
c.Status(http.StatusNoContent)
|
||||
}
|
||||
|
||||
func (h *Creative) QueueRedrawAnalysis(c *gin.Context) {
|
||||
projectID, ok := h.projectID(c)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
task, err := h.service.QueueRedrawAnalysis(currentWebUser(c).ID, projectID)
|
||||
if err != nil {
|
||||
h.respondError(c, err)
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusAccepted, gin.H{"data": task})
|
||||
}
|
||||
|
||||
func (h *Creative) ReanalyzeRedraw(c *gin.Context) {
|
||||
if !h.objectStorageAvailable(c) {
|
||||
return
|
||||
}
|
||||
projectID, ok := h.projectID(c)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
objectKeys, err := h.service.ResetRedrawAnalysis(currentWebUser(c).ID, projectID)
|
||||
if err != nil {
|
||||
h.respondError(c, err)
|
||||
return
|
||||
}
|
||||
if !h.deleteStoredObjects(c, objectKeys) {
|
||||
return
|
||||
}
|
||||
task, err := h.service.QueueRedrawAnalysis(currentWebUser(c).ID, projectID)
|
||||
if err != nil {
|
||||
h.respondError(c, err)
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusAccepted, gin.H{"data": task})
|
||||
}
|
||||
|
||||
func (h *Creative) QueueRedrawStoryboardAnalysis(c *gin.Context) {
|
||||
projectID, ok := h.projectID(c)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
storyboardID, err := service.ParseUUID(c.Param("storyboard_id"), "分镜")
|
||||
if err != nil {
|
||||
h.bad(c, err)
|
||||
return
|
||||
}
|
||||
task, err := h.service.QueueRedrawStoryboardAnalysis(currentWebUser(c).ID, projectID, storyboardID)
|
||||
if err != nil {
|
||||
h.respondError(c, err)
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusAccepted, gin.H{"data": task})
|
||||
}
|
||||
|
||||
func (h *Creative) QueueRedrawScript(c *gin.Context) {
|
||||
projectID, ok := h.projectID(c)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
task, err := h.service.QueueRedrawScript(currentWebUser(c).ID, projectID)
|
||||
if err != nil {
|
||||
h.respondError(c, err)
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusAccepted, gin.H{"data": task})
|
||||
}
|
||||
|
||||
func (h *Creative) UploadRedrawSource(c *gin.Context) {
|
||||
h.uploadRedrawMedia(c, "source")
|
||||
}
|
||||
|
||||
func (h *Creative) UploadRedrawSubtitle(c *gin.Context) {
|
||||
h.uploadRedrawMedia(c, "subtitle")
|
||||
}
|
||||
|
||||
func (h *Creative) uploadRedrawMedia(c *gin.Context, kind string) {
|
||||
if h.cos == nil {
|
||||
fail(c, http.StatusServiceUnavailable, "cos_not_configured", "COS 对象存储未配置")
|
||||
return
|
||||
}
|
||||
projectID, ok := h.projectID(c)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
target, err := h.service.PrepareRedrawMediaUpload(currentWebUser(c).ID, projectID, kind)
|
||||
if err != nil {
|
||||
h.respondError(c, err)
|
||||
return
|
||||
}
|
||||
project := target.Project
|
||||
file, header, err := c.Request.FormFile("file")
|
||||
if err != nil {
|
||||
h.bad(c, errors.New("请选择文件"))
|
||||
return
|
||||
}
|
||||
defer file.Close()
|
||||
contentType := cleanContentType(header.Header.Get("Content-Type"))
|
||||
maxBytes := int64(50 * 1024 * 1024)
|
||||
ext := strings.ToLower(filepath.Ext(header.Filename))
|
||||
allowedVideoTypes := map[string]bool{"video/mp4": true, "video/webm": true, "video/quicktime": true, "video/x-m4v": true}
|
||||
allowed := allowedVideoTypes[contentType] && map[string]bool{".mp4": true, ".webm": true, ".mov": true, ".m4v": true}[ext]
|
||||
if kind == "subtitle" {
|
||||
maxBytes = 2 * 1024 * 1024
|
||||
allowed = ext == ".srt" || ext == ".vtt" || ext == ".txt"
|
||||
if contentType == "" {
|
||||
contentType = "text/plain"
|
||||
}
|
||||
}
|
||||
if !allowed || header.Size <= 0 || header.Size > maxBytes {
|
||||
h.bad(c, fmt.Errorf("%s文件格式或大小不符合要求", map[bool]string{true: "字幕", false: "视频"}[kind == "subtitle"]))
|
||||
return
|
||||
}
|
||||
oldMediaID := target.OldMediaID
|
||||
if kind == "source" && project.SourceVideoAssetID != nil {
|
||||
objectKeys, err := h.service.ResetRedrawAnalysis(currentWebUser(c).ID, projectID)
|
||||
if err != nil {
|
||||
h.respondError(c, err)
|
||||
return
|
||||
}
|
||||
if !h.deleteStoredObjects(c, objectKeys) {
|
||||
return
|
||||
}
|
||||
}
|
||||
mediaID := uuid.New()
|
||||
key := mediakey.ProjectSource(projectID, mediaID, header.Filename, contentType)
|
||||
displayName := project.Name + "原视频"
|
||||
if kind == "subtitle" {
|
||||
key = mediakey.ProjectSubtitle(projectID, mediaID, header.Filename, contentType)
|
||||
displayName = project.Name + "字幕"
|
||||
}
|
||||
asset, err := h.storeUpload(c, mediaID, currentWebUser(c).ID, file, header, key, displayName, contentType, nil)
|
||||
if err != nil {
|
||||
h.internal(c, err)
|
||||
return
|
||||
}
|
||||
if err := h.service.PersistRedrawMediaUpload(projectID, kind, project.RedrawStatus, asset, oldMediaID); err != nil {
|
||||
_ = h.cos.Delete(c.Request.Context(), asset.ObjectKey)
|
||||
h.internal(c, err)
|
||||
return
|
||||
}
|
||||
if target.OldObjectKey != "" && !h.deleteStoredObjects(c, []string{target.OldObjectKey}) {
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusCreated, gin.H{"data": asset})
|
||||
}
|
||||
@@ -0,0 +1,166 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"io"
|
||||
"net/http"
|
||||
|
||||
"juhe-factory/api/internal/service"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
func (h *Creative) ListScriptAnalyses(c *gin.Context) {
|
||||
items, err := h.service.ListScriptAnalyses(currentWebUser(c).ID)
|
||||
if err != nil {
|
||||
h.respondError(c, err)
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"data": items})
|
||||
}
|
||||
|
||||
func (h *Creative) CreateScriptAnalysis(c *gin.Context) {
|
||||
item, err := h.service.CreateScriptAnalysis(currentWebUser(c).ID)
|
||||
if err != nil {
|
||||
h.respondError(c, err)
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusCreated, gin.H{"data": item})
|
||||
}
|
||||
|
||||
func (h *Creative) ListScriptAnalysisImportProjects(c *gin.Context) {
|
||||
items, err := h.service.ListScriptAnalysisImportProjects(currentWebUser(c).ID)
|
||||
if err != nil {
|
||||
h.respondError(c, err)
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"data": items})
|
||||
}
|
||||
|
||||
func (h *Creative) ImportScriptAnalysisProject(c *gin.Context) {
|
||||
id, err := service.ParseUUID(c.Param("script_id"), "剧本")
|
||||
if err != nil {
|
||||
h.bad(c, err)
|
||||
return
|
||||
}
|
||||
var body struct {
|
||||
ProjectID uuid.UUID `json:"project_id"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&body); err != nil || body.ProjectID == uuid.Nil {
|
||||
h.bad(c, errors.New("请选择要导入的剧本反推项目"))
|
||||
return
|
||||
}
|
||||
item, err := h.service.ImportScriptAnalysisProject(currentWebUser(c).ID, id, body.ProjectID)
|
||||
if err != nil {
|
||||
h.respondError(c, err)
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"data": item})
|
||||
}
|
||||
|
||||
func (h *Creative) ImportScriptAnalysisFile(c *gin.Context) {
|
||||
id, err := service.ParseUUID(c.Param("script_id"), "剧本")
|
||||
if err != nil {
|
||||
h.bad(c, err)
|
||||
return
|
||||
}
|
||||
c.Request.Body = http.MaxBytesReader(c.Writer, c.Request.Body, 4<<20)
|
||||
if err := c.Request.ParseMultipartForm(4 << 20); err != nil {
|
||||
h.bad(c, errors.New("导入文件无效"))
|
||||
return
|
||||
}
|
||||
if c.Request.MultipartForm != nil {
|
||||
defer c.Request.MultipartForm.RemoveAll()
|
||||
}
|
||||
file, header, err := c.Request.FormFile("file")
|
||||
if err != nil {
|
||||
h.bad(c, errors.New("请选择剧本文件"))
|
||||
return
|
||||
}
|
||||
defer file.Close()
|
||||
content, err := io.ReadAll(io.LimitReader(file, 3*1024*1024+1))
|
||||
if err != nil {
|
||||
h.bad(c, errors.New("读取剧本文件失败"))
|
||||
return
|
||||
}
|
||||
item, err := h.service.ImportScriptAnalysisFile(currentWebUser(c).ID, id, header.Filename, content)
|
||||
if err != nil {
|
||||
h.respondError(c, err)
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"data": item})
|
||||
}
|
||||
|
||||
func (h *Creative) GetScriptAnalysis(c *gin.Context) {
|
||||
id, err := service.ParseUUID(c.Param("script_id"), "剧本")
|
||||
if err != nil {
|
||||
h.bad(c, err)
|
||||
return
|
||||
}
|
||||
item, err := h.service.GetScriptAnalysis(currentWebUser(c).ID, id)
|
||||
if err != nil {
|
||||
h.respondError(c, err)
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"data": item})
|
||||
}
|
||||
|
||||
func (h *Creative) UpdateScriptAnalysis(c *gin.Context) {
|
||||
id, err := service.ParseUUID(c.Param("script_id"), "剧本")
|
||||
if err != nil {
|
||||
h.bad(c, err)
|
||||
return
|
||||
}
|
||||
var input service.ScriptAnalysisInput
|
||||
if err := c.ShouldBindJSON(&input); err != nil {
|
||||
h.bad(c, err)
|
||||
return
|
||||
}
|
||||
item, err := h.service.UpdateScriptAnalysis(currentWebUser(c).ID, id, input)
|
||||
if err != nil {
|
||||
h.respondError(c, err)
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"data": item})
|
||||
}
|
||||
|
||||
func (h *Creative) DeleteScriptAnalysis(c *gin.Context) {
|
||||
id, err := service.ParseUUID(c.Param("script_id"), "剧本")
|
||||
if err != nil {
|
||||
h.bad(c, err)
|
||||
return
|
||||
}
|
||||
if err := h.service.DeleteScriptAnalysis(currentWebUser(c).ID, id); err != nil {
|
||||
h.respondError(c, err)
|
||||
return
|
||||
}
|
||||
c.Status(http.StatusNoContent)
|
||||
}
|
||||
|
||||
func (h *Creative) QueueScriptAnalysis(c *gin.Context) {
|
||||
id, err := service.ParseUUID(c.Param("script_id"), "剧本")
|
||||
if err != nil {
|
||||
h.bad(c, err)
|
||||
return
|
||||
}
|
||||
task, err := h.service.QueueScriptAnalysis(currentWebUser(c).ID, id)
|
||||
if err != nil {
|
||||
h.respondError(c, err)
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusAccepted, gin.H{"data": task})
|
||||
}
|
||||
|
||||
func (h *Creative) CancelScriptAnalysis(c *gin.Context) {
|
||||
id, err := service.ParseUUID(c.Param("script_id"), "剧本")
|
||||
if err != nil {
|
||||
h.bad(c, err)
|
||||
return
|
||||
}
|
||||
if err := h.service.CancelScriptAnalysis(currentWebUser(c).ID, id); err != nil {
|
||||
h.respondError(c, err)
|
||||
return
|
||||
}
|
||||
c.Status(http.StatusNoContent)
|
||||
}
|
||||
@@ -0,0 +1,226 @@
|
||||
// WEB 用户接口处理器,负责认证、账户资料、积分查询、头像上传和兑换请求。
|
||||
package handler
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
mediakey "juhe-factory/api/internal/media"
|
||||
"juhe-factory/api/internal/model"
|
||||
"juhe-factory/api/internal/service"
|
||||
"juhe-factory/api/internal/storage"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
const webUserContextKey = "web_user"
|
||||
|
||||
type Web struct {
|
||||
service *service.Web
|
||||
cos *storage.COS
|
||||
}
|
||||
|
||||
func NewWeb(service *service.Web, cos *storage.COS) *Web { return &Web{service: service, cos: cos} }
|
||||
|
||||
func (h *Web) Login(c *gin.Context) {
|
||||
var body struct {
|
||||
Account string `json:"account" binding:"required"`
|
||||
Password string `json:"password" binding:"required"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&body); err != nil {
|
||||
fail(c, http.StatusBadRequest, "invalid_request", "请输入账号和密码")
|
||||
return
|
||||
}
|
||||
pair, err := h.service.Login(body.Account, body.Password)
|
||||
if err != nil {
|
||||
fail(c, http.StatusUnauthorized, "invalid_credentials", err.Error())
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"data": pair})
|
||||
}
|
||||
func (h *Web) Refresh(c *gin.Context) {
|
||||
var body struct {
|
||||
RefreshToken string `json:"refresh_token" binding:"required"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&body); err != nil {
|
||||
fail(c, http.StatusBadRequest, "invalid_request", "缺少刷新令牌")
|
||||
return
|
||||
}
|
||||
pair, err := h.service.Refresh(body.RefreshToken)
|
||||
if err != nil {
|
||||
fail(c, http.StatusUnauthorized, "invalid_refresh_token", err.Error())
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"data": pair})
|
||||
}
|
||||
func (h *Web) Logout(c *gin.Context) {
|
||||
var body struct {
|
||||
RefreshToken string `json:"refresh_token"`
|
||||
}
|
||||
_ = c.ShouldBindJSON(&body)
|
||||
if err := h.service.Logout(body.RefreshToken); err != nil {
|
||||
fail(c, http.StatusInternalServerError, "logout_failed", "退出失败,请重试")
|
||||
return
|
||||
}
|
||||
c.Status(http.StatusNoContent)
|
||||
}
|
||||
func (h *Web) Middleware() gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
raw := strings.TrimSpace(strings.TrimPrefix(c.GetHeader("Authorization"), "Bearer "))
|
||||
user, err := h.service.Authenticate(raw)
|
||||
if err != nil {
|
||||
fail(c, http.StatusUnauthorized, "unauthorized", err.Error())
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
c.Set(webUserContextKey, user)
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
|
||||
// Account 返回当前用户账户资料、积分额度和头像上传限制。
|
||||
func (h *Web) Account(c *gin.Context) {
|
||||
data, err := h.service.Account(currentWebUser(c).ID)
|
||||
if err != nil {
|
||||
fail(c, http.StatusInternalServerError, "account_failed", "账户信息加载失败")
|
||||
return
|
||||
}
|
||||
if h.cos != nil {
|
||||
data["avatar_max_size_bytes"] = h.cos.MaxImageBytes()
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"data": data})
|
||||
}
|
||||
func (h *Web) UpdateProfile(c *gin.Context) {
|
||||
var body struct {
|
||||
Username string `json:"username" binding:"required"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&body); err != nil {
|
||||
fail(c, http.StatusBadRequest, "invalid_request", "请输入用户名")
|
||||
return
|
||||
}
|
||||
data, err := h.service.UpdateProfile(currentWebUser(c).ID, body.Username)
|
||||
if err != nil {
|
||||
if errors.Is(err, service.ErrInvalidUsername) || errors.Is(err, service.ErrUsernameExists) {
|
||||
fail(c, http.StatusBadRequest, "profile_invalid", err.Error())
|
||||
return
|
||||
}
|
||||
fail(c, http.StatusInternalServerError, "profile_update_failed", "资料更新失败,请稍后重试")
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"data": data})
|
||||
}
|
||||
func (h *Web) ChangePassword(c *gin.Context) {
|
||||
var body struct {
|
||||
OriginalPassword string `json:"original_password" binding:"required"`
|
||||
NewPassword string `json:"new_password" binding:"required"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&body); err != nil {
|
||||
fail(c, http.StatusBadRequest, "invalid_request", "请完整填写密码")
|
||||
return
|
||||
}
|
||||
if err := h.service.ChangePassword(currentWebUser(c).ID, body.OriginalPassword, body.NewPassword); err != nil {
|
||||
if errors.Is(err, service.ErrOriginalPassword) || strings.Contains(err.Error(), "密码长度") {
|
||||
fail(c, http.StatusBadRequest, "password_invalid", err.Error())
|
||||
return
|
||||
}
|
||||
fail(c, http.StatusInternalServerError, "password_update_failed", "密码修改失败,请稍后重试")
|
||||
return
|
||||
}
|
||||
c.Status(http.StatusNoContent)
|
||||
}
|
||||
func (h *Web) UploadAvatar(c *gin.Context) {
|
||||
if h.cos == nil {
|
||||
fail(c, http.StatusServiceUnavailable, "cos_not_configured", "对象存储未配置")
|
||||
return
|
||||
}
|
||||
file, header, err := c.Request.FormFile("file")
|
||||
if err != nil {
|
||||
fail(c, http.StatusBadRequest, "invalid_request", "请选择头像图片")
|
||||
return
|
||||
}
|
||||
defer file.Close()
|
||||
contentType := strings.TrimSpace(strings.Split(header.Header.Get("Content-Type"), ";")[0])
|
||||
extensions := map[string]string{"image/jpeg": ".jpg", "image/png": ".png", "image/webp": ".webp"}
|
||||
extension, allowed := extensions[contentType]
|
||||
if !allowed || header.Size <= 0 || header.Size > h.cos.MaxImageBytes() {
|
||||
fail(c, http.StatusBadRequest, "avatar_invalid", "头像仅支持 JPEG、PNG 或 WebP,且不得超过图片大小限制")
|
||||
return
|
||||
}
|
||||
headerBytes := make([]byte, 512)
|
||||
read, readErr := file.Read(headerBytes)
|
||||
if readErr != nil && !errors.Is(readErr, io.EOF) {
|
||||
fail(c, http.StatusBadRequest, "avatar_invalid", "头像图片内容无法解析")
|
||||
return
|
||||
}
|
||||
if detected := http.DetectContentType(headerBytes[:read]); detected != contentType {
|
||||
fail(c, http.StatusBadRequest, "avatar_invalid", "头像图片内容无法解析")
|
||||
return
|
||||
}
|
||||
if _, err := file.Seek(0, io.SeekStart); err != nil {
|
||||
fail(c, http.StatusInternalServerError, "avatar_upload_failed", "头像上传失败")
|
||||
return
|
||||
}
|
||||
userID := currentWebUser(c).ID
|
||||
oldKey, err := h.service.AvatarKey(userID)
|
||||
if err != nil {
|
||||
fail(c, http.StatusInternalServerError, "avatar_upload_failed", "头像上传失败")
|
||||
return
|
||||
}
|
||||
key := mediakey.UserAvatar(userID, uuid.New(), extension)
|
||||
url, err := h.cos.Put(c.Request.Context(), key, contentType, file, header.Size)
|
||||
if err != nil {
|
||||
fail(c, http.StatusInternalServerError, "avatar_upload_failed", "头像上传失败")
|
||||
return
|
||||
}
|
||||
if err := h.service.UpdateAvatar(userID, url, key); err != nil {
|
||||
_ = h.cos.Delete(c.Request.Context(), key)
|
||||
fail(c, http.StatusInternalServerError, "avatar_upload_failed", "头像上传失败")
|
||||
return
|
||||
}
|
||||
if oldKey != "" {
|
||||
_ = h.cos.Delete(c.Request.Context(), oldKey)
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"data": gin.H{"avatar_url": url}})
|
||||
}
|
||||
func (h *Web) Usage30Days(c *gin.Context) {
|
||||
items, err := h.service.Usage30Days(currentWebUser(c).ID)
|
||||
if err != nil {
|
||||
fail(c, http.StatusInternalServerError, "usage_failed", "消耗数据加载失败")
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"data": items})
|
||||
}
|
||||
|
||||
// ConsumptionRecords 按请求页码返回当前用户的积分变动流水。
|
||||
func (h *Web) ConsumptionRecords(c *gin.Context) {
|
||||
page, pageSize := parsePage(c)
|
||||
result, err := h.service.ConsumptionRecords(currentWebUser(c).ID, page, pageSize)
|
||||
if err != nil {
|
||||
fail(c, http.StatusInternalServerError, "consumption_records_failed", "消耗记录加载失败")
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"data": result})
|
||||
}
|
||||
func (h *Web) Redeem(c *gin.Context) {
|
||||
var body struct {
|
||||
Code string `json:"code" binding:"required"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&body); err != nil {
|
||||
fail(c, http.StatusBadRequest, "invalid_request", "请输入兑换码")
|
||||
return
|
||||
}
|
||||
result, err := h.service.Redeem(currentWebUser(c).ID, body.Code)
|
||||
if err != nil {
|
||||
fail(c, http.StatusBadRequest, "redeem_failed", err.Error())
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"data": result})
|
||||
}
|
||||
func currentWebUser(c *gin.Context) *model.WebUser {
|
||||
value, _ := c.Get(webUserContextKey)
|
||||
user, _ := value.(*model.WebUser)
|
||||
return user
|
||||
}
|
||||
@@ -0,0 +1,156 @@
|
||||
// Package media 统一生成可长期稳定引用的 COS 媒体对象键和扩展名。
|
||||
package media
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"path/filepath"
|
||||
"regexp"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
const objectRoot = "juyou_ran"
|
||||
|
||||
var unsafeKeyPart = regexp.MustCompile(`[^a-z0-9-]+`)
|
||||
var safeExtension = regexp.MustCompile(`^\.[a-z0-9]+$`)
|
||||
|
||||
// objectKey 为所有媒体对象添加统一的业务根目录,避免不同业务直接污染存储桶根目录。
|
||||
func objectKey(format string, args ...any) string {
|
||||
return fmt.Sprintf(objectRoot+"/"+format, args...)
|
||||
}
|
||||
|
||||
// Slug 将业务名称转换为受控英文路径片段;展示名称保留在数据库中。
|
||||
func Slug(value, fallback string) string {
|
||||
value = strings.ToLower(strings.TrimSpace(value))
|
||||
value = unsafeKeyPart.ReplaceAllString(value, "-")
|
||||
value = strings.Trim(value, "-")
|
||||
if value == "" {
|
||||
value = fallback
|
||||
}
|
||||
if len(value) > 48 {
|
||||
value = strings.Trim(value[:48], "-")
|
||||
}
|
||||
return value
|
||||
}
|
||||
|
||||
// Extension 根据媒体类型或安全的原文件扩展名生成对象扩展名。
|
||||
func Extension(filename, contentType string) string {
|
||||
switch strings.ToLower(contentType) {
|
||||
case "image/jpeg":
|
||||
return ".jpg"
|
||||
case "image/png":
|
||||
return ".png"
|
||||
case "image/webp":
|
||||
return ".webp"
|
||||
case "image/gif":
|
||||
return ".gif"
|
||||
case "video/mp4":
|
||||
return ".mp4"
|
||||
case "video/webm":
|
||||
return ".webm"
|
||||
case "audio/mpeg":
|
||||
return ".mp3"
|
||||
case "audio/wav", "audio/x-wav":
|
||||
return ".wav"
|
||||
case "text/vtt":
|
||||
return ".vtt"
|
||||
}
|
||||
ext := strings.ToLower(filepath.Ext(filename))
|
||||
if len(ext) > 1 && len(ext) <= 8 && safeExtension.MatchString(ext) {
|
||||
return ext
|
||||
}
|
||||
return ".bin"
|
||||
}
|
||||
|
||||
// EpisodeSource 返回剧集原视频的对象键。
|
||||
func EpisodeSource(projectID uuid.UUID, episodeNo int, episodeID, assetID uuid.UUID, filename, contentType string) string {
|
||||
return objectKey("video-redraw/projects/%s/episodes/%03d-%s/source/%s-source%s", projectID, episodeNo, episodeID, assetID, Extension(filename, contentType))
|
||||
}
|
||||
|
||||
// EpisodeSubtitle 返回剧集字幕的对象键。
|
||||
func EpisodeSubtitle(projectID uuid.UUID, episodeNo int, episodeID, assetID uuid.UUID, filename, contentType string) string {
|
||||
return objectKey("video-redraw/projects/%s/episodes/%03d-%s/subtitles/%s-subtitle%s", projectID, episodeNo, episodeID, assetID, Extension(filename, contentType))
|
||||
}
|
||||
|
||||
// ProjectSource 返回无剧集项目原视频的对象键。
|
||||
func ProjectSource(projectID, assetID uuid.UUID, filename, contentType string) string {
|
||||
return objectKey("video-redraw/projects/%s/source/%s-source%s", projectID, assetID, Extension(filename, contentType))
|
||||
}
|
||||
|
||||
// ProjectSubtitle 返回无剧集项目字幕的对象键。
|
||||
func ProjectSubtitle(projectID, assetID uuid.UUID, filename, contentType string) string {
|
||||
return objectKey("video-redraw/projects/%s/subtitles/%s-subtitle%s", projectID, assetID, Extension(filename, contentType))
|
||||
}
|
||||
|
||||
// AssetReference 返回项目资产参考媒体的对象键。
|
||||
func AssetReference(projectID uuid.UUID, assetType string, assetID, mediaID uuid.UUID, filename, contentType string) string {
|
||||
return objectKey("video-redraw/projects/%s/assets/%s/%s/%s-reference%s", projectID, Slug(assetType, "custom"), assetID, mediaID, Extension(filename, contentType))
|
||||
}
|
||||
|
||||
// ProductImageReference 返回独立图片生成工具参考图片的稳定对象键。
|
||||
func ProductImageReference(userID, mediaID uuid.UUID, filename, contentType string) string {
|
||||
return objectKey("product-images/users/%s/references/%s%s", userID, mediaID, Extension(filename, contentType))
|
||||
}
|
||||
|
||||
// AssetOutput 返回项目资产生成结果的对象键。
|
||||
func AssetOutput(projectID uuid.UUID, assetType string, assetID, taskID, mediaID uuid.UUID, contentType string) string {
|
||||
return objectKey("video-redraw/projects/%s/assets/%s/%s/outputs/%s-%s-result%s", projectID, Slug(assetType, "custom"), assetID, taskID, mediaID, Extension("", contentTypeFor("image", contentType)))
|
||||
}
|
||||
|
||||
// StoryboardFrame 返回剧集分镜抽帧的对象键。
|
||||
func StoryboardFrame(projectID, episodeID, storyboardID, mediaID uuid.UUID, sequence int, timestampMS int64) string {
|
||||
return objectKey("video-redraw/projects/%s/episodes/%s/storyboards/%04d-%s/frames/%s-%010d.jpg", projectID, episodeID, sequence, storyboardID, mediaID, timestampMS)
|
||||
}
|
||||
|
||||
// StoryboardAnalysis 返回剧集分镜反推原始响应的对象键。
|
||||
func StoryboardAnalysis(projectID, episodeID, storyboardID, taskID, mediaID uuid.UUID, sequence int) string {
|
||||
return objectKey("video-redraw/projects/%s/episodes/%s/storyboards/%04d-%s/analysis/%s-%s-reverse.json", projectID, episodeID, sequence, storyboardID, taskID, mediaID)
|
||||
}
|
||||
|
||||
// TaskOutput 返回剧集分镜生成结果的对象键。
|
||||
func TaskOutput(projectID, episodeID uuid.UUID, sequence int, taskID, mediaID uuid.UUID, outputType, contentType string) string {
|
||||
return objectKey("video-redraw/projects/%s/episodes/%s/storyboards/%04d/outputs/%s-%s-result%s", projectID, episodeID, sequence, taskID, mediaID, Extension("", contentTypeFor(outputType, contentType)))
|
||||
}
|
||||
|
||||
// ProjectStoryboardFrame 返回无剧集项目分镜抽帧的对象键。
|
||||
func ProjectStoryboardFrame(projectID, storyboardID, mediaID uuid.UUID, sequence int, timestampMS int64) string {
|
||||
return objectKey("video-redraw/projects/%s/storyboards/%04d-%s/frames/%s-%010d.jpg", projectID, sequence, storyboardID, mediaID, timestampMS)
|
||||
}
|
||||
|
||||
// ProjectStoryboardAnalysis 返回无剧集项目分镜反推原始响应的对象键。
|
||||
func ProjectStoryboardAnalysis(projectID, storyboardID, taskID, mediaID uuid.UUID, sequence int) string {
|
||||
return objectKey("video-redraw/projects/%s/storyboards/%04d-%s/analysis/%s-%s-reverse.json", projectID, sequence, storyboardID, taskID, mediaID)
|
||||
}
|
||||
|
||||
// ProjectTaskOutput 返回无剧集项目分镜生成结果的对象键。
|
||||
func ProjectTaskOutput(projectID uuid.UUID, sequence int, taskID, mediaID uuid.UUID, outputType, contentType string) string {
|
||||
return objectKey("video-redraw/projects/%s/storyboards/%04d/outputs/%s-%s-result%s", projectID, sequence, taskID, mediaID, Extension("", contentTypeFor(outputType, contentType)))
|
||||
}
|
||||
|
||||
// StyleImage 返回管理端风格图片的统一对象键。
|
||||
func StyleImage(mediaID uuid.UUID, extension string) string {
|
||||
return objectKey("styles/%s/%s%s", time.Now().Format("2006/01"), mediaID, extension)
|
||||
}
|
||||
|
||||
// UserAvatar 返回用户头像的统一对象键。
|
||||
func UserAvatar(userID, mediaID uuid.UUID, extension string) string {
|
||||
return objectKey("users/%s/avatar/%s%s", userID, mediaID, extension)
|
||||
}
|
||||
|
||||
// StandaloneTaskOutput 返回独立生成任务的结果对象键。
|
||||
func StandaloneTaskOutput(userID, taskID uuid.UUID, outputType, contentType string) string {
|
||||
return objectKey("users/%s/tasks/%s/output-1%s", userID, taskID, Extension("", contentTypeFor(outputType, contentType)))
|
||||
}
|
||||
|
||||
// contentTypeFor 返回已知内容类型,缺省时根据输出类型选择默认媒体类型。
|
||||
func contentTypeFor(outputType, contentType string) string {
|
||||
if contentType != "" {
|
||||
return contentType
|
||||
}
|
||||
if outputType == "image" {
|
||||
return "image/png"
|
||||
}
|
||||
return "video/mp4"
|
||||
}
|
||||
@@ -0,0 +1,17 @@
|
||||
package media
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
func TestAssetKeysUseControlledSegments(t *testing.T) {
|
||||
projectID, assetID, taskID, mediaID := uuid.New(), uuid.New(), uuid.New(), uuid.New()
|
||||
key := AssetOutput(projectID, "角色 中文", assetID, taskID, mediaID, "image/webp")
|
||||
wantPrefix := "juyou_ran/video-redraw/projects/" + projectID.String() + "/assets/custom/" + assetID.String() + "/outputs/"
|
||||
if !strings.HasPrefix(key, wantPrefix) || !strings.HasSuffix(key, ".webp") {
|
||||
t.Fatalf("unexpected key: %s", key)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,46 @@
|
||||
package model
|
||||
|
||||
import (
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
type AdminUser struct {
|
||||
ID uuid.UUID `gorm:"type:uuid;default:gen_random_uuid();primaryKey" json:"id"`
|
||||
Username string `json:"username"`
|
||||
PasswordHash string `json:"-"`
|
||||
Enabled bool `json:"enabled"`
|
||||
LastLoginAt *time.Time `json:"last_login_at"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
UpdatedAt time.Time `json:"updated_at"`
|
||||
DeletedAt gorm.DeletedAt `json:"-"`
|
||||
}
|
||||
|
||||
type AdminRefreshToken struct {
|
||||
ID uuid.UUID `gorm:"type:uuid;default:gen_random_uuid();primaryKey"`
|
||||
AdminID uuid.UUID `gorm:"type:uuid"`
|
||||
TokenHash string
|
||||
ExpiresAt time.Time
|
||||
RevokedAt *time.Time
|
||||
CreatedAt time.Time
|
||||
}
|
||||
|
||||
type AdminAuditLog struct {
|
||||
ID int64 `json:"id"`
|
||||
AdminID *uuid.UUID `gorm:"type:uuid" json:"admin_id"`
|
||||
AdminUsername string `json:"admin_username"`
|
||||
Action string `json:"action"`
|
||||
ResourceType string `json:"resource_type"`
|
||||
ResourceID string `json:"resource_id"`
|
||||
Reason string `json:"reason"`
|
||||
Detail string `gorm:"type:jsonb" json:"detail"`
|
||||
IPAddress string `gorm:"type:inet" json:"ip_address"`
|
||||
TraceID string `json:"trace_id"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
}
|
||||
|
||||
func (AdminUser) TableName() string { return "admin_users" }
|
||||
func (AdminRefreshToken) TableName() string { return "admin_refresh_tokens" }
|
||||
func (AdminAuditLog) TableName() string { return "admin_audit_logs" }
|
||||
@@ -0,0 +1,299 @@
|
||||
package model
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
type MediaAsset struct {
|
||||
ID uuid.UUID `gorm:"type:uuid;default:gen_random_uuid();primaryKey" json:"id"`
|
||||
OwnerUserID *uuid.UUID `gorm:"type:uuid" json:"owner_user_id"`
|
||||
StorageProvider string `json:"storage_provider"`
|
||||
ObjectKey string `json:"object_key"`
|
||||
PublicURL string `json:"public_url"`
|
||||
OriginalName string `json:"original_name"`
|
||||
DisplayName string `json:"display_name"`
|
||||
MimeType string `json:"mime_type"`
|
||||
SizeBytes int64 `json:"size_bytes"`
|
||||
SHA256 string `json:"sha256"`
|
||||
Width *int `json:"width"`
|
||||
Height *int `json:"height"`
|
||||
DurationMS *int64 `json:"duration_ms"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
DeletedAt *time.Time `json:"-"`
|
||||
}
|
||||
|
||||
type CreativeProject struct {
|
||||
ID uuid.UUID `gorm:"type:uuid;default:gen_random_uuid();primaryKey" json:"id"`
|
||||
UserID uuid.UUID `gorm:"type:uuid" json:"user_id"`
|
||||
ProjectType string `json:"project_type"`
|
||||
Name string `json:"name"`
|
||||
CoverAssetID *uuid.UUID `gorm:"type:uuid" json:"cover_asset_id"`
|
||||
StyleID uuid.UUID `gorm:"type:uuid" json:"style_id"`
|
||||
EraType string `json:"era_type"`
|
||||
CustomEra *string `json:"custom_era"`
|
||||
AspectRatio string `json:"aspect_ratio"`
|
||||
Localization *string `json:"localization"`
|
||||
SourceVideoAssetID *uuid.UUID `gorm:"type:uuid" json:"source_video_asset_id"`
|
||||
SubtitleAssetID *uuid.UUID `gorm:"type:uuid" json:"subtitle_asset_id"`
|
||||
AudioSource string `json:"audio_source"`
|
||||
SourceLanguage *string `json:"source_language"`
|
||||
RedrawStatus string `json:"redraw_status"`
|
||||
AnalysisMessage string `json:"analysis_message"`
|
||||
RedrawScript string `json:"redraw_script"`
|
||||
ShortDramaType *string `json:"short_drama_type"`
|
||||
PlayCount *string `json:"play_count"`
|
||||
AudienceProfile *string `json:"audience_profile"`
|
||||
Producer *string `json:"producer"`
|
||||
CastMembers json.RawMessage `gorm:"type:jsonb;not null;default:'[]'" json:"cast_members"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
UpdatedAt time.Time `json:"updated_at"`
|
||||
DeletedAt *time.Time `json:"-"`
|
||||
}
|
||||
|
||||
func (project *CreativeProject) BeforeCreate(_ *gorm.DB) error {
|
||||
if len(project.CastMembers) == 0 {
|
||||
project.CastMembers = json.RawMessage("[]")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
type ProjectModelConfig struct {
|
||||
ID uuid.UUID `gorm:"type:uuid;default:gen_random_uuid();primaryKey" json:"id"`
|
||||
ProjectID uuid.UUID `gorm:"type:uuid" json:"project_id"`
|
||||
Purpose string `json:"purpose"`
|
||||
ModelType string `json:"model_type"`
|
||||
ModelID uuid.UUID `gorm:"type:uuid" json:"model_id"`
|
||||
PromptID *uuid.UUID `gorm:"type:uuid" json:"prompt_id"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
UpdatedAt time.Time `json:"updated_at"`
|
||||
}
|
||||
|
||||
type UserModelConfig struct {
|
||||
ID uuid.UUID `gorm:"type:uuid;default:gen_random_uuid();primaryKey" json:"id"`
|
||||
UserID uuid.UUID `gorm:"type:uuid" json:"user_id"`
|
||||
Purpose string `json:"purpose"`
|
||||
ProjectType string `json:"project_type"`
|
||||
ModelType string `json:"model_type"`
|
||||
ModelID uuid.UUID `gorm:"type:uuid" json:"model_id"`
|
||||
PromptID *uuid.UUID `gorm:"type:uuid" json:"prompt_id"`
|
||||
Settings json.RawMessage `gorm:"type:jsonb" json:"settings"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
UpdatedAt time.Time `json:"updated_at"`
|
||||
}
|
||||
|
||||
type ProjectEpisode struct {
|
||||
ID uuid.UUID `gorm:"type:uuid;default:gen_random_uuid();primaryKey" json:"id"`
|
||||
ProjectID uuid.UUID `gorm:"type:uuid" json:"project_id"`
|
||||
EpisodeNo int `json:"episode_no"`
|
||||
Name string `json:"name"`
|
||||
CoverAssetID *uuid.UUID `gorm:"type:uuid" json:"cover_asset_id"`
|
||||
SourceVideoAssetID *uuid.UUID `gorm:"type:uuid" json:"source_video_asset_id"`
|
||||
SubtitleAssetID *uuid.UUID `gorm:"type:uuid" json:"subtitle_asset_id"`
|
||||
AudioSource string `json:"audio_source"`
|
||||
SourceLanguage *string `json:"source_language"`
|
||||
Status string `json:"status"`
|
||||
AnalysisMessage string `json:"analysis_message"`
|
||||
RedrawScript string `json:"redraw_script"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
UpdatedAt time.Time `json:"updated_at"`
|
||||
DeletedAt *time.Time `json:"-"`
|
||||
}
|
||||
|
||||
type ProjectAsset struct {
|
||||
ID uuid.UUID `gorm:"type:uuid;default:gen_random_uuid();primaryKey" json:"id"`
|
||||
ProjectID uuid.UUID `gorm:"type:uuid" json:"project_id"`
|
||||
SourceEpisodeID *uuid.UUID `gorm:"type:uuid" json:"source_episode_id"`
|
||||
AssetType string `json:"asset_type"`
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description"`
|
||||
ImagePrompt string `json:"image_prompt"`
|
||||
Appearances json.RawMessage `gorm:"type:jsonb" json:"appearances"`
|
||||
ImageAssetID *uuid.UUID `gorm:"type:uuid" json:"image_asset_id"`
|
||||
AudioAssetID *uuid.UUID `gorm:"type:uuid" json:"audio_asset_id"`
|
||||
UserEdited bool `json:"user_edited"`
|
||||
Aliases json.RawMessage `gorm:"type:jsonb" json:"aliases"`
|
||||
Attributes json.RawMessage `gorm:"type:jsonb" json:"attributes"`
|
||||
AIDescription string `gorm:"column:ai_description" json:"ai_description"`
|
||||
SourceParseTaskID *uuid.UUID `gorm:"type:uuid" json:"source_parse_task_id"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
UpdatedAt time.Time `json:"updated_at"`
|
||||
DeletedAt *time.Time `json:"-"`
|
||||
}
|
||||
|
||||
func (asset *ProjectAsset) BeforeCreate(_ *gorm.DB) error {
|
||||
if len(asset.Appearances) == 0 {
|
||||
asset.Appearances = json.RawMessage("[]")
|
||||
}
|
||||
if len(asset.Aliases) == 0 {
|
||||
asset.Aliases = json.RawMessage("[]")
|
||||
}
|
||||
if len(asset.Attributes) == 0 {
|
||||
asset.Attributes = json.RawMessage("{}")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
type EpisodeStoryboard struct {
|
||||
ID uuid.UUID `gorm:"type:uuid;default:gen_random_uuid();primaryKey" json:"id"`
|
||||
EpisodeID *uuid.UUID `gorm:"type:uuid" json:"episode_id"`
|
||||
ProjectID *uuid.UUID `gorm:"type:uuid" json:"project_id"`
|
||||
SequenceNo int `json:"sequence_no"`
|
||||
StableKey string `json:"stable_key"`
|
||||
StartMS int64 `json:"start_ms"`
|
||||
EndMS int64 `json:"end_ms"`
|
||||
DurationSeconds int `json:"duration_seconds"`
|
||||
Title string `json:"title"`
|
||||
ThumbnailAssetID *uuid.UUID `gorm:"type:uuid" json:"thumbnail_asset_id"`
|
||||
PromptContent string `json:"prompt_content"`
|
||||
ImagePrompt string `json:"image_prompt"`
|
||||
ScriptContent string `json:"script_content"`
|
||||
SourceExcerpt string `json:"source_excerpt"`
|
||||
Dialogue json.RawMessage `gorm:"type:jsonb" json:"dialogue"`
|
||||
AssetRefs json.RawMessage `gorm:"type:jsonb" json:"asset_refs"`
|
||||
Locked bool `json:"locked"`
|
||||
ActiveOutputID *uuid.UUID `gorm:"type:uuid" json:"active_output_id"`
|
||||
Status string `json:"status"`
|
||||
SourceParseTaskID *uuid.UUID `gorm:"type:uuid" json:"source_parse_task_id"`
|
||||
UserEdited bool `json:"user_edited"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
UpdatedAt time.Time `json:"updated_at"`
|
||||
DeletedAt *time.Time `json:"-"`
|
||||
}
|
||||
|
||||
type EpisodeSource struct {
|
||||
ID uuid.UUID `gorm:"type:uuid;default:gen_random_uuid();primaryKey" json:"id"`
|
||||
EpisodeID uuid.UUID `gorm:"type:uuid;uniqueIndex" json:"episode_id"`
|
||||
Title string `json:"title"`
|
||||
RawContent string `json:"raw_content"`
|
||||
CharCount int `json:"char_count"`
|
||||
SourceType string `json:"source_type"`
|
||||
SourceFilename string `json:"source_filename"`
|
||||
ContentSHA256 string `gorm:"column:content_sha256" json:"content_sha256"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
UpdatedAt time.Time `json:"updated_at"`
|
||||
}
|
||||
|
||||
type DramaImportSession struct {
|
||||
ID uuid.UUID `gorm:"type:uuid;default:gen_random_uuid();primaryKey" json:"id"`
|
||||
UserID uuid.UUID `gorm:"type:uuid" json:"user_id"`
|
||||
ProjectID uuid.UUID `gorm:"type:uuid" json:"project_id"`
|
||||
SourceType string `json:"source_type"`
|
||||
SourceFilename string `json:"source_filename"`
|
||||
RawContent string `json:"-"`
|
||||
ContentSHA256 string `gorm:"column:content_sha256" json:"content_sha256"`
|
||||
PreviewData json.RawMessage `gorm:"type:jsonb" json:"preview_data"`
|
||||
ExpiresAt time.Time `json:"expires_at"`
|
||||
ConsumedAt *time.Time `json:"consumed_at"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
}
|
||||
|
||||
type DramaParseBatch struct {
|
||||
ID uuid.UUID `gorm:"type:uuid;default:gen_random_uuid();primaryKey" json:"id"`
|
||||
ProjectID uuid.UUID `gorm:"type:uuid" json:"project_id"`
|
||||
UserID uuid.UUID `gorm:"type:uuid" json:"user_id"`
|
||||
Mode string `json:"mode"`
|
||||
Status string `json:"status"`
|
||||
TotalCount int `json:"total_count"`
|
||||
CompletedCount int `json:"completed_count"`
|
||||
FailedCount int `json:"failed_count"`
|
||||
CurrentEpisodeID *uuid.UUID `gorm:"type:uuid" json:"current_episode_id"`
|
||||
CancelRequestedAt *time.Time `json:"cancel_requested_at"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
FinishedAt *time.Time `json:"finished_at"`
|
||||
}
|
||||
|
||||
type DramaParseTask struct {
|
||||
ID uuid.UUID `gorm:"type:uuid;default:gen_random_uuid();primaryKey" json:"id"`
|
||||
RequestID string `json:"request_id"`
|
||||
BatchID *uuid.UUID `gorm:"type:uuid" json:"batch_id"`
|
||||
UserID uuid.UUID `gorm:"type:uuid" json:"user_id"`
|
||||
ProjectID uuid.UUID `gorm:"type:uuid" json:"project_id"`
|
||||
EpisodeID uuid.UUID `gorm:"type:uuid" json:"episode_id"`
|
||||
ChannelID uuid.UUID `gorm:"type:uuid" json:"channel_id"`
|
||||
ModelID uuid.UUID `gorm:"type:uuid" json:"model_id"`
|
||||
SourceSHA256 string `gorm:"column:source_sha256" json:"source_sha256"`
|
||||
Status string `json:"status"`
|
||||
AttemptCount int `json:"attempt_count"`
|
||||
ModelSnapshot json.RawMessage `gorm:"type:jsonb" json:"model_snapshot"`
|
||||
PromptSnapshot string `json:"prompt_snapshot"`
|
||||
ContextSnapshot json.RawMessage `gorm:"type:jsonb" json:"context_snapshot"`
|
||||
NormalizedResult json.RawMessage `gorm:"type:jsonb" json:"normalized_result"`
|
||||
InputTokens *int64 `json:"input_tokens"`
|
||||
OutputTokens *int64 `json:"output_tokens"`
|
||||
TotalTokens *int64 `json:"total_tokens"`
|
||||
TokenCountSource string `json:"token_count_source"`
|
||||
UsageRaw json.RawMessage `gorm:"type:jsonb" json:"usage_raw"`
|
||||
BillingSnapshot json.RawMessage `gorm:"type:jsonb;not null;default:'{}'" json:"billing_snapshot"`
|
||||
EstimatedPoints string `gorm:"column:estimated_points" json:"estimated_points"`
|
||||
PrepaidPoints string `gorm:"column:prepaid_points" json:"prepaid_points"`
|
||||
ActualPoints *string `gorm:"column:actual_points" json:"actual_points"`
|
||||
CostRefunded bool `json:"cost_refunded"`
|
||||
RequestCharCount int64 `json:"request_char_count"`
|
||||
ResponseCharCount int64 `json:"response_char_count"`
|
||||
LeaseOwner string `json:"lease_owner"`
|
||||
LeaseUntil *time.Time `json:"lease_until"`
|
||||
CancelRequestedAt *time.Time `json:"cancel_requested_at"`
|
||||
ErrorCode string `json:"error_code"`
|
||||
ErrorMessage string `json:"error_message"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
StartedAt *time.Time `json:"started_at"`
|
||||
FinishedAt *time.Time `json:"finished_at"`
|
||||
}
|
||||
|
||||
type GenerationTask struct {
|
||||
ID uuid.UUID `gorm:"type:uuid;default:gen_random_uuid();primaryKey" json:"id"`
|
||||
RequestID string `json:"request_id"`
|
||||
UserID uuid.UUID `gorm:"type:uuid" json:"user_id"`
|
||||
ChannelID *uuid.UUID `gorm:"type:uuid" json:"channel_id"`
|
||||
ModelID *uuid.UUID `gorm:"type:uuid" json:"model_id"`
|
||||
ProjectID *uuid.UUID `gorm:"type:uuid" json:"project_id"`
|
||||
EpisodeID *uuid.UUID `gorm:"type:uuid" json:"episode_id"`
|
||||
StoryboardID *uuid.UUID `gorm:"type:uuid" json:"storyboard_id"`
|
||||
ScriptAnalysisID *uuid.UUID `gorm:"type:uuid" json:"script_analysis_id"`
|
||||
TaskType string `json:"task_type"`
|
||||
Status string `json:"status"`
|
||||
UpstreamTaskID string `json:"upstream_task_id"`
|
||||
InputData json.RawMessage `gorm:"type:jsonb" json:"input_data"`
|
||||
BillingSnapshot json.RawMessage `gorm:"type:jsonb;not null;default:'{}'" json:"billing_snapshot"`
|
||||
InputTokens *int64 `json:"input_tokens"`
|
||||
OutputTokens *int64 `json:"output_tokens"`
|
||||
TotalTokens *int64 `json:"total_tokens"`
|
||||
TokenCountSource *string `json:"token_count_source"`
|
||||
UsageRaw json.RawMessage `gorm:"type:jsonb" json:"usage_raw"`
|
||||
ResultData json.RawMessage `gorm:"type:jsonb;not null;default:'{}'" json:"result_data"`
|
||||
PromptSnapshot string `json:"prompt_snapshot"`
|
||||
EstimatedPoints string `gorm:"column:estimated_points" json:"estimated_points"`
|
||||
ActualPoints *string `gorm:"column:actual_points" json:"actual_points"`
|
||||
PrepaidPoints string `gorm:"column:prepaid_points" json:"prepaid_points"`
|
||||
CostRefunded bool `json:"cost_refunded"`
|
||||
SlotReserved bool `json:"slot_reserved"`
|
||||
SlotReservedAt *time.Time `json:"slot_reserved_at"`
|
||||
SlotReleasedAt *time.Time `json:"slot_released_at"`
|
||||
SubmitAttempts int `json:"submit_attempts"`
|
||||
PollAttempts int `json:"poll_attempts"`
|
||||
DownloadAttempts int `json:"download_attempts"`
|
||||
NextSubmitAt *time.Time `json:"next_submit_at"`
|
||||
NextPollAt *time.Time `json:"next_poll_at"`
|
||||
LeaseOwner string `json:"lease_owner"`
|
||||
LeaseUntil *time.Time `json:"lease_until"`
|
||||
UpstreamResultURL string `json:"-"`
|
||||
ErrorCode string `json:"error_code"`
|
||||
ErrorMessage string `json:"error_message"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
SubmittedAt *time.Time `json:"submitted_at"`
|
||||
FinishedAt *time.Time `json:"finished_at"`
|
||||
}
|
||||
|
||||
type GenerationOutput struct {
|
||||
ID uuid.UUID `gorm:"type:uuid;default:gen_random_uuid();primaryKey" json:"id"`
|
||||
TaskID uuid.UUID `gorm:"type:uuid" json:"task_id"`
|
||||
MediaAssetID uuid.UUID `gorm:"type:uuid" json:"media_asset_id"`
|
||||
OutputType string `json:"output_type"`
|
||||
SequenceNo int `json:"sequence_no"`
|
||||
Metadata json.RawMessage `gorm:"type:jsonb" json:"metadata"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
}
|
||||
@@ -0,0 +1,19 @@
|
||||
package model
|
||||
|
||||
import "testing"
|
||||
|
||||
func TestProjectAssetBeforeCreateInitializesJSONFields(t *testing.T) {
|
||||
asset := &ProjectAsset{}
|
||||
if err := asset.BeforeCreate(nil); err != nil {
|
||||
t.Fatalf("BeforeCreate returned error: %v", err)
|
||||
}
|
||||
if string(asset.Appearances) != "[]" {
|
||||
t.Fatalf("unexpected appearances default: %s", asset.Appearances)
|
||||
}
|
||||
if string(asset.Aliases) != "[]" {
|
||||
t.Fatalf("unexpected aliases default: %s", asset.Aliases)
|
||||
}
|
||||
if string(asset.Attributes) != "{}" {
|
||||
t.Fatalf("unexpected attributes default: %s", asset.Attributes)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,36 @@
|
||||
package model
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
type ScriptAnalysis struct {
|
||||
ID uuid.UUID `gorm:"type:uuid;default:gen_random_uuid();primaryKey" json:"id"`
|
||||
UserID uuid.UUID `gorm:"type:uuid;index" json:"user_id"`
|
||||
Name string `json:"name"`
|
||||
SourceContent string `json:"source_content"`
|
||||
ResultContent string `json:"result_content"`
|
||||
AnalysisResult json.RawMessage `gorm:"type:jsonb;not null;default:'{}'" json:"analysis_result"`
|
||||
AnalysisStatus string `json:"analysis_status"`
|
||||
AnalysisMessage string `json:"analysis_message"`
|
||||
Characters []ScriptAnalysisCharacter `gorm:"foreignKey:ScriptAnalysisID" json:"characters,omitempty"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
UpdatedAt time.Time `json:"updated_at"`
|
||||
DeletedAt *time.Time `json:"-"`
|
||||
}
|
||||
|
||||
type ScriptAnalysisCharacter struct {
|
||||
ID uuid.UUID `gorm:"type:uuid;default:gen_random_uuid();primaryKey" json:"id"`
|
||||
ScriptAnalysisID uuid.UUID `gorm:"type:uuid;index" json:"script_analysis_id"`
|
||||
Name string `json:"name"`
|
||||
Faction string `json:"faction"`
|
||||
Biography string `json:"biography"`
|
||||
Motivation string `json:"motivation"`
|
||||
Relationship string `json:"relationship"`
|
||||
SortOrder int `json:"sort_order"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
UpdatedAt time.Time `json:"updated_at"`
|
||||
}
|
||||
@@ -0,0 +1,39 @@
|
||||
package model
|
||||
|
||||
import (
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
type WebUser struct {
|
||||
ID uuid.UUID `gorm:"type:uuid;default:gen_random_uuid();primaryKey" json:"id"`
|
||||
UID string `json:"uid"`
|
||||
Username string `json:"username"`
|
||||
Account string `json:"account"`
|
||||
AvatarURL string `json:"avatar_url"`
|
||||
AvatarKey string `json:"-"`
|
||||
PasswordHash string `json:"-"`
|
||||
PointBalance string `gorm:"column:point_balance" json:"point_balance"`
|
||||
DailyLimit string `json:"daily_limit"`
|
||||
Enabled bool `json:"enabled"`
|
||||
SessionVersion int64 `json:"-"`
|
||||
LastOnlineAt *time.Time `json:"last_online_at"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
UpdatedAt time.Time `json:"updated_at"`
|
||||
DeletedAt gorm.DeletedAt `json:"-"`
|
||||
}
|
||||
|
||||
type WebRefreshToken struct {
|
||||
ID uuid.UUID `gorm:"type:uuid;default:gen_random_uuid();primaryKey"`
|
||||
UserID uuid.UUID `gorm:"type:uuid"`
|
||||
TokenHash string
|
||||
SessionVersion int64
|
||||
ExpiresAt time.Time
|
||||
RevokedAt *time.Time
|
||||
CreatedAt time.Time
|
||||
}
|
||||
|
||||
func (WebUser) TableName() string { return "web_users" }
|
||||
func (WebRefreshToken) TableName() string { return "web_refresh_tokens" }
|
||||
@@ -0,0 +1,306 @@
|
||||
// 管理后台业务模块,封装模型配置、兑换码、用户积分和操作审计的持久化事务。
|
||||
package admin
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"math"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"juhe-factory/api/internal/billing"
|
||||
"juhe-factory/api/internal/model"
|
||||
legacyservice "juhe-factory/api/internal/service"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
// Service 提供管理后台模块的公开业务接口,并隐藏底层数据库连接。
|
||||
type Service struct {
|
||||
db *gorm.DB
|
||||
}
|
||||
|
||||
// ModelPriceInput 表示模型的单项计费配置。
|
||||
type ModelPriceInput struct {
|
||||
PriceKey string `json:"price_key"`
|
||||
Unit string `json:"unit"`
|
||||
Price any `json:"price"`
|
||||
}
|
||||
|
||||
// ModelInput 表示管理端新增或修改模型时提交的完整配置。
|
||||
type ModelInput struct {
|
||||
ChannelID string `json:"channel_id"`
|
||||
Name string `json:"name"`
|
||||
ModelType string `json:"model_type"`
|
||||
Multimodal *bool `json:"multimodal"`
|
||||
TextBillingMode string `json:"text_billing_mode"`
|
||||
Enabled bool `json:"enabled"`
|
||||
Prices []ModelPriceInput `json:"prices"`
|
||||
}
|
||||
|
||||
// NewService 创建管理后台业务服务。
|
||||
func NewService(db *gorm.DB) *Service {
|
||||
return &Service{db: db}
|
||||
}
|
||||
|
||||
// ListModels 按管理端筛选条件分页查询模型及价格配置。
|
||||
func (s *Service) ListModels(keyword, modelType, channelID string, page, size int) (legacyservice.Page, error) {
|
||||
query := s.db.Table("models m").Joins("JOIN channels c ON c.id=m.channel_id").Where("m.deleted_at IS NULL")
|
||||
if keyword = strings.TrimSpace(keyword); keyword != "" {
|
||||
query = query.Where("m.name::text ILIKE ?", "%"+keyword+"%")
|
||||
}
|
||||
if modelType = strings.TrimSpace(modelType); modelType != "" {
|
||||
query = query.Where("m.model_type=?", modelType)
|
||||
}
|
||||
if channelID = strings.TrimSpace(channelID); channelID != "" {
|
||||
query = query.Where("m.channel_id=?", channelID)
|
||||
}
|
||||
var total int64
|
||||
if err := query.Count(&total).Error; err != nil {
|
||||
return legacyservice.Page{}, err
|
||||
}
|
||||
items := make([]map[string]any, 0)
|
||||
err := query.Select(`m.id,m.channel_id,c.name AS channel_name,m.name,m.model_type,m.multimodal,m.text_billing_mode,m.enabled,c.enabled AS channel_enabled,
|
||||
coalesce((SELECT jsonb_agg(jsonb_build_object('price_key',p.price_key,'unit',p.unit,'price',p.price) ORDER BY CASE p.price_key WHEN 'default' THEN 0 WHEN '1K' THEN 1 WHEN '2K' THEN 2 WHEN '4K' THEN 3 WHEN '480p' THEN 4 WHEN '720p' THEN 5 WHEN '1080p' THEN 6 ELSE 99 END) FROM model_prices p WHERE p.model_id=m.id),'[]'::jsonb) AS prices`).
|
||||
Order("m.created_at DESC").Offset((page - 1) * size).Limit(size).Find(&items).Error
|
||||
if err != nil {
|
||||
return legacyservice.Page{}, err
|
||||
}
|
||||
for i := range items {
|
||||
items[i]["prices"] = decodePriceList(items[i]["prices"])
|
||||
}
|
||||
return legacyservice.Page{Items: items, Total: total, Page: page, PageSize: size}, nil
|
||||
}
|
||||
|
||||
// SaveModel 校验模型计费配置,并在同一事务中保存模型及全部价格项。
|
||||
func (s *Service) SaveModel(modelID string, input ModelInput) (string, bool, ModelInput, error) {
|
||||
if err := normalizeModelPrices(&input); err != nil {
|
||||
return "", false, input, err
|
||||
}
|
||||
created := modelID == ""
|
||||
if created {
|
||||
modelID = uuid.NewString()
|
||||
}
|
||||
err := s.db.Transaction(func(tx *gorm.DB) error {
|
||||
values := map[string]any{"channel_id": input.ChannelID, "name": input.Name, "model_type": input.ModelType, "enabled": input.Enabled, "text_billing_mode": nil}
|
||||
switch input.ModelType {
|
||||
case "text":
|
||||
values["multimodal"] = input.Multimodal
|
||||
values["text_billing_mode"] = input.TextBillingMode
|
||||
default:
|
||||
values["multimodal"] = nil
|
||||
}
|
||||
if created {
|
||||
values["id"] = modelID
|
||||
if err := tx.Table("models").Create(values).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
} else if err := tx.Table("models").Where("id=? AND deleted_at IS NULL", modelID).Updates(values).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if err := tx.Exec("DELETE FROM model_prices WHERE model_id=?", modelID).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
for _, price := range input.Prices {
|
||||
if err := tx.Exec("INSERT INTO model_prices(id,model_id,price_key,unit,price) VALUES(?,?,?,?,?)", uuid.NewString(), modelID, price.PriceKey, price.Unit, price.Price).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
})
|
||||
return modelID, created, input, err
|
||||
}
|
||||
|
||||
// ListRedemptions 按关键字和状态分页查询兑换码记录。
|
||||
func (s *Service) ListRedemptions(keyword, status string, page, size int) (legacyservice.Page, error) {
|
||||
query := s.db.Table("redemption_codes r").Joins("JOIN redemption_batches b ON b.id=r.batch_id").Joins("LEFT JOIN web_users u ON u.id=r.redeemed_by")
|
||||
if keyword = strings.TrimSpace(keyword); keyword != "" {
|
||||
query = query.Where("r.code ILIKE ? OR r.code_mask ILIKE ? OR b.name::text ILIKE ? OR u.uid=?", "%"+keyword+"%", "%"+keyword+"%", "%"+keyword+"%", keyword)
|
||||
}
|
||||
if status = strings.TrimSpace(status); status != "" {
|
||||
query = query.Where("CASE WHEN r.status='unused' AND r.expires_at<CURRENT_TIMESTAMP THEN 'expired' ELSE r.status END=?", status)
|
||||
}
|
||||
var total int64
|
||||
if err := query.Count(&total).Error; err != nil {
|
||||
return legacyservice.Page{}, err
|
||||
}
|
||||
items := make([]map[string]any, 0)
|
||||
err := query.Select("r.id,r.code,b.name AS batch_name,r.points,CASE WHEN r.status='unused' AND r.expires_at<CURRENT_TIMESTAMP THEN 'expired' ELSE r.status END AS status,u.uid AS redeemed_uid,r.redeemed_at,r.expires_at,a.username AS created_by,r.created_at").
|
||||
Joins("JOIN admin_users a ON a.id=b.created_by").Order("r.created_at DESC").Offset((page - 1) * size).Limit(size).Find(&items).Error
|
||||
if err != nil {
|
||||
return legacyservice.Page{}, err
|
||||
}
|
||||
return legacyservice.Page{Items: items, Total: total, Page: page, PageSize: size}, nil
|
||||
}
|
||||
|
||||
// ListAuditLogs 按关键字分页查询管理端操作审计记录,仅提供只读访问。
|
||||
func (s *Service) ListAuditLogs(keyword string, page, size int) (legacyservice.Page, error) {
|
||||
query := s.db.Table("admin_audit_logs")
|
||||
if keyword = strings.TrimSpace(keyword); keyword != "" {
|
||||
pattern := "%" + keyword + "%"
|
||||
query = query.Where(`admin_username ILIKE ? OR action::text ILIKE ? OR resource_type::text ILIKE ?
|
||||
OR coalesce(resource_id,'') ILIKE ? OR coalesce(reason,'') ILIKE ? OR coalesce(trace_id,'') ILIKE ?`,
|
||||
pattern, pattern, pattern, pattern, pattern, pattern)
|
||||
if resourceTypes := matchingAuditResourceTypes(keyword); len(resourceTypes) > 0 {
|
||||
query = query.Or("resource_type IN ?", resourceTypes)
|
||||
}
|
||||
}
|
||||
var total int64
|
||||
if err := query.Count(&total).Error; err != nil {
|
||||
return legacyservice.Page{}, err
|
||||
}
|
||||
items := make([]map[string]any, 0)
|
||||
err := query.Select(`id,admin_id,admin_username,action,resource_type,coalesce(resource_id,'') AS resource_id,
|
||||
coalesce(reason,'') AS reason,coalesce(trace_id,'') AS trace_id,created_at`).
|
||||
Order("created_at DESC,id DESC").Offset((page - 1) * size).Limit(size).Find(&items).Error
|
||||
return legacyservice.Page{Items: items, Total: total, Page: page, PageSize: size}, err
|
||||
}
|
||||
|
||||
// matchingAuditResourceTypes 将中文资源名称转换为数据库内部代码,支持管理端直接使用中文搜索。
|
||||
func matchingAuditResourceTypes(keyword string) []string {
|
||||
keyword = strings.TrimSpace(keyword)
|
||||
if keyword == "" {
|
||||
return nil
|
||||
}
|
||||
labels := map[string]string{
|
||||
"users": "用户",
|
||||
"redemption-codes": "兑换码",
|
||||
"styles": "风格",
|
||||
"prompts": "提示词",
|
||||
"channels": "渠道",
|
||||
"models": "模型",
|
||||
"admin-auth": "管理员账号",
|
||||
}
|
||||
result := make([]string, 0, len(labels))
|
||||
for resourceType, label := range labels {
|
||||
if strings.Contains(label, keyword) {
|
||||
result = append(result, resourceType)
|
||||
}
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
// CreateRedemptionBatch 校验生成限制,并在事务中创建兑换批次和明文返回值。
|
||||
func (s *Service) CreateRedemptionBatch(adminID uuid.UUID, name string, pointsValue any, quantity int, expiresAt time.Time) (string, []string, string, error) {
|
||||
points, err := ValidateRedemptionBatchLimits(pointsValue, quantity)
|
||||
if err != nil {
|
||||
return "", nil, "", err
|
||||
}
|
||||
batchID := uuid.NewString()
|
||||
codes := make([]string, 0, quantity)
|
||||
err = s.db.Transaction(func(tx *gorm.DB) error {
|
||||
if err := tx.Exec("INSERT INTO redemption_batches(id,name,points,quantity,expires_at,created_by) VALUES(?,?,?,?,?,?)", batchID, name, points, quantity, expiresAt, adminID).Error; err != nil {
|
||||
// 命中 redemption_batches_name_key 唯一约束时返回业务友好提示,避免把底层 SQL 错误透传给前端
|
||||
if strings.Contains(err.Error(), "23505") {
|
||||
return errors.New("批次名称已存在,请更换后重试")
|
||||
}
|
||||
return err
|
||||
}
|
||||
for i := 0; i < quantity; i++ {
|
||||
plain, hash, mask, err := legacyservice.GenerateRedemptionCode()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err := tx.Exec("INSERT INTO redemption_codes(id,batch_id,code_hash,code_mask,code,points,expires_at) VALUES(?,?,?,?,?,?,?)", uuid.NewString(), batchID, hash, mask, plain, points, expiresAt).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
codes = append(codes, plain)
|
||||
}
|
||||
return nil
|
||||
})
|
||||
return batchID, codes, points, err
|
||||
}
|
||||
|
||||
// GrantUserPoints 通过永久积分批次为指定用户增加积分,并使用请求号保证重复提交不会重复入账。
|
||||
func (s *Service) GrantUserPoints(userID uuid.UUID, requestID string, pointsValue any) (string, bool, error) {
|
||||
requestID = strings.TrimSpace(requestID)
|
||||
if _, err := uuid.Parse(requestID); err != nil {
|
||||
return "", false, errors.New("积分发放请求号无效")
|
||||
}
|
||||
points := strings.TrimSpace(fmt.Sprint(pointsValue))
|
||||
sourceID := userID.String() + ":" + requestID
|
||||
var balance string
|
||||
var credited bool
|
||||
err := s.db.Transaction(func(tx *gorm.DB) error {
|
||||
var err error
|
||||
balance, credited, err = billing.CreditPoints(tx, userID, "admin_grant", sourceID, points, "后台发放积分")
|
||||
return err
|
||||
})
|
||||
return balance, credited, err
|
||||
}
|
||||
|
||||
// WriteAudit 记录管理端操作审计信息,调用方无需接触数据库连接。
|
||||
func (s *Service) WriteAudit(adminUser *model.AdminUser, action, resource, resourceID, reason, ip, traceID string, detail any) error {
|
||||
return legacyservice.WriteAudit(s.db, adminUser, action, resource, resourceID, reason, ip, traceID, detail)
|
||||
}
|
||||
|
||||
// normalizeModelPrices 校验并标准化文本模型的计费键、单位和金额。
|
||||
func normalizeModelPrices(input *ModelInput) error {
|
||||
if input.ModelType != "text" {
|
||||
input.TextBillingMode = ""
|
||||
return nil
|
||||
}
|
||||
expected := map[string]string{}
|
||||
switch input.TextBillingMode {
|
||||
case "per_request":
|
||||
expected["default"] = "次"
|
||||
case "per_token":
|
||||
expected["input"] = "M Token"
|
||||
expected["output"] = "M Token"
|
||||
default:
|
||||
return errors.New("文本模型计费模式无效")
|
||||
}
|
||||
if len(input.Prices) != len(expected) {
|
||||
return errors.New("文本模型价格配置不完整")
|
||||
}
|
||||
seen := map[string]bool{}
|
||||
for index := range input.Prices {
|
||||
key := strings.ToLower(strings.TrimSpace(input.Prices[index].PriceKey))
|
||||
unit, ok := expected[key]
|
||||
value, err := strconv.ParseFloat(fmt.Sprint(input.Prices[index].Price), 64)
|
||||
if !ok || seen[key] || err != nil || value < 0 || math.IsNaN(value) || math.IsInf(value, 0) {
|
||||
return errors.New("文本模型价格配置无效")
|
||||
}
|
||||
seen[key] = true
|
||||
input.Prices[index].PriceKey = key
|
||||
input.Prices[index].Unit = unit
|
||||
input.Prices[index].Price = math.Round(value*100) / 100
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// decodePriceList 将 PostgreSQL JSON 聚合结果转换为稳定的价格数组。
|
||||
func decodePriceList(value any) []map[string]any {
|
||||
var raw []byte
|
||||
switch data := value.(type) {
|
||||
case []byte:
|
||||
raw = data
|
||||
case json.RawMessage:
|
||||
raw = data
|
||||
case string:
|
||||
raw = []byte(data)
|
||||
default:
|
||||
return []map[string]any{}
|
||||
}
|
||||
prices := make([]map[string]any, 0)
|
||||
if err := json.Unmarshal(raw, &prices); err != nil {
|
||||
return []map[string]any{}
|
||||
}
|
||||
return prices
|
||||
}
|
||||
|
||||
// ValidateRedemptionBatchLimits 校验单码积分和单批生成数量的安全范围。
|
||||
func ValidateRedemptionBatchLimits(value any, quantity int) (string, error) {
|
||||
points := strings.TrimSpace(fmt.Sprint(value))
|
||||
parsed, err := strconv.ParseFloat(points, 64)
|
||||
if err != nil || math.IsNaN(parsed) || math.IsInf(parsed, 0) || parsed < 1 || parsed > 100 {
|
||||
return "", errors.New("单码积分数量必须为 1 至 100")
|
||||
}
|
||||
if quantity < 1 || quantity > 10 {
|
||||
return "", errors.New("生成数量必须为 1 至 10")
|
||||
}
|
||||
return points, nil
|
||||
}
|
||||
@@ -0,0 +1,67 @@
|
||||
// 管理后台模块单元测试,验证模型计费、兑换码批次和审计搜索的纯业务规则。
|
||||
package admin
|
||||
|
||||
import (
|
||||
"slices"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// TestNormalizeModelPrices 验证按 Token 计费时价格键、单位和金额标准化。
|
||||
func TestNormalizeModelPrices(t *testing.T) {
|
||||
input := ModelInput{
|
||||
ModelType: "text",
|
||||
TextBillingMode: "per_token",
|
||||
Prices: []ModelPriceInput{
|
||||
{PriceKey: " INPUT ", Price: "1.236"},
|
||||
{PriceKey: "output", Price: 2},
|
||||
},
|
||||
}
|
||||
if err := normalizeModelPrices(&input); err != nil {
|
||||
t.Fatalf("normalize model prices: %v", err)
|
||||
}
|
||||
if input.Prices[0].PriceKey != "input" || input.Prices[0].Unit != "M Token" || input.Prices[0].Price != 1.24 {
|
||||
t.Fatalf("unexpected normalized input price: %#v", input.Prices[0])
|
||||
}
|
||||
}
|
||||
|
||||
// TestNormalizeModelPricesRejectsIncompleteConfig 验证文本模型缺少价格项时拒绝保存。
|
||||
func TestNormalizeModelPricesRejectsIncompleteConfig(t *testing.T) {
|
||||
input := ModelInput{ModelType: "text", TextBillingMode: "per_token", Prices: []ModelPriceInput{{PriceKey: "input", Price: 1}}}
|
||||
if err := normalizeModelPrices(&input); err == nil {
|
||||
t.Fatal("expected incomplete token prices to be rejected")
|
||||
}
|
||||
}
|
||||
|
||||
// TestValidateRedemptionBatchLimits 验证兑换码单码积分与批次数量边界。
|
||||
func TestValidateRedemptionBatchLimits(t *testing.T) {
|
||||
if points, err := ValidateRedemptionBatchLimits("10.5", 10); err != nil || points != "10.5" {
|
||||
t.Fatalf("expected valid redemption limits, points=%q err=%v", points, err)
|
||||
}
|
||||
for _, test := range []struct {
|
||||
points any
|
||||
quantity int
|
||||
}{{0, 1}, {101, 1}, {10, 0}, {10, 11}} {
|
||||
if _, err := ValidateRedemptionBatchLimits(test.points, test.quantity); err == nil {
|
||||
t.Fatalf("expected limits to be rejected: %#v", test)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestMatchingAuditResourceTypes 验证中文资源名称可转换为审计记录使用的内部代码。
|
||||
func TestMatchingAuditResourceTypes(t *testing.T) {
|
||||
for _, test := range []struct {
|
||||
keyword string
|
||||
want string
|
||||
}{
|
||||
{keyword: "渠道", want: "channels"},
|
||||
{keyword: "兑换", want: "redemption-codes"},
|
||||
{keyword: "管理员账号", want: "admin-auth"},
|
||||
} {
|
||||
if got := matchingAuditResourceTypes(test.keyword); !slices.Contains(got, test.want) {
|
||||
t.Fatalf("keyword %q matched %v, want %q", test.keyword, got, test.want)
|
||||
}
|
||||
}
|
||||
if got := matchingAuditResourceTypes("不存在的资源"); len(got) != 0 {
|
||||
t.Fatalf("unexpected resource matches: %v", got)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,323 @@
|
||||
// 图片生成生图业务模块,负责提交用户级图片任务、查询历史并清理生成资源。
|
||||
package productimage
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
"unicode/utf8"
|
||||
|
||||
"juhe-factory/api/internal/billing"
|
||||
"juhe-factory/api/internal/model"
|
||||
queuepkg "juhe-factory/api/internal/queue"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/hibiken/asynq"
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/clause"
|
||||
)
|
||||
|
||||
var activeTaskStatuses = []string{"pending_submission", "submitting", "submitted", "processing", "result_ready", "downloading", "cancel_requested"}
|
||||
|
||||
// ReferenceUpload 表示已经上传到对象存储、等待绑定到生图任务的参考图。
|
||||
type ReferenceUpload struct {
|
||||
Asset *model.MediaAsset
|
||||
Name string
|
||||
}
|
||||
|
||||
// GenerateInput 是图片生成生图接口的配置参数和临时参考图。
|
||||
type GenerateInput struct {
|
||||
Prompt string
|
||||
ModelID uuid.UUID
|
||||
AspectRatio string
|
||||
Resolution string
|
||||
References []ReferenceUpload
|
||||
}
|
||||
|
||||
// GenerationView 汇总任务状态和已经落库的图片结果。
|
||||
type GenerationView struct {
|
||||
TaskID uuid.UUID `json:"task_id"`
|
||||
Status string `json:"status"`
|
||||
Prompt string `json:"prompt"`
|
||||
ModelName string `json:"model_name"`
|
||||
ErrorMessage string `json:"error_message"`
|
||||
MediaAssetID *uuid.UUID `json:"media_asset_id"`
|
||||
PublicURL string `json:"public_url"`
|
||||
MimeType string `json:"mime_type"`
|
||||
EstimatedPoints string `json:"estimated_points"`
|
||||
ActualPoints *string `json:"actual_points"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
FinishedAt *time.Time `json:"finished_at"`
|
||||
}
|
||||
|
||||
// Service 组合图片生成模块需要的数据库和任务队列依赖。
|
||||
type Service struct {
|
||||
DB *gorm.DB
|
||||
Queue *asynq.Client
|
||||
}
|
||||
|
||||
// NewService 创建图片生成生图服务。
|
||||
func NewService(db *gorm.DB, queue *asynq.Client) *Service { return &Service{DB: db, Queue: queue} }
|
||||
|
||||
// QueueGeneration 校验配置、保存临时参考图并提交一个独立图片任务。
|
||||
func (s *Service) QueueGeneration(userID uuid.UUID, input GenerateInput) (*model.GenerationTask, error) {
|
||||
if s.Queue == nil {
|
||||
return nil, errors.New("生成任务队列不可用")
|
||||
}
|
||||
prompt := strings.TrimSpace(input.Prompt)
|
||||
if prompt == "" {
|
||||
return nil, errors.New("提示词不能为空")
|
||||
}
|
||||
if utf8.RuneCountInString(prompt) > 5000 {
|
||||
return nil, errors.New("提示词不能超过 5000 个字符")
|
||||
}
|
||||
if len(input.References) > 4 {
|
||||
return nil, errors.New("最多上传 4 张参考图")
|
||||
}
|
||||
if input.ModelID == uuid.Nil {
|
||||
return nil, errors.New("请选择图片模型")
|
||||
}
|
||||
if !validRatio(input.AspectRatio) || !validResolution(input.Resolution) {
|
||||
return nil, errors.New("图片生成参数无效")
|
||||
}
|
||||
var selected struct {
|
||||
ModelID uuid.UUID
|
||||
ChannelID uuid.UUID
|
||||
ModelName string
|
||||
Price string
|
||||
PriceExists bool
|
||||
}
|
||||
if err := s.DB.Raw(`SELECT model.id AS model_id,model.channel_id,model.name AS model_name,
|
||||
coalesce(price.price,0)::text AS price,(price.id IS NOT NULL) AS price_exists
|
||||
FROM models model JOIN channels channel ON channel.id=model.channel_id
|
||||
LEFT JOIN model_prices price ON price.model_id=model.id AND lower(price.price_key)=lower(?)
|
||||
WHERE model.id=? AND model.model_type='image' AND model.enabled=true AND model.deleted_at IS NULL
|
||||
AND channel.enabled=true AND channel.deleted_at IS NULL`, input.Resolution, input.ModelID).Scan(&selected).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if selected.ModelID == uuid.Nil {
|
||||
return nil, errors.New("图片模型不可用")
|
||||
}
|
||||
if !selected.PriceExists {
|
||||
return nil, fmt.Errorf("当前图片模型未配置 %s 价格", input.Resolution)
|
||||
}
|
||||
imageURLs := make([]string, 0, len(input.References))
|
||||
referenceIDs := make([]string, 0, len(input.References))
|
||||
referenceNames := make([]string, 0, len(input.References))
|
||||
seenNames := make(map[string]struct{}, len(input.References))
|
||||
for _, reference := range input.References {
|
||||
if reference.Asset == nil || strings.TrimSpace(reference.Asset.PublicURL) == "" {
|
||||
return nil, errors.New("参考图信息无效")
|
||||
}
|
||||
name := strings.TrimSpace(reference.Name)
|
||||
if name == "" {
|
||||
return nil, errors.New("参考图名称不能为空")
|
||||
}
|
||||
if _, exists := seenNames[name]; exists {
|
||||
return nil, errors.New("参考图名称不能重复")
|
||||
}
|
||||
seenNames[name] = struct{}{}
|
||||
imageURLs = append(imageURLs, reference.Asset.PublicURL)
|
||||
referenceIDs = append(referenceIDs, reference.Asset.ID.String())
|
||||
referenceNames = append(referenceNames, name)
|
||||
}
|
||||
providerPrompt := buildProviderPrompt(prompt, referenceNames)
|
||||
payload, err := json.Marshal(map[string]any{
|
||||
"product_image": true, "prompt": providerPrompt, "display_prompt": prompt,
|
||||
"model": selected.ModelName, "size": input.AspectRatio, "resolution": input.Resolution,
|
||||
"image_urls": imageURLs, "reference_media_asset_ids": referenceIDs, "reference_names": referenceNames, "n": 1,
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
task := &model.GenerationTask{ID: uuid.New(), RequestID: "product_image_" + uuid.NewString(), UserID: userID, ChannelID: &selected.ChannelID, ModelID: &selected.ModelID, TaskType: "image_generation", Status: "pending_submission", InputData: payload}
|
||||
err = s.DB.Transaction(func(tx *gorm.DB) error {
|
||||
for _, reference := range input.References {
|
||||
owner := userID
|
||||
reference.Asset.OwnerUserID = &owner
|
||||
if err := tx.Create(reference.Asset).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return billing.PrechargeGenerationTask(tx, task, selected.Price, 1, "图片生成")
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := queuepkg.EnqueueID(s.Queue, queuepkg.TypeDispatchChannel, selected.ChannelID, 0); err != nil {
|
||||
if refundErr := s.failQueuedTask(task.ID, err.Error()); refundErr != nil {
|
||||
return nil, fmt.Errorf("生成任务入队失败且预扣返还失败: %w", refundErr)
|
||||
}
|
||||
return nil, errors.New("生成任务队列暂时不可用,请稍后重试")
|
||||
}
|
||||
return task, nil
|
||||
}
|
||||
|
||||
// Generations 返回当前用户最近五十条图片生成记录。
|
||||
func (s *Service) Generations(userID uuid.UUID) ([]GenerationView, error) {
|
||||
items := make([]GenerationView, 0)
|
||||
err := s.DB.Raw(`SELECT task.id AS task_id,task.status,
|
||||
coalesce(task.input_data->>'display_prompt',task.input_data->>'prompt','') AS prompt,
|
||||
coalesce(model.name,task.input_data->>'model','') AS model_name,task.error_message,
|
||||
output.media_asset_id,coalesce(media.public_url,'') AS public_url,coalesce(media.mime_type,'') AS mime_type,
|
||||
task.estimated_points,task.actual_points,task.created_at,task.finished_at
|
||||
FROM generation_tasks task
|
||||
LEFT JOIN models model ON model.id=task.model_id
|
||||
LEFT JOIN generation_outputs output ON output.task_id=task.id AND output.sequence_no=1
|
||||
LEFT JOIN media_assets media ON media.id=output.media_asset_id AND media.deleted_at IS NULL
|
||||
WHERE task.user_id=? AND task.task_type='image_generation' AND task.input_data->>'product_image'='true'
|
||||
ORDER BY task.created_at DESC LIMIT 50`, userID).Scan(&items).Error
|
||||
return items, err
|
||||
}
|
||||
|
||||
// DeleteGeneration 清理当前用户的图片任务、生成媒体和参考媒体;活动任务也允许主动删除。
|
||||
func (s *Service) DeleteGeneration(userID, taskID uuid.UUID, deleteObjects func([]string) error) error {
|
||||
objectKeys := make([]string, 0)
|
||||
var channelID *uuid.UUID
|
||||
err := s.DB.Transaction(func(tx *gorm.DB) error {
|
||||
var task model.GenerationTask
|
||||
if err := tx.Where("id=? AND user_id=? AND task_type='image_generation' AND input_data->>'product_image'='true'", taskID, userID).Take(&task).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if !canDeleteGeneration(task.Status, task.ErrorMessage) {
|
||||
return errors.New("图片仍在正常生成中,暂时不能删除")
|
||||
}
|
||||
channelID = task.ChannelID
|
||||
var input struct {
|
||||
ReferenceMediaAssetIDs []uuid.UUID `json:"reference_media_asset_ids"`
|
||||
}
|
||||
_ = json.Unmarshal(task.InputData, &input)
|
||||
mediaIDs := append([]uuid.UUID(nil), input.ReferenceMediaAssetIDs...)
|
||||
var outputs []struct {
|
||||
MediaAssetID uuid.UUID
|
||||
ObjectKey string
|
||||
}
|
||||
if err := tx.Table("generation_outputs output").Select("output.media_asset_id,media.object_key").Joins("JOIN media_assets media ON media.id=output.media_asset_id").Where("output.task_id=?", taskID).Scan(&outputs).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
for _, output := range outputs {
|
||||
mediaIDs = append(mediaIDs, output.MediaAssetID)
|
||||
if output.ObjectKey != "" {
|
||||
objectKeys = append(objectKeys, output.ObjectKey)
|
||||
}
|
||||
}
|
||||
if len(mediaIDs) > 0 {
|
||||
var refs []struct{ ObjectKey string }
|
||||
if err := tx.Table("media_assets").Select("object_key").Where("id IN ?", mediaIDs).Scan(&refs).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
for _, ref := range refs {
|
||||
if ref.ObjectKey != "" {
|
||||
objectKeys = append(objectKeys, ref.ObjectKey)
|
||||
}
|
||||
}
|
||||
}
|
||||
if deleteObjects != nil {
|
||||
if err := deleteObjects(uniqueStrings(objectKeys)); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
if err := tx.Exec("DELETE FROM generation_outputs WHERE task_id=?", taskID).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if len(mediaIDs) > 0 {
|
||||
if err := tx.Exec("DELETE FROM channel_asset_cache WHERE media_asset_id IN ?", mediaIDs).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if err := tx.Exec("DELETE FROM media_assets WHERE id IN ?", mediaIDs).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return tx.Exec("DELETE FROM generation_tasks WHERE id=?", taskID).Error
|
||||
})
|
||||
if err == nil && channelID != nil && s.Queue != nil {
|
||||
_ = queuepkg.EnqueueID(s.Queue, queuepkg.TypeDispatchChannel, *channelID, 0)
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
||||
// canDeleteGeneration 判断图片任务是否已结束,或因错误停留在活动状态而允许用户清理。
|
||||
func canDeleteGeneration(status, errorMessage string) bool {
|
||||
for _, activeStatus := range activeTaskStatuses {
|
||||
if status == activeStatus {
|
||||
return strings.TrimSpace(errorMessage) != ""
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
// failQueuedTask 在调度队列不可用时结束任务、返还积分并删除参考媒体记录。
|
||||
func (s *Service) failQueuedTask(taskID uuid.UUID, message string) error {
|
||||
return s.DB.Transaction(func(tx *gorm.DB) error {
|
||||
var task model.GenerationTask
|
||||
if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Where("id=?", taskID).Take(&task).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
refunded, err := billing.RefundGenerationTask(tx, &task, "图片生成入队失败返还")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
var input struct {
|
||||
ReferenceMediaAssetIDs []uuid.UUID `json:"reference_media_asset_ids"`
|
||||
}
|
||||
_ = json.Unmarshal(task.InputData, &input)
|
||||
if len(input.ReferenceMediaAssetIDs) > 0 {
|
||||
if err := tx.Exec("DELETE FROM media_assets WHERE id IN ?", input.ReferenceMediaAssetIDs).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
updates := map[string]any{"status": "failed", "error_code": "queue_unavailable", "error_message": message, "finished_at": time.Now(), "actual_points": "0.00"}
|
||||
if refunded {
|
||||
updates["cost_refunded"] = true
|
||||
}
|
||||
return tx.Model(&model.GenerationTask{}).Where("id=?", taskID).Updates(updates).Error
|
||||
})
|
||||
}
|
||||
|
||||
// validRatio 判断画幅比例是否属于图片生成工具公开支持的取值。
|
||||
func validRatio(value string) bool {
|
||||
for _, allowed := range []string{"9:21", "1:1", "3:4", "4:3", "9:16", "16:9"} {
|
||||
if value == allowed {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// validResolution 判断分辨率是否属于图片生成工具公开支持的档位。
|
||||
func validResolution(value string) bool {
|
||||
for _, allowed := range []string{"1k", "2k", "4k"} {
|
||||
if value == allowed {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// buildProviderPrompt 将文件名提及与参考图数组顺序建立明确对应关系。
|
||||
func buildProviderPrompt(prompt string, names []string) string {
|
||||
if len(names) == 0 {
|
||||
return prompt
|
||||
}
|
||||
mappings := make([]string, 0, len(names))
|
||||
for index, name := range names {
|
||||
mappings = append(mappings, fmt.Sprintf("第%d张参考图名称为“%s”", index+1, name))
|
||||
}
|
||||
return "参考图对应关系:" + strings.Join(mappings, ";") + "。\n" + prompt
|
||||
}
|
||||
|
||||
// uniqueStrings 按首次出现顺序去除空对象键和重复对象键。
|
||||
func uniqueStrings(values []string) []string {
|
||||
seen := make(map[string]struct{}, len(values))
|
||||
result := make([]string, 0, len(values))
|
||||
for _, value := range values {
|
||||
if _, ok := seen[value]; ok || value == "" {
|
||||
continue
|
||||
}
|
||||
seen[value] = struct{}{}
|
||||
result = append(result, value)
|
||||
}
|
||||
return result
|
||||
}
|
||||
@@ -0,0 +1,56 @@
|
||||
// 图片生成服务测试,验证参考图文件名映射和对象键去重规则。
|
||||
package productimage
|
||||
|
||||
import "testing"
|
||||
|
||||
// TestBuildProviderPrompt 验证文件名提及会与参考图数组顺序建立稳定对应关系。
|
||||
func TestBuildProviderPrompt(t *testing.T) {
|
||||
t.Parallel()
|
||||
prompt := "让@正面图.png中的商品使用@包装图.webp的包装"
|
||||
got := buildProviderPrompt(prompt, []string{"正面图.png", "包装图.webp"})
|
||||
want := "参考图对应关系:第1张参考图名称为“正面图.png”;第2张参考图名称为“包装图.webp”。\n" + prompt
|
||||
if got != want {
|
||||
t.Fatalf("buildProviderPrompt() = %q, want %q", got, want)
|
||||
}
|
||||
if withoutReferences := buildProviderPrompt(prompt, nil); withoutReferences != prompt {
|
||||
t.Fatalf("buildProviderPrompt() without references = %q, want %q", withoutReferences, prompt)
|
||||
}
|
||||
}
|
||||
|
||||
// TestUniqueStrings 验证重复或空对象键不会触发重复删除请求。
|
||||
func TestUniqueStrings(t *testing.T) {
|
||||
t.Parallel()
|
||||
got := uniqueStrings([]string{"a", "", "b", "a", "b", "c"})
|
||||
want := []string{"a", "b", "c"}
|
||||
if len(got) != len(want) {
|
||||
t.Fatalf("uniqueStrings() length = %d, want %d", len(got), len(want))
|
||||
}
|
||||
for index := range want {
|
||||
if got[index] != want[index] {
|
||||
t.Fatalf("uniqueStrings()[%d] = %q, want %q", index, got[index], want[index])
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestCanDeleteGeneration 验证失败停滞的活动任务可删除,而正常生成中的任务仍受保护。
|
||||
func TestCanDeleteGeneration(t *testing.T) {
|
||||
t.Parallel()
|
||||
tests := []struct {
|
||||
name string
|
||||
status string
|
||||
errorMessage string
|
||||
want bool
|
||||
}{
|
||||
{name: "正常生成", status: "processing", want: false},
|
||||
{name: "连接超时", status: "pending_submission", errorMessage: "dial tcp: i/o timeout", want: true},
|
||||
{name: "生成失败", status: "failed", want: true},
|
||||
{name: "生成成功", status: "succeeded", want: true},
|
||||
}
|
||||
for _, test := range tests {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
if got := canDeleteGeneration(test.status, test.errorMessage); got != test.want {
|
||||
t.Fatalf("canDeleteGeneration(%q, %q) = %t, want %t", test.status, test.errorMessage, got, test.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,328 @@
|
||||
// 提示词业务模块,统一封装用户提示词规则、选择关系和持久化操作。
|
||||
package prompt
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"strings"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
// Service 提供提示词模块的公开业务接口,并隐藏底层数据库实现。
|
||||
type Service struct {
|
||||
db *gorm.DB
|
||||
}
|
||||
|
||||
// CustomPromptInput 表示用户自定义提示词的可编辑内容。
|
||||
type CustomPromptInput struct {
|
||||
Name string
|
||||
Type string
|
||||
Content string
|
||||
}
|
||||
|
||||
// CustomPrompt 表示返回给接口层的用户自定义提示词。
|
||||
type CustomPrompt struct {
|
||||
ID uuid.UUID `json:"id"`
|
||||
Name string `json:"name"`
|
||||
Type string `json:"type"`
|
||||
Content string `json:"content"`
|
||||
Scope string `json:"scope"`
|
||||
Editable bool `json:"editable"`
|
||||
}
|
||||
|
||||
// SystemPromptInput 表示管理端可维护的系统提示词内容。
|
||||
type SystemPromptInput struct {
|
||||
Name string `json:"name"`
|
||||
Type string `json:"type"`
|
||||
Content string `json:"content"`
|
||||
}
|
||||
|
||||
// SystemPrompt 表示管理端保存后返回的系统提示词。
|
||||
type SystemPrompt struct {
|
||||
ID string `json:"id"`
|
||||
Name string `json:"name"`
|
||||
Type string `json:"type"`
|
||||
Content string `json:"content"`
|
||||
}
|
||||
|
||||
// VersionedPromptInput 表示旧版提示词工作流中的草稿和版本字段。
|
||||
type VersionedPromptInput struct {
|
||||
Code string
|
||||
Name string
|
||||
Category string
|
||||
Type string
|
||||
Content string
|
||||
Variables []string
|
||||
VersionNote string
|
||||
}
|
||||
|
||||
// UserPromptList 表示指定类型的可用提示词及用户当前选择。
|
||||
type UserPromptList struct {
|
||||
Prompts []map[string]any
|
||||
SelectedID *uuid.UUID
|
||||
}
|
||||
|
||||
// PersistenceError 标识应由接口层转换为内部错误的数据库故障。
|
||||
type PersistenceError struct {
|
||||
Err error
|
||||
}
|
||||
|
||||
// Error 返回底层数据库错误文本,供日志与测试定位。
|
||||
func (e PersistenceError) Error() string {
|
||||
return e.Err.Error()
|
||||
}
|
||||
|
||||
// Unwrap 暴露底层错误,保留 errors.Is 和 errors.As 语义。
|
||||
func (e PersistenceError) Unwrap() error {
|
||||
return e.Err
|
||||
}
|
||||
|
||||
var fixedDramaParsePromptTypes = []string{"剧本解析", "角色、场景、道具解析"}
|
||||
|
||||
// NewService 创建提示词服务,数据库连接仅在模块内部使用。
|
||||
func NewService(db *gorm.DB) *Service {
|
||||
return &Service{db: db}
|
||||
}
|
||||
|
||||
// IsFixedDramaParseType 判断提示词类型是否由后端固定配置且禁止用户编辑。
|
||||
func IsFixedDramaParseType(promptType string) bool {
|
||||
for _, fixedType := range fixedDramaParsePromptTypes {
|
||||
if strings.TrimSpace(promptType) == fixedType {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// IsPersistenceError 判断错误是否属于需要隐藏细节的数据库故障。
|
||||
func IsPersistenceError(err error) bool {
|
||||
var target PersistenceError
|
||||
return errors.As(err, &target)
|
||||
}
|
||||
|
||||
// ListSystemPrompts 查询管理端可维护的系统提示词,不返回后端固定配置。
|
||||
func (s *Service) ListSystemPrompts(keyword, promptType string, page, size int) ([]map[string]any, int64, error) {
|
||||
query := s.db.Table("prompts p").Where("p.deleted_at IS NULL AND p.scope='system' AND p.type NOT IN ?", fixedDramaParsePromptTypes)
|
||||
if keyword = strings.TrimSpace(keyword); keyword != "" {
|
||||
query = query.Where("p.name::text ILIKE ? OR p.content ILIKE ?", "%"+keyword+"%", "%"+keyword+"%")
|
||||
}
|
||||
if promptType = strings.TrimSpace(promptType); promptType != "" {
|
||||
query = query.Where("p.type=?", promptType)
|
||||
}
|
||||
var total int64
|
||||
if err := query.Count(&total).Error; err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
items := make([]map[string]any, 0)
|
||||
err := query.Select("p.id,p.name,p.type,p.content,p.scope,p.owner_user_id,p.created_at,p.updated_at").
|
||||
Order("p.updated_at DESC").Offset((page - 1) * size).Limit(size).Find(&items).Error
|
||||
return items, total, err
|
||||
}
|
||||
|
||||
// SaveSystemPrompt 校验并新增或更新管理端系统提示词。
|
||||
func (s *Service) SaveSystemPrompt(promptID string, input SystemPromptInput) (SystemPrompt, error) {
|
||||
input.Name = strings.TrimSpace(input.Name)
|
||||
input.Type = strings.TrimSpace(input.Type)
|
||||
input.Content = strings.TrimSpace(input.Content)
|
||||
if input.Name == "" || input.Type == "" || input.Content == "" {
|
||||
return SystemPrompt{}, errors.New("name、type、content不能为空")
|
||||
}
|
||||
if IsFixedDramaParseType(input.Type) {
|
||||
return SystemPrompt{}, errors.New("该提示词已固定在后端代码中,不支持管理")
|
||||
}
|
||||
if promptID == "" {
|
||||
promptID = uuid.NewString()
|
||||
if err := s.db.Exec("INSERT INTO prompts(id,name,type,content,scope,owner_user_id) VALUES(?,?,?,?,'system',NULL)", promptID, input.Name, input.Type, input.Content).Error; err != nil {
|
||||
return SystemPrompt{}, err
|
||||
}
|
||||
} else if err := s.db.Exec("UPDATE prompts SET name=?,type=?,content=?,updated_at=CURRENT_TIMESTAMP WHERE id=? AND scope='system' AND deleted_at IS NULL", input.Name, input.Type, input.Content, promptID).Error; err != nil {
|
||||
return SystemPrompt{}, err
|
||||
}
|
||||
return SystemPrompt{ID: promptID, Name: input.Name, Type: input.Type, Content: input.Content}, nil
|
||||
}
|
||||
|
||||
// DeleteSystemPrompt 删除管理端指定的系统提示词。
|
||||
func (s *Service) DeleteSystemPrompt(promptID string) error {
|
||||
return s.db.Exec("DELETE FROM prompts WHERE id=? AND scope='system'", promptID).Error
|
||||
}
|
||||
|
||||
// SaveVersionedPrompt 保存旧版管理接口使用的提示词草稿及版本快照。
|
||||
func (s *Service) SaveVersionedPrompt(adminID uuid.UUID, promptID string, input VersionedPromptInput) (string, error) {
|
||||
if strings.TrimSpace(input.Category) == "" {
|
||||
input.Category = input.Type
|
||||
}
|
||||
if strings.TrimSpace(input.Type) == "" {
|
||||
input.Type = input.Category
|
||||
}
|
||||
err := s.db.Transaction(func(tx *gorm.DB) error {
|
||||
version := 1
|
||||
if promptID == "" {
|
||||
promptID = uuid.NewString()
|
||||
if err := tx.Exec("INSERT INTO prompts(id,code,name,category,type,content,status) VALUES(?,?,?,?,?,?,'draft')", promptID, input.Code, input.Name, input.Category, input.Type, input.Content).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
} else {
|
||||
if err := tx.Raw("SELECT coalesce(max(version),0)+1 FROM prompt_versions WHERE prompt_id=?", promptID).Scan(&version).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if err := tx.Exec("UPDATE prompts SET code=?,name=?,category=?,type=?,content=? WHERE id=? AND deleted_at IS NULL", input.Code, input.Name, input.Category, input.Type, input.Content, promptID).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
variables, _ := json.Marshal(input.Variables)
|
||||
versionID := uuid.NewString()
|
||||
if err := tx.Exec("INSERT INTO prompt_versions(id,prompt_id,version,content,variables,version_note,status,operator_id) VALUES(?,?,?,?,?::jsonb,?,'draft',?)", versionID, promptID, version, input.Content, string(variables), input.VersionNote, adminID).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Exec("UPDATE prompts SET current_version_id=?,status='draft' WHERE id=?", versionID, promptID).Error
|
||||
})
|
||||
return promptID, err
|
||||
}
|
||||
|
||||
// ApplyPromptAction 执行旧版提示词接口的发布、停用或版本回滚事务。
|
||||
func (s *Service) ApplyPromptAction(adminID uuid.UUID, promptID, action, versionID, versionNote string) error {
|
||||
return s.db.Transaction(func(tx *gorm.DB) error {
|
||||
switch action {
|
||||
case "publish":
|
||||
if err := tx.Exec("UPDATE prompt_versions SET status=CASE WHEN id=(SELECT current_version_id FROM prompts WHERE id=?) THEN 'published' ELSE 'retired' END WHERE prompt_id=?", promptID, promptID).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Exec("UPDATE prompts SET status='published' WHERE id=?", promptID).Error
|
||||
case "disable":
|
||||
return tx.Exec("UPDATE prompts SET status='disabled' WHERE id=?", promptID).Error
|
||||
case "rollback":
|
||||
var source struct {
|
||||
Content string
|
||||
Variables string
|
||||
}
|
||||
if err := tx.Raw("SELECT content,variables::text FROM prompt_versions WHERE id=? AND prompt_id=?", versionID, promptID).Scan(&source).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
var version int
|
||||
if err := tx.Raw("SELECT coalesce(max(version),0)+1 FROM prompt_versions WHERE prompt_id=?", promptID).Scan(&version).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
newVersionID := uuid.NewString()
|
||||
if err := tx.Exec("INSERT INTO prompt_versions(id,prompt_id,version,content,variables,version_note,status,operator_id) VALUES(?,?,?,?,?::jsonb,?,'draft',?)", newVersionID, promptID, version, source.Content, source.Variables, versionNote, adminID).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Exec("UPDATE prompts SET current_version_id=?,status='draft' WHERE id=?", newVersionID, promptID).Error
|
||||
default:
|
||||
return errors.New("不支持的提示词操作")
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// ListPromptHistory 查询旧版提示词接口保留的版本历史。
|
||||
func (s *Service) ListPromptHistory(promptID string) ([]map[string]any, error) {
|
||||
items := make([]map[string]any, 0)
|
||||
err := s.db.Table("prompt_versions v").
|
||||
Select("v.id,v.version,v.content,v.variables,v.version_note,v.status,v.created_at,a.username AS operator").
|
||||
Joins("LEFT JOIN admin_users a ON a.id=v.operator_id").Where("v.prompt_id=?", promptID).
|
||||
Order("v.version DESC").Find(&items).Error
|
||||
return items, err
|
||||
}
|
||||
|
||||
// ListUserPrompts 查询系统提示词、当前用户自定义提示词及用户选择。
|
||||
func (s *Service) ListUserPrompts(userID uuid.UUID, promptType string) (UserPromptList, error) {
|
||||
if IsFixedDramaParseType(promptType) {
|
||||
return UserPromptList{Prompts: []map[string]any{}}, nil
|
||||
}
|
||||
var prompts []map[string]any
|
||||
err := s.db.Table("prompts").Select("id,name,type,content,scope,owner_user_id,created_at,updated_at,(scope='user') AS editable").
|
||||
Where("deleted_at IS NULL AND type=? AND (scope='system' OR (scope='user' AND owner_user_id=?))", promptType, userID).
|
||||
Order("CASE WHEN scope='system' THEN 0 ELSE 1 END, updated_at DESC").Find(&prompts).Error
|
||||
if err != nil {
|
||||
return UserPromptList{}, err
|
||||
}
|
||||
var selected struct{ PromptID *uuid.UUID }
|
||||
// 用户未设置偏好时返回空属于正常情况,用 Limit(1).Find 避免触发 GORM 的 record not found 日志
|
||||
s.db.Table("user_prompt_preferences").Select("prompt_id").Where("user_id=? AND prompt_type=?", userID, promptType).Limit(1).Find(&selected)
|
||||
return UserPromptList{Prompts: prompts, SelectedID: selected.PromptID}, nil
|
||||
}
|
||||
|
||||
// SelectUserPrompt 更新用户对指定提示词类型的选择,空值表示关闭选择。
|
||||
func (s *Service) SelectUserPrompt(userID uuid.UUID, promptType string, promptID *uuid.UUID) error {
|
||||
if IsFixedDramaParseType(promptType) {
|
||||
return errors.New("该提示词由平台固定配置,不支持选择")
|
||||
}
|
||||
if promptID == nil {
|
||||
if err := s.db.Exec("DELETE FROM user_prompt_preferences WHERE user_id=? AND prompt_type=?", userID, promptType).Error; err != nil {
|
||||
return PersistenceError{Err: err}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
var count int64
|
||||
if err := s.db.Table("prompts").Where("id=? AND type=? AND deleted_at IS NULL AND (scope='system' OR (scope='user' AND owner_user_id=?))", *promptID, promptType, userID).Count(&count).Error; err != nil {
|
||||
return PersistenceError{Err: err}
|
||||
}
|
||||
if count == 0 {
|
||||
return gorm.ErrRecordNotFound
|
||||
}
|
||||
if err := s.db.Exec(`INSERT INTO user_prompt_preferences(id,user_id,prompt_id,prompt_type) VALUES(gen_random_uuid(),?,?,?) ON CONFLICT(user_id,prompt_type) DO UPDATE SET prompt_id=excluded.prompt_id,updated_at=CURRENT_TIMESTAMP`, userID, *promptID, promptType).Error; err != nil {
|
||||
return PersistenceError{Err: err}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// CreateCustomPrompt 校验并创建当前用户拥有的自定义提示词。
|
||||
func (s *Service) CreateCustomPrompt(userID uuid.UUID, input CustomPromptInput) (CustomPrompt, error) {
|
||||
input = normalizeCustomPromptInput(input)
|
||||
if input.Name == "" || input.Type == "" || input.Content == "" {
|
||||
return CustomPrompt{}, errors.New("name、type、content不能为空")
|
||||
}
|
||||
if IsFixedDramaParseType(input.Type) {
|
||||
return CustomPrompt{}, errors.New("该提示词由平台固定配置,不支持自定义")
|
||||
}
|
||||
item := CustomPrompt{ID: uuid.New(), Name: input.Name, Type: input.Type, Content: input.Content, Scope: "user", Editable: true}
|
||||
if err := s.db.Exec("INSERT INTO prompts(id,name,type,content,scope,owner_user_id) VALUES(?,?,?,?,'user',?)", item.ID, item.Name, item.Type, item.Content, userID).Error; err != nil {
|
||||
return CustomPrompt{}, err
|
||||
}
|
||||
return item, nil
|
||||
}
|
||||
|
||||
// UpdateCustomPrompt 校验并更新当前用户拥有的自定义提示词。
|
||||
func (s *Service) UpdateCustomPrompt(userID, promptID uuid.UUID, input CustomPromptInput) (CustomPrompt, error) {
|
||||
input = normalizeCustomPromptInput(input)
|
||||
if input.Name == "" || input.Content == "" {
|
||||
return CustomPrompt{}, errors.New("name、content不能为空")
|
||||
}
|
||||
if IsFixedDramaParseType(input.Type) {
|
||||
return CustomPrompt{}, errors.New("该提示词由平台固定配置,不支持自定义")
|
||||
}
|
||||
result := s.db.Exec("UPDATE prompts SET name=?,type=COALESCE(NULLIF(?,''),type),content=?,updated_at=CURRENT_TIMESTAMP WHERE id=? AND scope='user' AND owner_user_id=? AND deleted_at IS NULL", input.Name, input.Type, input.Content, promptID, userID)
|
||||
if result.Error != nil {
|
||||
return CustomPrompt{}, result.Error
|
||||
}
|
||||
if result.RowsAffected == 0 {
|
||||
return CustomPrompt{}, gorm.ErrRecordNotFound
|
||||
}
|
||||
return CustomPrompt{ID: promptID, Name: input.Name, Type: input.Type, Content: input.Content, Scope: "user", Editable: true}, nil
|
||||
}
|
||||
|
||||
// DeleteCustomPrompt 在同一事务中删除用户选择关系及其自定义提示词。
|
||||
func (s *Service) DeleteCustomPrompt(userID, promptID uuid.UUID) error {
|
||||
return s.db.Transaction(func(tx *gorm.DB) error {
|
||||
if err := tx.Exec("DELETE FROM user_prompt_preferences WHERE prompt_id=? AND user_id=?", promptID, userID).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
result := tx.Exec("DELETE FROM prompts WHERE id=? AND scope='user' AND owner_user_id=?", promptID, userID)
|
||||
if result.Error != nil {
|
||||
return result.Error
|
||||
}
|
||||
if result.RowsAffected == 0 {
|
||||
return gorm.ErrRecordNotFound
|
||||
}
|
||||
return nil
|
||||
})
|
||||
}
|
||||
|
||||
// normalizeCustomPromptInput 清理用户输入两端空白,保证校验和写入使用一致值。
|
||||
func normalizeCustomPromptInput(input CustomPromptInput) CustomPromptInput {
|
||||
input.Name = strings.TrimSpace(input.Name)
|
||||
input.Type = strings.TrimSpace(input.Type)
|
||||
input.Content = strings.TrimSpace(input.Content)
|
||||
return input
|
||||
}
|
||||
@@ -0,0 +1,24 @@
|
||||
// 提示词模块单元测试,验证固定类型规则和用户输入标准化行为。
|
||||
package prompt
|
||||
|
||||
import "testing"
|
||||
|
||||
// TestIsFixedDramaParseType 验证固定提示词类型及空白清理规则。
|
||||
func TestIsFixedDramaParseType(t *testing.T) {
|
||||
for _, value := range []string{"剧本解析", " 角色、场景、道具解析 "} {
|
||||
if !IsFixedDramaParseType(value) {
|
||||
t.Fatalf("expected %q to be fixed", value)
|
||||
}
|
||||
}
|
||||
if IsFixedDramaParseType("剧本分析") {
|
||||
t.Fatal("剧本分析不应被识别为固定解析提示词")
|
||||
}
|
||||
}
|
||||
|
||||
// TestNormalizeCustomPromptInput 验证提示词字段在校验和写入前统一清理两端空白。
|
||||
func TestNormalizeCustomPromptInput(t *testing.T) {
|
||||
input := normalizeCustomPromptInput(CustomPromptInput{Name: " 名称 ", Type: " 类型\n", Content: " 内容 "})
|
||||
if input.Name != "名称" || input.Type != "类型" || input.Content != "内容" {
|
||||
t.Fatalf("unexpected normalized input: %#v", input)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,64 @@
|
||||
package apimart
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"juhe-factory/api/internal/provider"
|
||||
)
|
||||
|
||||
type BalanceResult struct {
|
||||
RemainBalance float64 `json:"remain_balance"`
|
||||
RemainCredits float64 `json:"remain_credits"`
|
||||
UsedBalance float64 `json:"used_balance"`
|
||||
UsedCredits float64 `json:"used_credits"`
|
||||
UnlimitedQuota bool `json:"unlimited_quota"`
|
||||
}
|
||||
|
||||
func (c Client) Balance(ctx context.Context, baseURL, apiKey string) (BalanceResult, error) {
|
||||
var result BalanceResult
|
||||
request, err := http.NewRequestWithContext(ctx, http.MethodGet, strings.TrimRight(baseURL, "/")+"/user/balance", nil)
|
||||
if err != nil {
|
||||
return result, err
|
||||
}
|
||||
request.Header.Set("Authorization", "Bearer "+apiKey)
|
||||
|
||||
client := c.HTTPClient
|
||||
if client == nil {
|
||||
client = http.DefaultClient
|
||||
}
|
||||
response, err := client.Do(request)
|
||||
if err != nil {
|
||||
return result, err
|
||||
}
|
||||
defer response.Body.Close()
|
||||
|
||||
data, err := io.ReadAll(io.LimitReader(response.Body, 1024*1024))
|
||||
if err != nil {
|
||||
return result, err
|
||||
}
|
||||
if response.StatusCode < 200 || response.StatusCode >= 300 {
|
||||
return result, &provider.HTTPError{StatusCode: response.StatusCode, Body: safeBody(data), RetryAfter: response.Header.Get("Retry-After")}
|
||||
}
|
||||
|
||||
var decoded struct {
|
||||
Success bool `json:"success"`
|
||||
Message string `json:"message"`
|
||||
BalanceResult
|
||||
}
|
||||
if err := json.Unmarshal(data, &decoded); err != nil {
|
||||
return result, fmt.Errorf("余额响应不是有效 JSON: %w", err)
|
||||
}
|
||||
if !decoded.Success {
|
||||
if decoded.Message == "" {
|
||||
return result, errors.New("获取渠道余额失败")
|
||||
}
|
||||
return result, errors.New(decoded.Message)
|
||||
}
|
||||
return decoded.BalanceResult, nil
|
||||
}
|
||||
@@ -0,0 +1,42 @@
|
||||
package apimart
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestBalanceUsesConfiguredBaseURLAndAPIKey(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) {
|
||||
if request.URL.Path != "/v1/user/balance" {
|
||||
t.Fatalf("unexpected path: %s", request.URL.Path)
|
||||
}
|
||||
if request.Header.Get("Authorization") != "Bearer test-key" {
|
||||
t.Fatalf("unexpected authorization header: %s", request.Header.Get("Authorization"))
|
||||
}
|
||||
fmt.Fprint(response, `{"success":true,"remain_balance":10.5,"remain_credits":105.25,"used_balance":2.3,"used_credits":2005.21388,"unlimited_quota":false}`)
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
result, err := NewClient(server.Client()).Balance(context.Background(), server.URL+"/v1", "test-key")
|
||||
if err != nil {
|
||||
t.Fatalf("Balance returned error: %v", err)
|
||||
}
|
||||
if result.RemainBalance != 10.5 || result.RemainCredits != 105.25 || result.UsedBalance != 2.3 || result.UsedCredits != 2005.21388 || result.UnlimitedQuota {
|
||||
t.Fatalf("unexpected result: %#v", result)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBalanceReturnsUpstreamMessage(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, _ *http.Request) {
|
||||
fmt.Fprint(response, `{"success":false,"message":"record not found"}`)
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
_, err := NewClient(server.Client()).Balance(context.Background(), server.URL, "test-key")
|
||||
if err == nil || err.Error() != "record not found" {
|
||||
t.Fatalf("expected upstream message, got %v", err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,141 @@
|
||||
package apimart
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"juhe-factory/api/internal/provider"
|
||||
)
|
||||
|
||||
type ChatResult struct {
|
||||
Content string
|
||||
Raw json.RawMessage
|
||||
FinishReason string
|
||||
InputTokens int64
|
||||
OutputTokens int64
|
||||
TotalTokens int64
|
||||
UsageRaw json.RawMessage
|
||||
}
|
||||
|
||||
func (c Client) Chat(ctx context.Context, baseURL, apiKey string, payload map[string]any) (string, json.RawMessage, error) {
|
||||
result, err := c.ChatWithUsage(ctx, baseURL, apiKey, payload)
|
||||
return result.Content, result.Raw, err
|
||||
}
|
||||
|
||||
func (c Client) ChatWithUsage(ctx context.Context, baseURL, apiKey string, payload map[string]any) (ChatResult, error) {
|
||||
var result ChatResult
|
||||
requestPayload := make(map[string]any, len(payload)+1)
|
||||
for key, value := range payload {
|
||||
requestPayload[key] = value
|
||||
}
|
||||
requestPayload["stream"] = false
|
||||
body, err := json.Marshal(requestPayload)
|
||||
if err != nil {
|
||||
return result, err
|
||||
}
|
||||
request, err := http.NewRequestWithContext(ctx, http.MethodPost, strings.TrimRight(baseURL, "/")+"/chat/completions", bytes.NewReader(body))
|
||||
if err != nil {
|
||||
return result, err
|
||||
}
|
||||
request.Header.Set("Authorization", "Bearer "+apiKey)
|
||||
request.Header.Set("Content-Type", "application/json")
|
||||
response, err := c.HTTPClient.Do(request)
|
||||
if err != nil {
|
||||
return result, err
|
||||
}
|
||||
defer response.Body.Close()
|
||||
data, err := io.ReadAll(io.LimitReader(response.Body, 16*1024*1024))
|
||||
if err != nil {
|
||||
return result, err
|
||||
}
|
||||
if response.StatusCode < 200 || response.StatusCode >= 300 {
|
||||
result.Raw = append(json.RawMessage(nil), data...)
|
||||
return result, &provider.HTTPError{StatusCode: response.StatusCode, Body: safeBody(data), RetryAfter: response.Header.Get("Retry-After")}
|
||||
}
|
||||
type chatChoice struct {
|
||||
Message struct {
|
||||
Content any `json:"content"`
|
||||
} `json:"message"`
|
||||
FinishReason string `json:"finish_reason"`
|
||||
}
|
||||
var decoded struct {
|
||||
Choices []chatChoice `json:"choices"`
|
||||
Usage json.RawMessage `json:"usage"`
|
||||
Data struct {
|
||||
Choices []chatChoice `json:"choices"`
|
||||
Usage json.RawMessage `json:"usage"`
|
||||
} `json:"data"`
|
||||
}
|
||||
if err := json.Unmarshal(data, &decoded); err != nil {
|
||||
result.Raw = append(json.RawMessage(nil), data...)
|
||||
return result, fmt.Errorf("文本模型响应不是有效 JSON: %w;上游响应: %s", err, safeBody(data))
|
||||
}
|
||||
result.Raw = append(json.RawMessage(nil), data...)
|
||||
result.UsageRaw = decoded.Usage
|
||||
choices := decoded.Choices
|
||||
if len(choices) == 0 {
|
||||
choices = decoded.Data.Choices
|
||||
if len(result.UsageRaw) == 0 {
|
||||
result.UsageRaw = decoded.Data.Usage
|
||||
}
|
||||
}
|
||||
if len(choices) == 0 {
|
||||
return result, fmt.Errorf("文本模型响应缺少 choices;上游响应: %s", safeBody(data))
|
||||
}
|
||||
result.FinishReason = choices[0].FinishReason
|
||||
if len(result.UsageRaw) > 0 {
|
||||
var usage map[string]any
|
||||
if json.Unmarshal(result.UsageRaw, &usage) == nil {
|
||||
result.InputTokens = intValue(usage["prompt_tokens"], usage["input_tokens"])
|
||||
result.OutputTokens = intValue(usage["completion_tokens"], usage["output_tokens"])
|
||||
result.TotalTokens = intValue(usage["total_tokens"])
|
||||
if result.TotalTokens == 0 {
|
||||
result.TotalTokens = result.InputTokens + result.OutputTokens
|
||||
}
|
||||
}
|
||||
}
|
||||
switch content := choices[0].Message.Content.(type) {
|
||||
case string:
|
||||
result.Content = content
|
||||
return result, nil
|
||||
case []any:
|
||||
parts := make([]string, 0)
|
||||
for _, item := range content {
|
||||
if block, ok := item.(map[string]any); ok {
|
||||
if text, ok := block["text"].(string); ok {
|
||||
parts = append(parts, text)
|
||||
}
|
||||
}
|
||||
}
|
||||
result.Content = strings.Join(parts, "\n")
|
||||
return result, nil
|
||||
default:
|
||||
return result, errors.New("文本模型响应内容无效")
|
||||
}
|
||||
}
|
||||
|
||||
func intValue(values ...any) int64 {
|
||||
for _, value := range values {
|
||||
switch number := value.(type) {
|
||||
case float64:
|
||||
if number > 0 {
|
||||
return int64(number)
|
||||
}
|
||||
case int64:
|
||||
if number > 0 {
|
||||
return number
|
||||
}
|
||||
case json.Number:
|
||||
if parsed, err := number.Int64(); err == nil && parsed > 0 {
|
||||
return parsed
|
||||
}
|
||||
}
|
||||
}
|
||||
return 0
|
||||
}
|
||||
@@ -0,0 +1,58 @@
|
||||
package apimart
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"juhe-factory/api/internal/provider"
|
||||
)
|
||||
|
||||
func TestChatWithUsageReadsOpenAICompatibleAndNestedUsage(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
response string
|
||||
content string
|
||||
input int64
|
||||
output int64
|
||||
total int64
|
||||
finish string
|
||||
}{
|
||||
{name: "top level", response: `{"choices":[{"message":{"content":"result"},"finish_reason":"stop"}],"usage":{"prompt_tokens":120,"completion_tokens":30,"total_tokens":150}}`, content: "result", input: 120, output: 30, total: 150, finish: "stop"},
|
||||
{name: "nested alternate names", response: `{"data":{"choices":[{"message":{"content":[{"type":"text","text":"part one"},{"type":"text","text":"part two"}]}}],"usage":{"input_tokens":44,"output_tokens":11}}}`, content: "part one\npart two", input: 44, output: 11, total: 55},
|
||||
}
|
||||
for _, test := range cases {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) {
|
||||
if request.URL.Path != "/chat/completions" {
|
||||
t.Fatalf("unexpected path: %s", request.URL.Path)
|
||||
}
|
||||
response.Header().Set("Content-Type", "application/json")
|
||||
fmt.Fprint(response, test.response)
|
||||
}))
|
||||
defer server.Close()
|
||||
result, err := NewClient(server.Client()).ChatWithUsage(context.Background(), server.URL, "key", map[string]any{"model": "text-model"})
|
||||
if err != nil {
|
||||
t.Fatalf("ChatWithUsage returned error: %v", err)
|
||||
}
|
||||
if result.Content != test.content || result.InputTokens != test.input || result.OutputTokens != test.output || result.TotalTokens != test.total || result.FinishReason != test.finish {
|
||||
t.Fatalf("unexpected result: %#v", result)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestChatWithUsageReturnsExplicitUpstreamFailure(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, _ *http.Request) {
|
||||
response.WriteHeader(http.StatusBadGateway)
|
||||
fmt.Fprint(response, `{"error":"relay unavailable"}`)
|
||||
}))
|
||||
defer server.Close()
|
||||
_, err := NewClient(server.Client()).ChatWithUsage(context.Background(), server.URL, "key", map[string]any{"model": "text-model"})
|
||||
httpErr, ok := err.(*provider.HTTPError)
|
||||
if !ok || httpErr.StatusCode != http.StatusBadGateway {
|
||||
t.Fatalf("expected HTTPError with relay status, got %T %v", err, err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,64 @@
|
||||
package apimart
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"juhe-factory/api/internal/provider"
|
||||
)
|
||||
|
||||
type Client struct {
|
||||
HTTPClient *http.Client
|
||||
}
|
||||
|
||||
func NewClient(client *http.Client) Client {
|
||||
return Client{HTTPClient: client}
|
||||
}
|
||||
|
||||
func (c Client) requestJSON(ctx context.Context, method, baseURL, path, apiKey, requestID string, payload map[string]any) ([]byte, error) {
|
||||
var body io.Reader
|
||||
if payload != nil {
|
||||
encoded, err := json.Marshal(payload)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
body = bytes.NewReader(encoded)
|
||||
}
|
||||
request, err := http.NewRequestWithContext(ctx, method, strings.TrimRight(baseURL, "/")+path, body)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
request.Header.Set("Authorization", "Bearer "+apiKey)
|
||||
if payload != nil {
|
||||
request.Header.Set("Content-Type", "application/json")
|
||||
}
|
||||
if requestID != "" {
|
||||
request.Header.Set("Idempotency-Key", requestID)
|
||||
request.Header.Set("X-Request-ID", requestID)
|
||||
}
|
||||
response, err := c.HTTPClient.Do(request)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer response.Body.Close()
|
||||
data, err := io.ReadAll(io.LimitReader(response.Body, 4*1024*1024))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if response.StatusCode < 200 || response.StatusCode >= 300 {
|
||||
return nil, &provider.HTTPError{StatusCode: response.StatusCode, Body: safeBody(data), RetryAfter: response.Header.Get("Retry-After")}
|
||||
}
|
||||
return data, nil
|
||||
}
|
||||
|
||||
func safeBody(data []byte) string {
|
||||
value := strings.TrimSpace(string(data))
|
||||
if len(value) > 2000 {
|
||||
value = value[:2000]
|
||||
}
|
||||
return value
|
||||
}
|
||||
@@ -0,0 +1,103 @@
|
||||
package apimart
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"juhe-factory/api/internal/provider"
|
||||
)
|
||||
|
||||
func (c Client) SubmitImage(ctx context.Context, baseURL, apiKey, requestID string, payload map[string]any) (provider.SubmitResult, error) {
|
||||
data, err := c.requestJSON(ctx, http.MethodPost, baseURL, "/images/generations", apiKey, requestID, payload)
|
||||
if err != nil {
|
||||
return provider.SubmitResult{}, err
|
||||
}
|
||||
result := provider.SubmitResult{Raw: append(json.RawMessage(nil), data...)}
|
||||
var decoded any
|
||||
if err := json.Unmarshal(data, &decoded); err != nil {
|
||||
return provider.SubmitResult{}, errors.New("中转站提交响应不是有效 JSON")
|
||||
}
|
||||
result.TaskID = findString(decoded, "task_id", "id", "request_id")
|
||||
result.URL = findURL(decoded)
|
||||
if result.TaskID == "" && result.URL == "" {
|
||||
return provider.SubmitResult{}, errors.New("中转站提交响应缺少任务 ID 或结果 URL")
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
// PollImage 通过 APIMart 统一任务接口查询图片生成状态和结果。
|
||||
func (c Client) PollImage(ctx context.Context, baseURL, apiKey, taskID string) (provider.PollResult, error) {
|
||||
data, err := c.requestJSON(ctx, http.MethodGet, baseURL, "/tasks/"+taskID, apiKey, "", nil)
|
||||
if err != nil {
|
||||
return provider.PollResult{}, err
|
||||
}
|
||||
var decoded any
|
||||
if err := json.Unmarshal(data, &decoded); err != nil {
|
||||
return provider.PollResult{}, errors.New("中转站轮询响应不是有效 JSON")
|
||||
}
|
||||
status := strings.ToLower(findString(decoded, "task_status", "status", "state"))
|
||||
if status == "" {
|
||||
status = "processing"
|
||||
}
|
||||
return provider.PollResult{Status: status, URL: findURL(decoded), Error: findString(decoded, "error_message", "error", "message", "detail"), Raw: append(json.RawMessage(nil), data...)}, nil
|
||||
}
|
||||
|
||||
func findURL(value any) string {
|
||||
var walk func(any, string) string
|
||||
walk = func(node any, parentKey string) string {
|
||||
switch data := node.(type) {
|
||||
case string:
|
||||
if (strings.Contains(strings.ToLower(parentKey), "url") || strings.HasPrefix(data, "http://") || strings.HasPrefix(data, "https://")) && (strings.HasPrefix(data, "https://") || strings.HasPrefix(data, "http://")) {
|
||||
return data
|
||||
}
|
||||
case map[string]any:
|
||||
for key, child := range data {
|
||||
if found := walk(child, key); found != "" {
|
||||
return found
|
||||
}
|
||||
}
|
||||
case []any:
|
||||
for _, child := range data {
|
||||
if found := walk(child, parentKey); found != "" {
|
||||
return found
|
||||
}
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
return walk(value, "")
|
||||
}
|
||||
|
||||
func findString(value any, keys ...string) string {
|
||||
wanted := map[string]bool{}
|
||||
for _, key := range keys {
|
||||
wanted[strings.ToLower(key)] = true
|
||||
}
|
||||
return findMatchingString(value, func(key, _ string) bool { return wanted[strings.ToLower(key)] })
|
||||
}
|
||||
|
||||
func findMatchingString(value any, match func(key, value string) bool) string {
|
||||
switch data := value.(type) {
|
||||
case map[string]any:
|
||||
for key, child := range data {
|
||||
if text, ok := child.(string); ok && match(key, text) {
|
||||
return text
|
||||
}
|
||||
}
|
||||
for _, child := range data {
|
||||
if found := findMatchingString(child, match); found != "" {
|
||||
return found
|
||||
}
|
||||
}
|
||||
case []any:
|
||||
for _, child := range data {
|
||||
if found := findMatchingString(child, match); found != "" {
|
||||
return found
|
||||
}
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
@@ -0,0 +1,48 @@
|
||||
// 图片生成中转站适配器测试,验证任务提交、状态轮询及成功结果解析契约。
|
||||
package apimart
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// TestSubmitImageUsesGenerationEndpoint 验证图片生成使用配置渠道下的标准提交接口。
|
||||
func TestSubmitImageUsesGenerationEndpoint(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) {
|
||||
if request.Method != http.MethodPost || request.URL.Path != "/v1/images/generations" {
|
||||
t.Fatalf("unexpected request: %s %s", request.Method, request.URL.Path)
|
||||
}
|
||||
fmt.Fprint(response, `{"code":200,"data":[{"status":"submitted","task_id":"task-image-1"}]}`)
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
result, err := NewClient(server.Client()).SubmitImage(context.Background(), server.URL+"/v1", "test-key", "request-1", map[string]any{"model": "gpt-image-2"})
|
||||
if err != nil {
|
||||
t.Fatalf("SubmitImage returned error: %v", err)
|
||||
}
|
||||
if result.TaskID != "task-image-1" {
|
||||
t.Fatalf("unexpected task id: %s", result.TaskID)
|
||||
}
|
||||
}
|
||||
|
||||
// TestPollImageUsesTaskEndpoint 验证图片任务使用 APIMart 统一任务接口查询并解析完成结果。
|
||||
func TestPollImageUsesTaskEndpoint(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) {
|
||||
if request.Method != http.MethodGet || request.URL.Path != "/v1/tasks/task-image-1" {
|
||||
t.Fatalf("unexpected request: %s %s", request.Method, request.URL.Path)
|
||||
}
|
||||
fmt.Fprint(response, `{"code":200,"data":{"id":"task-image-1","status":"completed","result":{"images":[{"url":["https://example.com/result.png"]}]}}}`)
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
result, err := NewClient(server.Client()).PollImage(context.Background(), server.URL+"/v1", "test-key", "task-image-1")
|
||||
if err != nil {
|
||||
t.Fatalf("PollImage returned error: %v", err)
|
||||
}
|
||||
if result.Status != "completed" || result.URL != "https://example.com/result.png" {
|
||||
t.Fatalf("unexpected result: %#v", result)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,225 @@
|
||||
package apimart
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"juhe-factory/api/internal/provider"
|
||||
)
|
||||
|
||||
type Seedance2 struct {
|
||||
Gateway Client
|
||||
}
|
||||
|
||||
type ReferenceCapability struct {
|
||||
Supported bool `json:"supported"`
|
||||
MaxCount int `json:"max_count"`
|
||||
MinTotalDurationMS int64 `json:"min_total_duration_ms,omitempty"`
|
||||
MaxTotalDurationMS int64 `json:"max_total_duration_ms,omitempty"`
|
||||
MinShortSidePixels int `json:"min_short_side_pixels,omitempty"`
|
||||
MaxShortSidePixels int `json:"max_short_side_pixels,omitempty"`
|
||||
}
|
||||
|
||||
type VideoModelCapabilities struct {
|
||||
// DurationMinSeconds / DurationMaxSeconds 表示视频模型支持的生成时长范围(秒),供前端可视化选择与后端校验。
|
||||
DurationMinSeconds int `json:"duration_min_seconds"`
|
||||
DurationMaxSeconds int `json:"duration_max_seconds"`
|
||||
ReferenceImages ReferenceCapability `json:"reference_images"`
|
||||
ReferenceAudios ReferenceCapability `json:"reference_audios"`
|
||||
ReferenceVideos ReferenceCapability `json:"reference_videos"`
|
||||
}
|
||||
|
||||
var seedance2Capabilities = map[string]VideoModelCapabilities{
|
||||
"doubao-seedance-2.0": seedance2ReferenceCapabilities(),
|
||||
"doubao-seedance-2.0-fast": seedance2ReferenceCapabilities(),
|
||||
"doubao-seedance-2.0-mini": seedance2ReferenceCapabilities(),
|
||||
}
|
||||
|
||||
func seedance2ReferenceCapabilities() VideoModelCapabilities {
|
||||
return VideoModelCapabilities{
|
||||
DurationMinSeconds: 5,
|
||||
DurationMaxSeconds: 15,
|
||||
ReferenceImages: ReferenceCapability{Supported: true, MaxCount: 9},
|
||||
ReferenceAudios: ReferenceCapability{Supported: true, MaxCount: 3, MaxTotalDurationMS: 15_000},
|
||||
ReferenceVideos: ReferenceCapability{Supported: true, MaxCount: 3, MinTotalDurationMS: 1_800, MaxTotalDurationMS: 15_200, MinShortSidePixels: 480, MaxShortSidePixels: 720},
|
||||
}
|
||||
}
|
||||
|
||||
func VideoCapabilities(modelName string) (VideoModelCapabilities, bool) {
|
||||
capabilities, ok := seedance2Capabilities[strings.TrimSpace(modelName)]
|
||||
return capabilities, ok
|
||||
}
|
||||
|
||||
func NewSeedance2(client *http.Client) Seedance2 {
|
||||
return Seedance2{Gateway: NewClient(client)}
|
||||
}
|
||||
|
||||
func BuildSeedance2Payload(modelName string, inputData json.RawMessage) (map[string]any, error) {
|
||||
modelName = strings.TrimSpace(modelName)
|
||||
capabilities, supported := VideoCapabilities(modelName)
|
||||
if !supported {
|
||||
return nil, fmt.Errorf("暂不支持视频模型 %q", modelName)
|
||||
}
|
||||
input := map[string]any{}
|
||||
if err := json.Unmarshal(inputData, &input); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
prompt := strings.TrimSpace(fmt.Sprint(input["prompt"]))
|
||||
if prompt == "" || prompt == "<nil>" {
|
||||
return nil, errors.New("生成提示词不能为空")
|
||||
}
|
||||
size := strings.TrimSpace(fmt.Sprint(input["size"]))
|
||||
if size != "16:9" && size != "9:16" {
|
||||
return nil, errors.New("项目画面比例不受 Seedance 2 支持")
|
||||
}
|
||||
duration, err := strconv.Atoi(fmt.Sprint(input["duration"]))
|
||||
if err != nil || duration < 5 || duration > 15 {
|
||||
return nil, errors.New("Seedance 2 视频时长必须为 5 到 15 秒")
|
||||
}
|
||||
images, err := seedanceURLList(input["image_urls"])
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("Seedance 2 参考图无效: %w", err)
|
||||
}
|
||||
if len(images) > capabilities.ReferenceImages.MaxCount {
|
||||
return nil, fmt.Errorf("Seedance 2 最多支持 %d 张参考图", capabilities.ReferenceImages.MaxCount)
|
||||
}
|
||||
audios, err := seedanceURLList(input["audio_urls"])
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("Seedance 2 参考音频无效: %w", err)
|
||||
}
|
||||
if len(audios) > capabilities.ReferenceAudios.MaxCount {
|
||||
return nil, fmt.Errorf("Seedance 2 最多支持 %d 个参考音频", capabilities.ReferenceAudios.MaxCount)
|
||||
}
|
||||
videos, err := seedanceURLList(input["video_urls"])
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("Seedance 2 参考视频无效: %w", err)
|
||||
}
|
||||
if len(videos) > capabilities.ReferenceVideos.MaxCount {
|
||||
return nil, fmt.Errorf("Seedance 2 最多支持 %d 个参考视频", capabilities.ReferenceVideos.MaxCount)
|
||||
}
|
||||
if len(audios) > 0 && len(images) == 0 && len(videos) == 0 {
|
||||
return nil, errors.New("Seedance 2 参考音频必须与参考图片或参考视频同时使用")
|
||||
}
|
||||
payload := map[string]any{
|
||||
"model": modelName, "prompt": prompt, "size": size, "duration": duration,
|
||||
"resolution": seedanceString(input["resolution"], "480p"),
|
||||
}
|
||||
if len(images) > 0 {
|
||||
payload["image_urls"] = images
|
||||
}
|
||||
if len(audios) > 0 {
|
||||
payload["audio_urls"] = audios
|
||||
}
|
||||
if len(videos) > 0 {
|
||||
payload["video_urls"] = videos
|
||||
}
|
||||
if value, ok := input["generate_audio"].(bool); ok {
|
||||
payload["generate_audio"] = value
|
||||
}
|
||||
if value, ok := input["return_last_frame"].(bool); ok {
|
||||
payload["return_last_frame"] = value
|
||||
}
|
||||
return payload, nil
|
||||
}
|
||||
|
||||
func (s Seedance2) Submit(ctx context.Context, baseURL, apiKey, requestID string, payload map[string]any) (provider.SubmitResult, error) {
|
||||
data, err := s.Gateway.requestJSON(ctx, http.MethodPost, baseURL, "/videos/generations", apiKey, requestID, payload)
|
||||
if err != nil {
|
||||
return provider.SubmitResult{}, err
|
||||
}
|
||||
var decoded any
|
||||
if err := json.Unmarshal(data, &decoded); err != nil {
|
||||
return provider.SubmitResult{}, errors.New("Seedance 2 提交响应不是有效 JSON")
|
||||
}
|
||||
result := provider.SubmitResult{TaskID: findString(decoded, "task_id", "id"), URL: findSeedanceVideoURL(decoded), Raw: append(json.RawMessage(nil), data...)}
|
||||
if result.TaskID == "" && result.URL == "" {
|
||||
return provider.SubmitResult{}, errors.New("Seedance 2 提交响应缺少任务 ID 或视频地址")
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func (s Seedance2) Poll(ctx context.Context, baseURL, apiKey, taskID string) (provider.PollResult, error) {
|
||||
data, err := s.Gateway.requestJSON(ctx, http.MethodGet, baseURL, "/tasks/"+taskID, apiKey, "", nil)
|
||||
if err != nil {
|
||||
return provider.PollResult{}, err
|
||||
}
|
||||
var decoded any
|
||||
if err := json.Unmarshal(data, &decoded); err != nil {
|
||||
return provider.PollResult{}, errors.New("Seedance 2 轮询响应不是有效 JSON")
|
||||
}
|
||||
status := strings.ToLower(findString(decoded, "task_status", "status", "state"))
|
||||
if status == "" {
|
||||
status = "processing"
|
||||
}
|
||||
return provider.PollResult{
|
||||
Status: status, URL: findSeedanceVideoURL(decoded),
|
||||
Error: findString(decoded, "error_message", "message", "detail"), Raw: append(json.RawMessage(nil), data...),
|
||||
}, nil
|
||||
}
|
||||
|
||||
func seedanceURLList(value any) ([]string, error) {
|
||||
if value == nil {
|
||||
return nil, nil
|
||||
}
|
||||
raw, ok := value.([]any)
|
||||
if !ok {
|
||||
return nil, errors.New("必须是 URL 数组")
|
||||
}
|
||||
result := make([]string, 0, len(raw))
|
||||
for _, item := range raw {
|
||||
url := strings.TrimSpace(fmt.Sprint(item))
|
||||
if !strings.HasPrefix(url, "https://") && !strings.HasPrefix(url, "http://") && !strings.HasPrefix(url, "asset://") {
|
||||
return nil, errors.New("包含不可访问的 URL")
|
||||
}
|
||||
result = append(result, url)
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func seedanceString(value any, fallback string) string {
|
||||
text := strings.TrimSpace(fmt.Sprint(value))
|
||||
if text == "" || text == "<nil>" {
|
||||
return fallback
|
||||
}
|
||||
return text
|
||||
}
|
||||
|
||||
func findSeedanceVideoURL(value any) string {
|
||||
if url := findURLInSeedanceVideoField(value); url != "" {
|
||||
return url
|
||||
}
|
||||
return findURL(value)
|
||||
}
|
||||
|
||||
func findURLInSeedanceVideoField(value any) string {
|
||||
switch data := value.(type) {
|
||||
case map[string]any:
|
||||
for key, child := range data {
|
||||
if strings.Contains(strings.ToLower(key), "video") {
|
||||
if text, ok := child.(string); ok && (strings.HasPrefix(text, "https://") || strings.HasPrefix(text, "http://")) {
|
||||
return text
|
||||
}
|
||||
if url := findURL(child); url != "" {
|
||||
return url
|
||||
}
|
||||
}
|
||||
}
|
||||
for _, child := range data {
|
||||
if url := findURLInSeedanceVideoField(child); url != "" {
|
||||
return url
|
||||
}
|
||||
}
|
||||
case []any:
|
||||
for _, child := range data {
|
||||
if url := findURLInSeedanceVideoField(child); url != "" {
|
||||
return url
|
||||
}
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
@@ -0,0 +1,48 @@
|
||||
package apimart
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestVideoCapabilities(t *testing.T) {
|
||||
capabilities, ok := VideoCapabilities("doubao-seedance-2.0")
|
||||
if !ok {
|
||||
t.Fatal("expected Seedance 2 capabilities")
|
||||
}
|
||||
if capabilities.ReferenceImages.MaxCount != 9 || capabilities.ReferenceAudios.MaxCount != 3 || capabilities.ReferenceVideos.MaxCount != 3 {
|
||||
t.Fatalf("unexpected reference limits: %+v", capabilities)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildSeedance2PayloadSupportsVideoAndAudioReferences(t *testing.T) {
|
||||
input := json.RawMessage(`{
|
||||
"prompt":"test",
|
||||
"size":"16:9",
|
||||
"duration":5,
|
||||
"video_urls":["https://example.com/reference.mp4"],
|
||||
"audio_urls":["https://example.com/reference.mp3"]
|
||||
}`)
|
||||
payload, err := BuildSeedance2Payload("doubao-seedance-2.0", input)
|
||||
if err != nil {
|
||||
t.Fatalf("BuildSeedance2Payload() error = %v", err)
|
||||
}
|
||||
if videos, ok := payload["video_urls"].([]string); !ok || len(videos) != 1 {
|
||||
t.Fatalf("video_urls = %#v", payload["video_urls"])
|
||||
}
|
||||
if audios, ok := payload["audio_urls"].([]string); !ok || len(audios) != 1 {
|
||||
t.Fatalf("audio_urls = %#v", payload["audio_urls"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildSeedance2PayloadRejectsTooManyVideos(t *testing.T) {
|
||||
input := json.RawMessage(`{
|
||||
"prompt":"test",
|
||||
"size":"16:9",
|
||||
"duration":5,
|
||||
"video_urls":["https://example.com/1.mp4","https://example.com/2.mp4","https://example.com/3.mp4","https://example.com/4.mp4"]
|
||||
}`)
|
||||
if _, err := BuildSeedance2Payload("doubao-seedance-2.0", input); err == nil {
|
||||
t.Fatal("expected too many reference videos to be rejected")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,29 @@
|
||||
package provider
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
)
|
||||
|
||||
type SubmitResult struct {
|
||||
TaskID string
|
||||
URL string
|
||||
Raw json.RawMessage
|
||||
}
|
||||
|
||||
type PollResult struct {
|
||||
Status string
|
||||
URL string
|
||||
Error string
|
||||
Raw json.RawMessage
|
||||
}
|
||||
|
||||
type HTTPError struct {
|
||||
StatusCode int
|
||||
Body string
|
||||
RetryAfter string
|
||||
}
|
||||
|
||||
func (e *HTTPError) Error() string {
|
||||
return fmt.Sprintf("中转站返回 HTTP %d: %s", e.StatusCode, e.Body)
|
||||
}
|
||||
@@ -0,0 +1,15 @@
|
||||
package queue
|
||||
|
||||
import (
|
||||
"juhe-factory/api/internal/config"
|
||||
|
||||
"github.com/hibiken/asynq"
|
||||
)
|
||||
|
||||
func NewClient(cfg config.Config) *asynq.Client {
|
||||
return asynq.NewClient(asynq.RedisClientOpt{
|
||||
Addr: cfg.RedisAddress,
|
||||
Password: cfg.RedisPassword,
|
||||
DB: cfg.AsynqRedisDB,
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,17 @@
|
||||
package queue
|
||||
|
||||
import (
|
||||
"juhe-factory/api/internal/config"
|
||||
|
||||
"github.com/hibiken/asynq"
|
||||
)
|
||||
|
||||
func NewServer(cfg config.Config) *asynq.Server {
|
||||
concurrency := cfg.AIWorkerConcurrency
|
||||
if concurrency < 1 {
|
||||
concurrency = 1
|
||||
}
|
||||
return asynq.NewServer(asynq.RedisClientOpt{
|
||||
Addr: cfg.RedisAddress, Password: cfg.RedisPassword, DB: cfg.AsynqRedisDB,
|
||||
}, asynq.Config{Concurrency: concurrency, Queues: map[string]int{"ai": 10, "default": 1}})
|
||||
}
|
||||
@@ -0,0 +1,48 @@
|
||||
package queue
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/hibiken/asynq"
|
||||
)
|
||||
|
||||
const (
|
||||
TypeDispatchChannel = "ai:dispatch-channel"
|
||||
TypeSubmitTask = "ai:submit-task"
|
||||
TypePollTask = "ai:poll-task"
|
||||
TypeDownloadTask = "ai:download-task"
|
||||
TypeAnalyzeEpisode = "media:analyze-episode"
|
||||
TypeParseDramaEpisode = "drama:parse-episode"
|
||||
TypeAnalyzeScript = "script:analyze"
|
||||
)
|
||||
|
||||
type IDPayload struct {
|
||||
ID uuid.UUID `json:"id"`
|
||||
}
|
||||
|
||||
func NewIDTask(taskType string, id uuid.UUID) (*asynq.Task, error) {
|
||||
payload, err := json.Marshal(IDPayload{ID: id})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return asynq.NewTask(taskType, payload), nil
|
||||
}
|
||||
|
||||
func EnqueueID(client *asynq.Client, taskType string, id uuid.UUID, delay time.Duration) error {
|
||||
task, err := NewIDTask(taskType, id)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
maxRetry := 0
|
||||
if taskType == TypeDispatchChannel || taskType == TypePollTask || taskType == TypeDownloadTask {
|
||||
maxRetry = 5
|
||||
}
|
||||
options := []asynq.Option{asynq.MaxRetry(maxRetry), asynq.Queue("ai")}
|
||||
if delay > 0 {
|
||||
options = append(options, asynq.ProcessIn(delay))
|
||||
}
|
||||
_, err = client.Enqueue(task, options...)
|
||||
return err
|
||||
}
|
||||
@@ -0,0 +1,57 @@
|
||||
package security
|
||||
|
||||
import (
|
||||
"crypto/aes"
|
||||
"crypto/cipher"
|
||||
"crypto/rand"
|
||||
"encoding/base64"
|
||||
"errors"
|
||||
"io"
|
||||
)
|
||||
|
||||
type Encryptor struct{ key []byte }
|
||||
|
||||
func NewEncryptor(encodedKey string) (*Encryptor, error) {
|
||||
key, err := base64.StdEncoding.DecodeString(encodedKey)
|
||||
if err != nil || len(key) != 32 {
|
||||
return nil, errors.New("CONFIG_ENCRYPTION_KEY 必须是 32 字节随机值的 Base64 编码")
|
||||
}
|
||||
return &Encryptor{key: key}, nil
|
||||
}
|
||||
|
||||
func (e *Encryptor) Encrypt(plain string) (string, error) {
|
||||
block, err := aes.NewCipher(e.key)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
gcm, err := cipher.NewGCM(block)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
nonce := make([]byte, gcm.NonceSize())
|
||||
if _, err = io.ReadFull(rand.Reader, nonce); err != nil {
|
||||
return "", err
|
||||
}
|
||||
sealed := gcm.Seal(nonce, nonce, []byte(plain), nil)
|
||||
return base64.RawStdEncoding.EncodeToString(sealed), nil
|
||||
}
|
||||
|
||||
func (e *Encryptor) Decrypt(encoded string) (string, error) {
|
||||
data, err := base64.RawStdEncoding.DecodeString(encoded)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
block, err := aes.NewCipher(e.key)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
gcm, err := cipher.NewGCM(block)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if len(data) < gcm.NonceSize() {
|
||||
return "", errors.New("密文无效")
|
||||
}
|
||||
plain, err := gcm.Open(nil, data[:gcm.NonceSize()], data[gcm.NonceSize():], nil)
|
||||
return string(plain), err
|
||||
}
|
||||
@@ -0,0 +1,55 @@
|
||||
package security
|
||||
|
||||
import (
|
||||
"crypto/rand"
|
||||
"crypto/subtle"
|
||||
"encoding/base64"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"golang.org/x/crypto/argon2"
|
||||
)
|
||||
|
||||
type PasswordHasher struct {
|
||||
Time uint32
|
||||
Memory uint32
|
||||
Parallelism uint8
|
||||
HashLength uint32
|
||||
SaltLength uint32
|
||||
}
|
||||
|
||||
func (h PasswordHasher) Hash(password string) (string, error) {
|
||||
if len(password) < 8 || len(password) > 128 {
|
||||
return "", errors.New("密码长度必须为 8 至 128 个字符")
|
||||
}
|
||||
salt := make([]byte, h.SaltLength)
|
||||
if _, err := rand.Read(salt); err != nil {
|
||||
return "", err
|
||||
}
|
||||
hash := argon2.IDKey([]byte(password), salt, h.Time, h.Memory, h.Parallelism, h.HashLength)
|
||||
return fmt.Sprintf("$argon2id$v=19$m=%d,t=%d,p=%d$%s$%s", h.Memory, h.Time, h.Parallelism,
|
||||
base64.RawStdEncoding.EncodeToString(salt), base64.RawStdEncoding.EncodeToString(hash)), nil
|
||||
}
|
||||
|
||||
func (h PasswordHasher) Verify(encoded, password string) bool {
|
||||
parts := strings.Split(encoded, "$")
|
||||
if len(parts) != 6 || parts[1] != "argon2id" {
|
||||
return false
|
||||
}
|
||||
var memory, iterations uint32
|
||||
var parallelism uint8
|
||||
if _, err := fmt.Sscanf(parts[3], "m=%d,t=%d,p=%d", &memory, &iterations, ¶llelism); err != nil {
|
||||
return false
|
||||
}
|
||||
salt, err := base64.RawStdEncoding.DecodeString(parts[4])
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
want, err := base64.RawStdEncoding.DecodeString(parts[5])
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
got := argon2.IDKey([]byte(password), salt, iterations, memory, parallelism, uint32(len(want)))
|
||||
return subtle.ConstantTimeCompare(got, want) == 1
|
||||
}
|
||||
@@ -0,0 +1,63 @@
|
||||
package security
|
||||
|
||||
import (
|
||||
"crypto/rand"
|
||||
"crypto/sha256"
|
||||
"encoding/base64"
|
||||
"encoding/hex"
|
||||
"errors"
|
||||
"time"
|
||||
|
||||
"github.com/golang-jwt/jwt/v5"
|
||||
)
|
||||
|
||||
type AdminClaims struct {
|
||||
AdminID string `json:"admin_id"`
|
||||
Username string `json:"username"`
|
||||
TokenType string `json:"token_type"`
|
||||
jwt.RegisteredClaims
|
||||
}
|
||||
|
||||
type TokenService struct {
|
||||
Secret []byte
|
||||
AccessTTL time.Duration
|
||||
RefreshTTL time.Duration
|
||||
}
|
||||
|
||||
func (s TokenService) NewAccessToken(adminID, username string) (string, time.Time, error) {
|
||||
now := time.Now()
|
||||
expires := now.Add(s.AccessTTL)
|
||||
claims := AdminClaims{AdminID: adminID, Username: username, TokenType: "access", RegisteredClaims: jwt.RegisteredClaims{Subject: adminID, IssuedAt: jwt.NewNumericDate(now), ExpiresAt: jwt.NewNumericDate(expires)}}
|
||||
token, err := jwt.NewWithClaims(jwt.SigningMethodHS256, claims).SignedString(s.Secret)
|
||||
return token, expires, err
|
||||
}
|
||||
|
||||
func (s TokenService) ParseAccessToken(raw string) (*AdminClaims, error) {
|
||||
claims := new(AdminClaims)
|
||||
token, err := jwt.ParseWithClaims(raw, claims, func(token *jwt.Token) (any, error) {
|
||||
if token.Method != jwt.SigningMethodHS256 {
|
||||
return nil, errors.New("不支持的签名算法")
|
||||
}
|
||||
return s.Secret, nil
|
||||
})
|
||||
if err != nil || !token.Valid || claims.TokenType != "access" {
|
||||
return nil, errors.New("登录凭证无效或已过期")
|
||||
}
|
||||
return claims, nil
|
||||
}
|
||||
|
||||
func (s TokenService) NewRefreshToken() (plain, hash string, expires time.Time, err error) {
|
||||
data := make([]byte, 32)
|
||||
if _, err = rand.Read(data); err != nil {
|
||||
return
|
||||
}
|
||||
plain = base64.RawURLEncoding.EncodeToString(data)
|
||||
hash = HashToken(plain)
|
||||
expires = time.Now().Add(s.RefreshTTL)
|
||||
return
|
||||
}
|
||||
|
||||
func HashToken(token string) string {
|
||||
sum := sha256.Sum256([]byte(token))
|
||||
return hex.EncodeToString(sum[:])
|
||||
}
|
||||
@@ -0,0 +1,47 @@
|
||||
package security
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"time"
|
||||
|
||||
"github.com/golang-jwt/jwt/v5"
|
||||
)
|
||||
|
||||
type WebClaims struct {
|
||||
UserID string `json:"user_id"`
|
||||
UID string `json:"uid"`
|
||||
Username string `json:"username"`
|
||||
SessionVersion int64 `json:"session_version"`
|
||||
TokenType string `json:"token_type"`
|
||||
jwt.RegisteredClaims
|
||||
}
|
||||
|
||||
type WebTokenService struct {
|
||||
Secret []byte
|
||||
AccessTTL time.Duration
|
||||
}
|
||||
|
||||
func (s WebTokenService) NewAccessToken(userID, uid, username string, sessionVersion int64) (string, time.Time, error) {
|
||||
now := time.Now()
|
||||
expires := now.Add(s.AccessTTL)
|
||||
claims := WebClaims{
|
||||
UserID: userID, UID: uid, Username: username, SessionVersion: sessionVersion, TokenType: "web_access",
|
||||
RegisteredClaims: jwt.RegisteredClaims{Subject: userID, IssuedAt: jwt.NewNumericDate(now), ExpiresAt: jwt.NewNumericDate(expires)},
|
||||
}
|
||||
token, err := jwt.NewWithClaims(jwt.SigningMethodHS256, claims).SignedString(s.Secret)
|
||||
return token, expires, err
|
||||
}
|
||||
|
||||
func (s WebTokenService) ParseAccessToken(raw string) (*WebClaims, error) {
|
||||
claims := new(WebClaims)
|
||||
token, err := jwt.ParseWithClaims(raw, claims, func(token *jwt.Token) (any, error) {
|
||||
if token.Method != jwt.SigningMethodHS256 {
|
||||
return nil, errors.New("不支持的签名算法")
|
||||
}
|
||||
return s.Secret, nil
|
||||
})
|
||||
if err != nil || !token.Valid || claims.TokenType != "web_access" {
|
||||
return nil, errors.New("登录凭证无效或已过期")
|
||||
}
|
||||
return claims, nil
|
||||
}
|
||||
@@ -0,0 +1,271 @@
|
||||
// API 服务组装模块,负责中间件、路由、静态资源和 HTTP 服务参数配置。
|
||||
package server
|
||||
|
||||
import (
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"juhe-factory/api/internal/config"
|
||||
"juhe-factory/api/internal/handler"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
func New(cfg config.Config, health *handler.Health, adminAuth *handler.AdminAuth, adminData *handler.AdminData, web *handler.Web, creative *handler.Creative, productImage *handler.ProductImage) *http.Server {
|
||||
if !cfg.Debug {
|
||||
gin.SetMode(gin.ReleaseMode)
|
||||
}
|
||||
router := gin.New()
|
||||
router.Use(traceID())
|
||||
router.Use(requestLogger())
|
||||
router.Use(gin.Recovery(), cors(cfg.CORSOrigins))
|
||||
|
||||
registerRoutes(router, health, adminAuth, adminData, web, creative, productImage)
|
||||
registerStatic(router, cfg.StaticDirectory)
|
||||
|
||||
return &http.Server{
|
||||
Addr: cfg.ServerAddress,
|
||||
Handler: router,
|
||||
ReadHeaderTimeout: 5 * time.Second,
|
||||
ReadTimeout: 30 * time.Second,
|
||||
WriteTimeout: 5 * time.Minute,
|
||||
IdleTimeout: 60 * time.Second,
|
||||
}
|
||||
}
|
||||
|
||||
// requestLogger 记录 HTTP 请求及其追踪 ID,便于通过审计记录关联服务器日志。
|
||||
// 正常响应(2xx)不输出,仅在出错(4xx/5xx)时记录,避免前端轮询请求刷屏。
|
||||
func requestLogger() gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
startedAt := time.Now()
|
||||
c.Next()
|
||||
if c.Writer.Status() < 400 {
|
||||
return
|
||||
}
|
||||
slog.InfoContext(c.Request.Context(), "http request",
|
||||
"trace_id", c.GetString("trace_id"),
|
||||
"method", c.Request.Method,
|
||||
"path", c.Request.URL.Path,
|
||||
"status", c.Writer.Status(),
|
||||
"latency_ms", time.Since(startedAt).Milliseconds(),
|
||||
"response_bytes", c.Writer.Size(),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
func registerRoutes(router *gin.Engine, health *handler.Health, auth *handler.AdminAuth, data *handler.AdminData, web *handler.Web, creative *handler.Creative, productImage *handler.ProductImage) {
|
||||
router.GET("/health", health.Check)
|
||||
webRoutes := router.Group("/api/web")
|
||||
webRoutes.GET("/health", health.Check)
|
||||
webRoutes.POST("/auth/login", web.Login)
|
||||
webRoutes.POST("/auth/refresh", web.Refresh)
|
||||
webRoutes.POST("/auth/logout", web.Logout)
|
||||
webProtected := webRoutes.Group("")
|
||||
webProtected.Use(web.Middleware())
|
||||
webProtected.GET("/account", web.Account)
|
||||
webProtected.PATCH("/account/profile", web.UpdateProfile)
|
||||
webProtected.POST("/account/avatar", web.UploadAvatar)
|
||||
webProtected.PUT("/account/password", web.ChangePassword)
|
||||
webProtected.GET("/usage/30-days", web.Usage30Days)
|
||||
webProtected.GET("/usage/consumption-records", web.ConsumptionRecords)
|
||||
webProtected.POST("/redemption/redeem", web.Redeem)
|
||||
webProtected.GET("/user/prompts/:prompt_type", creative.ListUserPrompts)
|
||||
webProtected.PUT("/user/prompts/:prompt_type/select", creative.SelectUserPrompt)
|
||||
webProtected.POST("/user/custom-prompts", creative.CreateCustomPrompt)
|
||||
webProtected.PUT("/user/custom-prompts/:id", creative.UpdateCustomPrompt)
|
||||
webProtected.DELETE("/user/custom-prompts/:id", creative.DeleteCustomPrompt)
|
||||
webProtected.GET("/creative/options", creative.ListOptions)
|
||||
webProtected.POST("/product-images/generations", productImage.Generate)
|
||||
webProtected.GET("/product-images/generations", productImage.ListGenerations)
|
||||
webProtected.DELETE("/product-images/generations/:task_id", productImage.DeleteGeneration)
|
||||
webProtected.GET("/user/model-configs/:scope", creative.ListUserModelPreferences)
|
||||
webProtected.PUT("/user/model-configs/:scope/:model_type", creative.SaveUserModelPreference)
|
||||
webProtected.GET("/script-analyses", creative.ListScriptAnalyses)
|
||||
webProtected.POST("/script-analyses", creative.CreateScriptAnalysis)
|
||||
webProtected.GET("/script-analyses/import-projects", creative.ListScriptAnalysisImportProjects)
|
||||
webProtected.GET("/script-analyses/:script_id", creative.GetScriptAnalysis)
|
||||
webProtected.PUT("/script-analyses/:script_id", creative.UpdateScriptAnalysis)
|
||||
webProtected.POST("/script-analyses/:script_id/import-project", creative.ImportScriptAnalysisProject)
|
||||
webProtected.POST("/script-analyses/:script_id/import-file", creative.ImportScriptAnalysisFile)
|
||||
webProtected.POST("/script-analyses/:script_id/analyze", creative.QueueScriptAnalysis)
|
||||
webProtected.POST("/script-analyses/:script_id/cancel", creative.CancelScriptAnalysis)
|
||||
webProtected.DELETE("/script-analyses/:script_id", creative.DeleteScriptAnalysis)
|
||||
webProtected.GET("/media/:media_id", creative.GetMedia)
|
||||
webProtected.GET("/projects", creative.ListProjects)
|
||||
webProtected.POST("/projects", creative.CreateProject)
|
||||
webProtected.GET("/projects/:project_id", creative.GetProject)
|
||||
webProtected.PUT("/projects/:project_id", creative.UpdateProject)
|
||||
webProtected.DELETE("/projects/:project_id", creative.DeleteProject)
|
||||
webProtected.PUT("/projects/:project_id/model-configs", creative.SaveModelConfigs)
|
||||
webProtected.POST("/projects/:project_id/episodes", creative.CreateEpisode)
|
||||
webProtected.PUT("/projects/:project_id/episodes/:episode_id", creative.UpdateEpisodeName)
|
||||
webProtected.POST("/projects/:project_id/episodes/import/preview", creative.PreviewDramaImport)
|
||||
webProtected.POST("/projects/:project_id/episodes/import/confirm", creative.ConfirmDramaImport)
|
||||
webProtected.DELETE("/projects/:project_id/episodes/:episode_id", creative.DeleteEpisode)
|
||||
webProtected.GET("/projects/:project_id/episodes/:episode_id/text-source", creative.GetDramaEpisodeSource)
|
||||
webProtected.PUT("/projects/:project_id/episodes/:episode_id/text-source", creative.SaveDramaEpisodeSource)
|
||||
webProtected.POST("/projects/:project_id/episodes/:episode_id/parse", creative.QueueDramaParse)
|
||||
webProtected.POST("/projects/:project_id/episodes/:episode_id/reparse", creative.ReparseDrama)
|
||||
webProtected.POST("/projects/:project_id/episodes/:episode_id/storyboards", creative.CreateDramaStoryboard)
|
||||
webProtected.GET("/projects/:project_id/parse-tasks", creative.ListDramaParseTasks)
|
||||
webProtected.POST("/projects/:project_id/parse-tasks/:task_id/cancel", creative.CancelDramaParse)
|
||||
webProtected.GET("/projects/:project_id/episodes/:episode_id/workbench", creative.GetWorkbench)
|
||||
webProtected.GET("/projects/:project_id/episodes/:episode_id/videos.zip", creative.DownloadEpisodeVideos)
|
||||
webProtected.GET("/projects/:project_id/episodes/:episode_id/redraw-workbench", creative.GetRedrawEpisodeWorkbench)
|
||||
webProtected.POST("/projects/:project_id/episodes/:episode_id/redraw-script", creative.QueueRedrawEpisodeScript)
|
||||
webProtected.PUT("/projects/:project_id/episodes/:episode_id/redraw-script", creative.SaveRedrawEpisodeScript)
|
||||
webProtected.GET("/projects/:project_id/redraw-workbench", creative.GetRedrawWorkbench)
|
||||
webProtected.POST("/projects/:project_id/redraw-source", creative.UploadRedrawSource)
|
||||
webProtected.POST("/projects/:project_id/redraw-subtitle", creative.UploadRedrawSubtitle)
|
||||
webProtected.POST("/projects/:project_id/redraw-analyze", creative.QueueRedrawAnalysis)
|
||||
webProtected.POST("/projects/:project_id/redraw-reanalyze", creative.ReanalyzeRedraw)
|
||||
webProtected.POST("/projects/:project_id/redraw-script", creative.QueueRedrawScript)
|
||||
webProtected.PUT("/projects/:project_id/redraw-analysis-settings", creative.UpdateRedrawAnalysisSettings)
|
||||
webProtected.POST("/projects/:project_id/episodes/:episode_id/source", creative.UploadEpisodeSource)
|
||||
webProtected.POST("/projects/:project_id/episodes/:episode_id/subtitle", creative.UploadEpisodeSubtitle)
|
||||
webProtected.POST("/projects/:project_id/episodes/:episode_id/analyze", creative.QueueAnalysis)
|
||||
webProtected.POST("/projects/:project_id/episodes/:episode_id/reanalyze", creative.Reanalyze)
|
||||
webProtected.PUT("/projects/:project_id/episodes/:episode_id/analysis-settings", creative.UpdateEpisodeAnalysisSettings)
|
||||
webProtected.POST("/projects/:project_id/assets", creative.CreateAsset)
|
||||
webProtected.PUT("/projects/:project_id/assets/:asset_id", creative.UpdateAsset)
|
||||
webProtected.DELETE("/projects/:project_id/assets/:asset_id", creative.DeleteAsset)
|
||||
webProtected.POST("/projects/:project_id/assets/:asset_id/image", creative.UploadAssetImage)
|
||||
webProtected.DELETE("/projects/:project_id/assets/:asset_id/image", creative.RemoveAssetImage)
|
||||
webProtected.POST("/projects/:project_id/assets/:asset_id/audio", creative.UploadAssetAudio)
|
||||
webProtected.PUT("/projects/:project_id/storyboards/:storyboard_id", creative.UpdateStoryboard)
|
||||
webProtected.POST("/projects/:project_id/storyboards/:storyboard_id/insert", creative.InsertStoryboard)
|
||||
webProtected.DELETE("/projects/:project_id/storyboards/:storyboard_id", creative.DeleteStoryboard)
|
||||
webProtected.POST("/projects/:project_id/episodes/:episode_id/storyboards/:storyboard_id/reinfer", creative.QueueStoryboardAnalysis)
|
||||
webProtected.POST("/projects/:project_id/episodes/:episode_id/storyboards/:storyboard_id/generate", creative.QueueGeneration)
|
||||
webProtected.POST("/projects/:project_id/storyboards/:storyboard_id/redraw-reinfer", creative.QueueRedrawStoryboardAnalysis)
|
||||
webProtected.GET("/projects/:project_id/tasks", creative.ListTasks)
|
||||
webProtected.POST("/projects/:project_id/tasks/:task_id/cancel", creative.CancelTask)
|
||||
webProtected.GET("/projects/:project_id/storyboards/:storyboard_id/outputs", creative.ListStoryboardOutputs)
|
||||
webProtected.DELETE("/projects/:project_id/storyboards/:storyboard_id/image", creative.RemoveStoryboardImage)
|
||||
webProtected.PUT("/projects/:project_id/storyboards/:storyboard_id/outputs/:output_id/activate", creative.ActivateStoryboardOutput)
|
||||
webProtected.PUT("/projects/:project_id/storyboards/:storyboard_id/outputs/:output_id/candidate", creative.AddStoryboardCandidate)
|
||||
webProtected.DELETE("/projects/:project_id/storyboards/:storyboard_id/outputs/:output_id/candidate", creative.RemoveStoryboardCandidate)
|
||||
webProtected.DELETE("/projects/:project_id/storyboards/:storyboard_id/outputs/:output_id", creative.DeleteStoryboardOutput)
|
||||
webProtected.GET("/projects/:project_id/assets/:asset_id/outputs", creative.ListAssetOutputs)
|
||||
webProtected.PUT("/projects/:project_id/assets/:asset_id/outputs/:output_id/activate", creative.ActivateAssetOutput)
|
||||
webProtected.DELETE("/projects/:project_id/assets/:asset_id/outputs/:output_id", creative.DeleteAssetOutput)
|
||||
admin := router.Group("/api/admin")
|
||||
admin.GET("/health", health.Check)
|
||||
admin.POST("/auth/login", auth.Login)
|
||||
admin.POST("/auth/refresh", auth.Refresh)
|
||||
admin.POST("/auth/logout", auth.Logout)
|
||||
protected := admin.Group("")
|
||||
protected.Use(auth.Middleware())
|
||||
protected.PUT("/auth/password", auth.ChangePassword)
|
||||
protected.GET("/audit-logs", data.ListAuditLogs)
|
||||
protected.GET("/users", data.ListUsers)
|
||||
protected.POST("/users", data.CreateUser)
|
||||
protected.POST("/users/batch", data.BatchCreateUsers)
|
||||
protected.PATCH("/users/batch", data.BatchUpdateUsers)
|
||||
protected.DELETE("/users/batch", data.BatchDeleteUsers)
|
||||
protected.POST("/users/:id/points", data.GrantUserPoints)
|
||||
protected.GET("/models", data.ListModels)
|
||||
protected.POST("/models", data.SaveModel)
|
||||
protected.PUT("/models/:id", data.SaveModel)
|
||||
protected.GET("/prompts", data.ListPrompts)
|
||||
protected.POST("/prompts", data.SavePromptSimple)
|
||||
protected.PUT("/prompts/:id", data.SavePromptSimple)
|
||||
protected.DELETE("/prompts/:id", data.DeletePromptSimple)
|
||||
protected.POST("/prompts/:id/actions", data.PromptActionSimple)
|
||||
protected.GET("/prompts/:id/history", data.PromptHistorySimple)
|
||||
protected.GET("/redemption-codes", data.ListRedemptions)
|
||||
protected.POST("/redemption-batches", data.CreateRedemptionBatch)
|
||||
protected.POST("/uploads/style-images", data.UploadStyleImage)
|
||||
protected.PUT("/resources/styles/reorder", data.ReorderStyles)
|
||||
resources := protected.Group("/resources/:resource")
|
||||
resources.Use(allowAdminResource("styles", "channels", "models"))
|
||||
resources.GET("", data.ListResource)
|
||||
resources.POST("", data.CreateResource)
|
||||
resources.PUT("/:id", data.UpdateResource)
|
||||
resources.PATCH("/:id/enabled", data.ToggleResource)
|
||||
resources.DELETE("/:id", data.DeleteResource)
|
||||
}
|
||||
|
||||
func allowAdminResource(resources ...string) gin.HandlerFunc {
|
||||
allowed := make(map[string]struct{}, len(resources))
|
||||
for _, resource := range resources {
|
||||
allowed[resource] = struct{}{}
|
||||
}
|
||||
return func(c *gin.Context) {
|
||||
if _, ok := allowed[c.Param("resource")]; !ok {
|
||||
c.AbortWithStatusJSON(http.StatusNotFound, gin.H{"code": "route_not_found", "message": "接口不存在"})
|
||||
return
|
||||
}
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
|
||||
func traceID() gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
id := strings.TrimSpace(c.GetHeader("X-Request-ID"))
|
||||
if id == "" {
|
||||
id = uuid.NewString()
|
||||
}
|
||||
c.Set("trace_id", id)
|
||||
c.Header("X-Request-ID", id)
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
|
||||
func registerStatic(router *gin.Engine, root string) {
|
||||
adminRoot := filepath.Join(root, "admin")
|
||||
webRoot := filepath.Join(root, "web")
|
||||
router.Static("/admin/assets", filepath.Join(adminRoot, "assets"))
|
||||
router.Static("/assets", filepath.Join(webRoot, "assets"))
|
||||
router.GET("/admin", func(c *gin.Context) { c.File(filepath.Join(adminRoot, "index.html")) })
|
||||
router.GET("/", func(c *gin.Context) { c.File(filepath.Join(webRoot, "index.html")) })
|
||||
|
||||
router.NoRoute(func(c *gin.Context) {
|
||||
path := c.Request.URL.Path
|
||||
if strings.HasPrefix(path, "/api/") {
|
||||
c.JSON(http.StatusNotFound, gin.H{"code": "route_not_found", "message": "接口不存在"})
|
||||
return
|
||||
}
|
||||
if strings.HasPrefix(path, "/admin") {
|
||||
c.File(filepath.Join(adminRoot, "index.html"))
|
||||
return
|
||||
}
|
||||
|
||||
relativePath := filepath.Clean(strings.TrimPrefix(path, "/"))
|
||||
if filepath.IsLocal(relativePath) {
|
||||
candidate := filepath.Join(webRoot, relativePath)
|
||||
if info, err := os.Stat(candidate); err == nil && !info.IsDir() {
|
||||
c.File(candidate)
|
||||
return
|
||||
}
|
||||
}
|
||||
c.File(filepath.Join(webRoot, "index.html"))
|
||||
})
|
||||
}
|
||||
|
||||
func cors(allowedOrigins []string) gin.HandlerFunc {
|
||||
allowed := make(map[string]struct{}, len(allowedOrigins))
|
||||
for _, origin := range allowedOrigins {
|
||||
allowed[origin] = struct{}{}
|
||||
}
|
||||
return func(c *gin.Context) {
|
||||
origin := c.GetHeader("Origin")
|
||||
if _, ok := allowed[origin]; ok {
|
||||
c.Header("Access-Control-Allow-Origin", origin)
|
||||
c.Header("Access-Control-Allow-Credentials", "true")
|
||||
c.Header("Vary", "Origin")
|
||||
}
|
||||
c.Header("Access-Control-Allow-Headers", "Content-Type, Authorization")
|
||||
c.Header("Access-Control-Allow-Methods", "GET, POST, PUT, PATCH, DELETE, OPTIONS")
|
||||
if c.Request.Method == http.MethodOptions {
|
||||
c.AbortWithStatus(http.StatusNoContent)
|
||||
return
|
||||
}
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,42 @@
|
||||
// 管理端路由测试,验证通用资源接口仅放行明确支持的资源类型。
|
||||
package server
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// TestAllowAdminResource 验证模型资源可进入处理器,未知资源仍返回接口不存在。
|
||||
func TestAllowAdminResource(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
tests := []struct {
|
||||
name string
|
||||
resource string
|
||||
statusCode int
|
||||
}{
|
||||
{name: "允许模型资源", resource: "models", statusCode: http.StatusNoContent},
|
||||
{name: "拒绝未知资源", resource: "unknown", statusCode: http.StatusNotFound},
|
||||
}
|
||||
|
||||
for _, test := range tests {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
router := gin.New()
|
||||
route := router.Group("/resources/:resource")
|
||||
route.Use(allowAdminResource("styles", "channels", "models"))
|
||||
route.PATCH("/:id/enabled", func(c *gin.Context) {
|
||||
c.Status(http.StatusNoContent)
|
||||
})
|
||||
|
||||
response := httptest.NewRecorder()
|
||||
request := httptest.NewRequest(http.MethodPatch, "/resources/"+test.resource+"/model-id/enabled", nil)
|
||||
router.ServeHTTP(response, request)
|
||||
|
||||
if response.Code != test.statusCode {
|
||||
t.Fatalf("unexpected status: got %d, want %d", response.Code, test.statusCode)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,575 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/rand"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"errors"
|
||||
"fmt"
|
||||
"math"
|
||||
"math/big"
|
||||
"net/http"
|
||||
"regexp"
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"juhe-factory/api/internal/security"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
type AdminData struct {
|
||||
DB *gorm.DB
|
||||
Passwords security.PasswordHasher
|
||||
Encryptor *security.Encryptor
|
||||
HTTPClient *http.Client
|
||||
EncryptionKeyVersion string
|
||||
balanceSyncing atomic.Bool
|
||||
}
|
||||
|
||||
type Page struct {
|
||||
Items []map[string]any `json:"items"`
|
||||
Total int64 `json:"total"`
|
||||
Page int `json:"page"`
|
||||
PageSize int `json:"page_size"`
|
||||
}
|
||||
|
||||
type ResourceSpec struct {
|
||||
Table string
|
||||
Select string
|
||||
SearchColumns []string
|
||||
Fields map[string]string
|
||||
Filters map[string]string
|
||||
Order string
|
||||
SoftDelete bool
|
||||
}
|
||||
|
||||
var resourceSpecs = map[string]ResourceSpec{
|
||||
"styles": {Table: "project_styles", Select: "id,name,image_key,image_url,image_mime,image_size,image_width,image_height,sort_order", Fields: map[string]string{"name": "name", "image_key": "image_key", "image_url": "image_url", "image_mime": "image_mime", "image_size": "image_size", "image_width": "image_width", "image_height": "image_height"}, Order: "sort_order,name", SoftDelete: true},
|
||||
"channels": {Table: "channels", Select: "c.id,c.name,c.channel_type,c.base_url,c.api_key_ciphertext,c.api_key_last4,c.warning_threshold,c.channel_points_per_cny,c.max_concurrency,c.max_user_concurrency,c.enabled,c.created_at,c.updated_at,b.balance,b.currency,b.synced_at", SearchColumns: []string{"c.name", "c.base_url"}, Fields: map[string]string{"name": "name", "channel_type": "channel_type", "base_url": "base_url", "warning_threshold": "warning_threshold", "channel_points_per_cny": "channel_points_per_cny", "max_concurrency": "max_concurrency", "max_user_concurrency": "max_user_concurrency", "enabled": "enabled"}, Filters: map[string]string{"channel_type": "channel_type", "enabled": "enabled"}, Order: "c.created_at DESC", SoftDelete: true},
|
||||
"models": {Table: "models", Select: "id", Fields: map[string]string{}, Order: "created_at DESC", SoftDelete: true},
|
||||
"prompts": {Table: "prompts", Select: "id", Fields: map[string]string{}, Order: "updated_at DESC", SoftDelete: true},
|
||||
}
|
||||
|
||||
func pageArgs(pageRaw, sizeRaw string) (int, int) {
|
||||
page, _ := strconv.Atoi(pageRaw)
|
||||
if page < 1 {
|
||||
page = 1
|
||||
}
|
||||
size, _ := strconv.Atoi(sizeRaw)
|
||||
if size < 1 {
|
||||
size = 20
|
||||
}
|
||||
if size > 100 {
|
||||
size = 100
|
||||
}
|
||||
return page, size
|
||||
}
|
||||
|
||||
func (s *AdminData) ListResource(ctx context.Context, resource, keyword string, filters map[string]string, pageRaw, sizeRaw string) (Page, error) {
|
||||
spec, ok := resourceSpecs[resource]
|
||||
if !ok {
|
||||
return Page{}, errors.New("不支持的资源类型")
|
||||
}
|
||||
page, size := pageArgs(pageRaw, sizeRaw)
|
||||
if resource == "channels" {
|
||||
s.triggerChannelBalanceSync()
|
||||
}
|
||||
query := s.DB.Table(spec.Table)
|
||||
if resource == "channels" {
|
||||
query = s.DB.Table("channels c").Joins("LEFT JOIN LATERAL (SELECT balance,currency,synced_at FROM channel_balance_snapshots WHERE channel_id=c.id ORDER BY synced_at DESC LIMIT 1) b ON true")
|
||||
}
|
||||
if spec.SoftDelete {
|
||||
prefix := strings.Split(spec.Table, " ")[0]
|
||||
if resource == "channels" {
|
||||
prefix = "c"
|
||||
}
|
||||
query = query.Where(prefix + ".deleted_at IS NULL")
|
||||
}
|
||||
if keyword = strings.TrimSpace(keyword); keyword != "" && len(spec.SearchColumns) > 0 {
|
||||
parts := make([]string, len(spec.SearchColumns))
|
||||
args := make([]any, len(spec.SearchColumns))
|
||||
for i, column := range spec.SearchColumns {
|
||||
parts[i] = column + "::text ILIKE ?"
|
||||
args[i] = "%" + keyword + "%"
|
||||
}
|
||||
query = query.Where("("+strings.Join(parts, " OR ")+")", args...)
|
||||
}
|
||||
for key, value := range filters {
|
||||
if column, exists := spec.Filters[key]; exists && value != "" {
|
||||
query = query.Where(column+" = ?", value)
|
||||
}
|
||||
}
|
||||
var total int64
|
||||
if err := query.Count(&total).Error; err != nil {
|
||||
return Page{}, err
|
||||
}
|
||||
items := make([]map[string]any, 0)
|
||||
if err := query.Select(spec.Select).Order(spec.Order).Offset((page - 1) * size).Limit(size).Find(&items).Error; err != nil {
|
||||
return Page{}, err
|
||||
}
|
||||
if resource == "channels" {
|
||||
s.maskChannelAPIKeys(items)
|
||||
}
|
||||
return Page{Items: items, Total: total, Page: page, PageSize: size}, nil
|
||||
}
|
||||
|
||||
func (s *AdminData) maskChannelAPIKeys(items []map[string]any) {
|
||||
for _, item := range items {
|
||||
ciphertext := strings.TrimSpace(fmt.Sprint(item["api_key_ciphertext"]))
|
||||
last4 := strings.TrimSpace(fmt.Sprint(item["api_key_last4"]))
|
||||
delete(item, "api_key_ciphertext")
|
||||
delete(item, "api_key_last4")
|
||||
item["api_key_masked"] = ""
|
||||
if ciphertext != "" && ciphertext != "<nil>" && s.Encryptor != nil {
|
||||
if plain, err := s.Encryptor.Decrypt(ciphertext); err == nil {
|
||||
item["api_key_masked"] = maskAPIKey(plain)
|
||||
continue
|
||||
}
|
||||
}
|
||||
if last4 != "" && last4 != "<nil>" {
|
||||
item["api_key_masked"] = "********" + last4
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func maskAPIKey(value string) string {
|
||||
characters := []rune(value)
|
||||
if len(characters) <= 10 {
|
||||
return strings.Repeat("*", max(len(characters), 8))
|
||||
}
|
||||
return string(characters[:5]) + "********" + string(characters[len(characters)-5:])
|
||||
}
|
||||
|
||||
func (s *AdminData) SaveResource(resource, id string, values map[string]any) (string, error) {
|
||||
spec, ok := resourceSpecs[resource]
|
||||
if !ok {
|
||||
return "", errors.New("不支持的资源类型")
|
||||
}
|
||||
data := map[string]any{}
|
||||
for input, column := range spec.Fields {
|
||||
if value, exists := values[input]; exists {
|
||||
data[column] = value
|
||||
}
|
||||
}
|
||||
if resource == "channels" {
|
||||
if err := s.prepareChannel(values, data); err != nil {
|
||||
return "", err
|
||||
}
|
||||
}
|
||||
if resource == "channels" && id == "" {
|
||||
apiKey, _ := values["api_key"].(string)
|
||||
if strings.TrimSpace(apiKey) == "" {
|
||||
return "", errors.New("API Key 不能为空")
|
||||
}
|
||||
data["channel_type"] = "relay"
|
||||
}
|
||||
if id == "" {
|
||||
newID := uuid.NewString()
|
||||
data["id"] = newID
|
||||
if resource == "styles" {
|
||||
var next int
|
||||
if err := s.DB.Table("project_styles").Select("coalesce(max(sort_order),-1)+1").Scan(&next).Error; err != nil {
|
||||
return "", err
|
||||
}
|
||||
data["sort_order"] = next
|
||||
}
|
||||
if err := s.DB.Table(spec.Table).Create(data).Error; err != nil {
|
||||
return "", err
|
||||
}
|
||||
return newID, nil
|
||||
}
|
||||
if len(data) == 0 {
|
||||
return id, nil
|
||||
}
|
||||
result := s.DB.Table(spec.Table).Where("id = ? AND deleted_at IS NULL", id).Updates(data)
|
||||
if result.Error != nil {
|
||||
return "", result.Error
|
||||
}
|
||||
if result.RowsAffected == 0 {
|
||||
return "", gorm.ErrRecordNotFound
|
||||
}
|
||||
return id, nil
|
||||
}
|
||||
|
||||
func (s *AdminData) ReorderStyles(ids []string) error {
|
||||
if len(ids) == 0 {
|
||||
return errors.New("风格排序不能为空")
|
||||
}
|
||||
seen := make(map[string]struct{}, len(ids))
|
||||
for _, id := range ids {
|
||||
if _, err := uuid.Parse(id); err != nil {
|
||||
return errors.New("风格 ID 无效")
|
||||
}
|
||||
if _, exists := seen[id]; exists {
|
||||
return errors.New("风格排序包含重复项")
|
||||
}
|
||||
seen[id] = struct{}{}
|
||||
}
|
||||
return s.DB.Transaction(func(tx *gorm.DB) error {
|
||||
var count int64
|
||||
if err := tx.Table("project_styles").Where("deleted_at IS NULL").Count(&count).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if int64(len(ids)) != count {
|
||||
return errors.New("风格列表已变化,请刷新后重试")
|
||||
}
|
||||
for position, id := range ids {
|
||||
result := tx.Table("project_styles").Where("id=? AND deleted_at IS NULL", id).Update("sort_order", position)
|
||||
if result.Error != nil {
|
||||
return result.Error
|
||||
}
|
||||
if result.RowsAffected == 0 {
|
||||
return errors.New("风格列表已变化,请刷新后重试")
|
||||
}
|
||||
}
|
||||
return nil
|
||||
})
|
||||
}
|
||||
|
||||
func (s *AdminData) prepareChannel(values, data map[string]any) error {
|
||||
rateText := strings.TrimSpace(fmt.Sprint(values["channel_points_per_cny"]))
|
||||
rate, err := strconv.ParseFloat(rateText, 64)
|
||||
if err != nil || math.IsNaN(rate) || math.IsInf(rate, 0) || rate <= 0 || rate > 1000000000000 {
|
||||
return errors.New("渠道汇率必须大于 0")
|
||||
}
|
||||
data["channel_points_per_cny"] = rateText
|
||||
maxConcurrency, err := positiveInt(values["max_concurrency"], 500)
|
||||
if err != nil || maxConcurrency > 5000 {
|
||||
return errors.New("最大并发数必须为 1 至 5000")
|
||||
}
|
||||
maxUserConcurrency, err := positiveInt(values["max_user_concurrency"], 10)
|
||||
if err != nil || maxUserConcurrency > 500 || maxUserConcurrency > maxConcurrency {
|
||||
return errors.New("单用户最高并发数必须为 1 至 500,且不能超过最大并发数")
|
||||
}
|
||||
data["max_concurrency"] = maxConcurrency
|
||||
data["max_user_concurrency"] = maxUserConcurrency
|
||||
plain, _ := values["api_key"].(string)
|
||||
plain = strings.TrimSpace(plain)
|
||||
if plain == "" {
|
||||
return nil
|
||||
}
|
||||
if s.Encryptor == nil {
|
||||
return errors.New("敏感配置加密未配置")
|
||||
}
|
||||
ciphertext, err := s.Encryptor.Encrypt(plain)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
last4 := plain
|
||||
if len(last4) > 4 {
|
||||
last4 = last4[len(last4)-4:]
|
||||
}
|
||||
data["api_key_ciphertext"] = ciphertext
|
||||
data["api_key_last4"] = last4
|
||||
data["encryption_key_version"] = s.EncryptionKeyVersion
|
||||
return nil
|
||||
}
|
||||
|
||||
func positiveInt(value any, fallback int) (int, error) {
|
||||
if value == nil || fmt.Sprint(value) == "" {
|
||||
return fallback, nil
|
||||
}
|
||||
parsed, err := strconv.Atoi(fmt.Sprint(value))
|
||||
if err != nil || parsed < 1 {
|
||||
return 0, errors.New("必须为正整数")
|
||||
}
|
||||
return parsed, nil
|
||||
}
|
||||
|
||||
func (s *AdminData) ToggleResource(resource, id string, enabled bool) error {
|
||||
spec, ok := resourceSpecs[resource]
|
||||
if !ok {
|
||||
return errors.New("不支持的资源类型")
|
||||
}
|
||||
if resource == "models" && enabled {
|
||||
var channelEnabled bool
|
||||
if err := s.DB.Raw(`SELECT c.enabled FROM models m JOIN channels c ON c.id=m.channel_id
|
||||
WHERE m.id=? AND m.deleted_at IS NULL AND c.deleted_at IS NULL`, id).Scan(&channelEnabled).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if !channelEnabled {
|
||||
return errors.New("所属渠道已禁用,无法启用该模型")
|
||||
}
|
||||
}
|
||||
if resource == "channels" {
|
||||
return s.DB.Transaction(func(tx *gorm.DB) error {
|
||||
result := tx.Table("channels").Where("id = ? AND deleted_at IS NULL", id).Update("enabled", enabled)
|
||||
if result.Error != nil {
|
||||
return result.Error
|
||||
}
|
||||
if result.RowsAffected == 0 {
|
||||
return gorm.ErrRecordNotFound
|
||||
}
|
||||
return tx.Table("models").Where("channel_id = ? AND deleted_at IS NULL", id).Update("enabled", enabled).Error
|
||||
})
|
||||
}
|
||||
result := s.DB.Table(spec.Table).Where("id = ? AND deleted_at IS NULL", id).Update("enabled", enabled)
|
||||
if result.Error != nil {
|
||||
return result.Error
|
||||
}
|
||||
if result.RowsAffected == 0 {
|
||||
return gorm.ErrRecordNotFound
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *AdminData) DeleteResource(resource, id string) error {
|
||||
spec, ok := resourceSpecs[resource]
|
||||
if !ok || !spec.SoftDelete {
|
||||
return errors.New("不支持删除该资源")
|
||||
}
|
||||
result := s.DB.Table(spec.Table).Where("id = ? AND deleted_at IS NULL", id).Update("deleted_at", time.Now())
|
||||
if result.Error != nil {
|
||||
return result.Error
|
||||
}
|
||||
if result.RowsAffected == 0 {
|
||||
return gorm.ErrRecordNotFound
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *AdminData) ListUsers(keyword, enabled, pageRaw, sizeRaw string) (Page, error) {
|
||||
page, size := pageArgs(pageRaw, sizeRaw)
|
||||
query := s.DB.Table("web_users u").Where("u.deleted_at IS NULL")
|
||||
if keyword = strings.TrimSpace(keyword); keyword != "" {
|
||||
query = query.Where("u.account::text ILIKE ? OR u.username::text ILIKE ? OR u.uid = ?", "%"+keyword+"%", "%"+keyword+"%", keyword)
|
||||
}
|
||||
if enabled != "" {
|
||||
query = query.Where("u.enabled = ?", enabled)
|
||||
}
|
||||
var total int64
|
||||
if err := query.Count(&total).Error; err != nil {
|
||||
return Page{}, err
|
||||
}
|
||||
items := make([]map[string]any, 0)
|
||||
err := query.Select(`u.id,u.account,u.username,u.uid,u.point_balance,u.daily_limit,u.enabled,u.last_online_at,u.created_at,
|
||||
coalesce((SELECT sum(-l.change_amount) FROM point_ledger l WHERE l.user_id=u.id AND l.change_amount<0 AND l.created_at>=CURRENT_DATE AND l.created_at<CURRENT_DATE+INTERVAL '1 day'),0) AS today_consumption`).
|
||||
Order("u.created_at DESC").Offset((page - 1) * size).Limit(size).Find(&items).Error
|
||||
return Page{Items: items, Total: total, Page: page, PageSize: size}, err
|
||||
}
|
||||
|
||||
func (s *AdminData) CreateUser(account, password string, dailyLimit any) (map[string]any, error) {
|
||||
account = strings.TrimSpace(account)
|
||||
if account == "" {
|
||||
return nil, errors.New("账号不能为空")
|
||||
}
|
||||
items, err := s.createUsers([]string{account}, password, dailyLimit)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return items[0], nil
|
||||
}
|
||||
|
||||
func (s *AdminData) BatchCreateUsers(prefix, startSequence, endSequence, password string, dailyLimit any) ([]map[string]any, error) {
|
||||
accounts, err := buildBatchAccounts(prefix, startSequence, endSequence)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return s.createUsers(accounts, password, dailyLimit)
|
||||
}
|
||||
|
||||
func (s *AdminData) createUsers(accounts []string, password string, dailyLimit any) ([]map[string]any, error) {
|
||||
hash, err := s.Passwords.Hash(password)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
created := make([]map[string]any, 0, len(accounts))
|
||||
err = s.DB.Transaction(func(tx *gorm.DB) error {
|
||||
var existing string
|
||||
if err := tx.Table("web_users").Select("account").Where("account IN ?", accounts).Limit(1).Scan(&existing).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if existing != "" {
|
||||
return fmt.Errorf("账号 %s 已存在", existing)
|
||||
}
|
||||
for _, account := range accounts {
|
||||
user, err := createUserWithAccount(tx, account, hash, dailyLimit)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
created = append(created, user)
|
||||
}
|
||||
return nil
|
||||
})
|
||||
return created, err
|
||||
}
|
||||
|
||||
const usernameAlphabet = "abcdefghijklmnopqrstuvwxyz0123456789"
|
||||
|
||||
var sequencePattern = regexp.MustCompile(`^[0-9]+$`)
|
||||
|
||||
func buildBatchAccounts(prefix, startSequence, endSequence string) ([]string, error) {
|
||||
prefix = strings.TrimSpace(prefix)
|
||||
startSequence = strings.TrimSpace(startSequence)
|
||||
endSequence = strings.TrimSpace(endSequence)
|
||||
if prefix == "" {
|
||||
return nil, errors.New("账号前缀不能为空")
|
||||
}
|
||||
if !sequencePattern.MatchString(startSequence) || !sequencePattern.MatchString(endSequence) {
|
||||
return nil, errors.New("起始序号和终止序号必须为数字")
|
||||
}
|
||||
start, err := strconv.Atoi(startSequence)
|
||||
if err != nil {
|
||||
return nil, errors.New("起始序号无效")
|
||||
}
|
||||
end, err := strconv.Atoi(endSequence)
|
||||
if err != nil {
|
||||
return nil, errors.New("终止序号无效")
|
||||
}
|
||||
if start > end {
|
||||
return nil, errors.New("起始序号不能大于终止序号")
|
||||
}
|
||||
count := end - start + 1
|
||||
if count > 50 {
|
||||
return nil, errors.New("每次最多批量创建 50 个用户")
|
||||
}
|
||||
width := max(len(startSequence), len(endSequence))
|
||||
accounts := make([]string, 0, count)
|
||||
for sequence := start; sequence <= end; sequence++ {
|
||||
accounts = append(accounts, fmt.Sprintf("%s-%0*d", prefix, width, sequence))
|
||||
}
|
||||
return accounts, nil
|
||||
}
|
||||
|
||||
func createUserWithAccount(tx *gorm.DB, account, passwordHash string, dailyLimit any) (map[string]any, error) {
|
||||
for range 20 {
|
||||
username, err := randomUsername()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
uid, err := newUID(tx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
id := uuid.NewString()
|
||||
result := tx.Exec(`INSERT INTO web_users(id,uid,username,account,password_hash,daily_limit,enabled)
|
||||
VALUES(?,?,?,?,?,?,true) ON CONFLICT DO NOTHING`, id, uid, username, account, passwordHash, dailyLimit)
|
||||
if result.Error != nil {
|
||||
return nil, result.Error
|
||||
}
|
||||
if result.RowsAffected == 1 {
|
||||
return map[string]any{"id": id, "uid": uid, "account": account, "username": username}, nil
|
||||
}
|
||||
var accountExists int64
|
||||
if err := tx.Table("web_users").Where("account = ?", account).Count(&accountExists).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if accountExists > 0 {
|
||||
return nil, fmt.Errorf("账号 %s 已存在", account)
|
||||
}
|
||||
}
|
||||
return nil, errors.New("用户名生成失败,请重试")
|
||||
}
|
||||
|
||||
func randomUsername() (string, error) {
|
||||
lengthOffset, err := rand.Int(rand.Reader, big.NewInt(8))
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
username := make([]byte, 5+lengthOffset.Int64())
|
||||
for i := range username {
|
||||
index, err := rand.Int(rand.Reader, big.NewInt(int64(len(usernameAlphabet))))
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
username[i] = usernameAlphabet[index.Int64()]
|
||||
}
|
||||
return string(username), nil
|
||||
}
|
||||
|
||||
func newUID(tx *gorm.DB) (string, error) {
|
||||
for i := 0; i < 20; i++ {
|
||||
n, err := rand.Int(rand.Reader, big.NewInt(90000000))
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
uid := fmt.Sprintf("%08d", n.Int64()+10000000)
|
||||
var count int64
|
||||
if err := tx.Table("web_users").Where("uid=?", uid).Count(&count).Error; err != nil {
|
||||
return "", err
|
||||
}
|
||||
if count == 0 {
|
||||
return uid, nil
|
||||
}
|
||||
}
|
||||
return "", errors.New("UID 生成失败,请重试")
|
||||
}
|
||||
|
||||
func (s *AdminData) UpdateUsers(ids []string, updates map[string]any) error {
|
||||
if len(ids) == 0 {
|
||||
return errors.New("请选择用户")
|
||||
}
|
||||
allowed := map[string]any{}
|
||||
if value, ok := updates["daily_limit"]; ok {
|
||||
allowed["daily_limit"] = value
|
||||
}
|
||||
if value, ok := updates["enabled"]; ok {
|
||||
allowed["enabled"] = value
|
||||
allowed["session_version"] = gorm.Expr("session_version + 1")
|
||||
}
|
||||
if len(allowed) == 0 {
|
||||
return errors.New("没有可更新的字段")
|
||||
}
|
||||
return s.DB.Table("web_users").Where("id IN ? AND deleted_at IS NULL", ids).Updates(allowed).Error
|
||||
}
|
||||
|
||||
func (s *AdminData) DeleteUsers(ids []string) error {
|
||||
if len(ids) == 0 {
|
||||
return errors.New("请选择用户")
|
||||
}
|
||||
return s.DB.Table("web_users").Where("id IN ? AND deleted_at IS NULL", ids).Updates(map[string]any{"deleted_at": time.Now(), "enabled": false, "session_version": gorm.Expr("session_version + 1")}).Error
|
||||
}
|
||||
|
||||
var placeholderPattern = regexp.MustCompile(`\{\{\s*([a-zA-Z_][a-zA-Z0-9_]*)\s*\}\}`)
|
||||
|
||||
func ValidatePromptVariables(content string, variables []string) error {
|
||||
allowed := map[string]bool{}
|
||||
for _, v := range variables {
|
||||
allowed[v] = true
|
||||
}
|
||||
missing := []string{}
|
||||
for _, m := range placeholderPattern.FindAllStringSubmatch(content, -1) {
|
||||
if !allowed[m[1]] {
|
||||
missing = append(missing, m[1])
|
||||
}
|
||||
}
|
||||
sort.Strings(missing)
|
||||
if len(missing) > 0 {
|
||||
return fmt.Errorf("正文使用了未声明变量:%s", strings.Join(missing, "、"))
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func GenerateRedemptionCode() (plain, hash, mask string, err error) {
|
||||
const alphabet = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"
|
||||
const firstAlphabet = "123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"
|
||||
raw := make([]byte, 30)
|
||||
for i := range raw {
|
||||
chars := alphabet
|
||||
if i == 0 {
|
||||
chars = firstAlphabet
|
||||
}
|
||||
index, randomErr := rand.Int(rand.Reader, big.NewInt(int64(len(chars))))
|
||||
if randomErr != nil {
|
||||
err = randomErr
|
||||
return
|
||||
}
|
||||
raw[i] = chars[index.Int64()]
|
||||
}
|
||||
normalized := string(raw)
|
||||
groups := make([]string, 0, 5)
|
||||
for start := 0; start < len(raw); start += 6 {
|
||||
groups = append(groups, string(raw[start:start+6]))
|
||||
}
|
||||
plain = strings.Join(groups, "-")
|
||||
sum := sha256.Sum256([]byte(normalized))
|
||||
hash = hex.EncodeToString(sum[:])
|
||||
mask = normalized[:3] + "***-******-******-******-" + normalized[len(normalized)-3:]
|
||||
return
|
||||
}
|
||||
@@ -0,0 +1,77 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"reflect"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestRandomUsernameLengthAndCharacters(t *testing.T) {
|
||||
for range 200 {
|
||||
username, err := randomUsername()
|
||||
if err != nil {
|
||||
t.Fatalf("randomUsername returned error: %v", err)
|
||||
}
|
||||
if len(username) < 5 || len(username) > 12 {
|
||||
t.Fatalf("username length must be 5 to 12, got %q", username)
|
||||
}
|
||||
for _, character := range username {
|
||||
if !strings.ContainsRune(usernameAlphabet, character) {
|
||||
t.Fatalf("username contains unsupported character: %q", username)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildBatchAccounts(t *testing.T) {
|
||||
accounts, err := buildBatchAccounts("RR", "001", "003")
|
||||
if err != nil {
|
||||
t.Fatalf("buildBatchAccounts returned error: %v", err)
|
||||
}
|
||||
want := []string{"RR-001", "RR-002", "RR-003"}
|
||||
if !reflect.DeepEqual(accounts, want) {
|
||||
t.Fatalf("accounts = %#v, want %#v", accounts, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildBatchAccountsUsesWidestSequence(t *testing.T) {
|
||||
accounts, err := buildBatchAccounts("A", "1", "003")
|
||||
if err != nil {
|
||||
t.Fatalf("buildBatchAccounts returned error: %v", err)
|
||||
}
|
||||
want := []string{"A-001", "A-002", "A-003"}
|
||||
if !reflect.DeepEqual(accounts, want) {
|
||||
t.Fatalf("accounts = %#v, want %#v", accounts, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildBatchAccountsAllowsFiftyUsers(t *testing.T) {
|
||||
accounts, err := buildBatchAccounts("RR", "001", "050")
|
||||
if err != nil {
|
||||
t.Fatalf("buildBatchAccounts returned error: %v", err)
|
||||
}
|
||||
if len(accounts) != 50 || accounts[0] != "RR-001" || accounts[49] != "RR-050" {
|
||||
t.Fatalf("unexpected account range: first=%q last=%q count=%d", accounts[0], accounts[len(accounts)-1], len(accounts))
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildBatchAccountsValidation(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
prefix string
|
||||
start string
|
||||
end string
|
||||
}{
|
||||
{name: "empty prefix", start: "001", end: "002"},
|
||||
{name: "non numeric start", prefix: "RR", start: "A01", end: "002"},
|
||||
{name: "start after end", prefix: "RR", start: "003", end: "002"},
|
||||
{name: "more than fifty", prefix: "RR", start: "001", end: "051"},
|
||||
}
|
||||
for _, test := range tests {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
if _, err := buildBatchAccounts(test.prefix, test.start, test.end); err == nil {
|
||||
t.Fatal("expected validation error")
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,179 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"juhe-factory/api/internal/model"
|
||||
"juhe-factory/api/internal/security"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
var ErrUnauthorized = errors.New("账号或密码错误")
|
||||
var ErrReauthFailed = errors.New("当前管理员密码错误")
|
||||
var ErrAdminPasswordConfirmation = errors.New("两次输入的新密码不一致")
|
||||
var ErrAdminPasswordLength = errors.New("新密码长度必须为 8~20 位")
|
||||
|
||||
type Auth struct {
|
||||
db *gorm.DB
|
||||
passwords security.PasswordHasher
|
||||
tokens security.TokenService
|
||||
}
|
||||
|
||||
type TokenPair struct {
|
||||
AccessToken string `json:"access_token"`
|
||||
RefreshToken string `json:"refresh_token"`
|
||||
ExpiresAt time.Time `json:"expires_at"`
|
||||
Admin map[string]any `json:"admin"`
|
||||
}
|
||||
|
||||
func NewAuth(db *gorm.DB, passwords security.PasswordHasher, tokens security.TokenService) *Auth {
|
||||
return &Auth{db: db, passwords: passwords, tokens: tokens}
|
||||
}
|
||||
|
||||
func (s *Auth) Bootstrap(username, password string) error {
|
||||
if strings.TrimSpace(password) == "" {
|
||||
return nil
|
||||
}
|
||||
username = strings.TrimSpace(username)
|
||||
var admin model.AdminUser
|
||||
err := s.db.Where("username = ? AND deleted_at IS NULL", username).First(&admin).Error
|
||||
if err == nil {
|
||||
if s.passwords.Verify(admin.PasswordHash, password) {
|
||||
return nil
|
||||
}
|
||||
hash, hashErr := s.passwords.Hash(password)
|
||||
if hashErr != nil {
|
||||
return hashErr
|
||||
}
|
||||
return s.db.Model(&admin).Update("password_hash", hash).Error
|
||||
}
|
||||
if !errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
return err
|
||||
}
|
||||
var count int64
|
||||
if err := s.db.Model(&model.AdminUser{}).Count(&count).Error; err != nil || count > 0 {
|
||||
return err
|
||||
}
|
||||
hash, err := s.passwords.Hash(password)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return s.db.Create(&model.AdminUser{Username: username, PasswordHash: hash, Enabled: true}).Error
|
||||
}
|
||||
|
||||
func (s *Auth) Login(username, password string) (*TokenPair, error) {
|
||||
var admin model.AdminUser
|
||||
if err := s.db.Where("username = ? AND deleted_at IS NULL", strings.TrimSpace(username)).First(&admin).Error; err != nil || !admin.Enabled || !s.passwords.Verify(admin.PasswordHash, password) {
|
||||
return nil, ErrUnauthorized
|
||||
}
|
||||
return s.issue(&admin, true)
|
||||
}
|
||||
|
||||
func (s *Auth) Refresh(refreshToken string) (*TokenPair, error) {
|
||||
hash := security.HashToken(refreshToken)
|
||||
var record model.AdminRefreshToken
|
||||
if err := s.db.Where("token_hash = ? AND revoked_at IS NULL AND expires_at > ?", hash, time.Now()).First(&record).Error; err != nil {
|
||||
return nil, ErrUnauthorized
|
||||
}
|
||||
var admin model.AdminUser
|
||||
if err := s.db.First(&admin, "id = ? AND enabled = true", record.AdminID).Error; err != nil {
|
||||
return nil, ErrUnauthorized
|
||||
}
|
||||
now := time.Now()
|
||||
if err := s.db.Model(&record).Update("revoked_at", now).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return s.issue(&admin, false)
|
||||
}
|
||||
|
||||
func (s *Auth) Logout(refreshToken string) error {
|
||||
if refreshToken == "" {
|
||||
return nil
|
||||
}
|
||||
now := time.Now()
|
||||
return s.db.Model(&model.AdminRefreshToken{}).Where("token_hash = ? AND revoked_at IS NULL", security.HashToken(refreshToken)).Update("revoked_at", now).Error
|
||||
}
|
||||
|
||||
func (s *Auth) Authenticate(raw string) (*model.AdminUser, error) {
|
||||
claims, err := s.tokens.ParseAccessToken(raw)
|
||||
if err != nil {
|
||||
return nil, ErrUnauthorized
|
||||
}
|
||||
id, err := uuid.Parse(claims.AdminID)
|
||||
if err != nil {
|
||||
return nil, ErrUnauthorized
|
||||
}
|
||||
var admin model.AdminUser
|
||||
if err := s.db.First(&admin, "id = ? AND enabled = true", id).Error; err != nil {
|
||||
return nil, ErrUnauthorized
|
||||
}
|
||||
return &admin, nil
|
||||
}
|
||||
|
||||
func (s *Auth) ChangePassword(adminID uuid.UUID, currentPassword, newPassword, confirmPassword string) error {
|
||||
if newPassword != confirmPassword {
|
||||
return ErrAdminPasswordConfirmation
|
||||
}
|
||||
if len(newPassword) < 8 || len(newPassword) > 20 {
|
||||
return ErrAdminPasswordLength
|
||||
}
|
||||
var admin model.AdminUser
|
||||
if err := s.db.Where("id = ? AND enabled = true AND deleted_at IS NULL", adminID).First(&admin).Error; err != nil {
|
||||
return ErrUnauthorized
|
||||
}
|
||||
if !s.passwords.Verify(admin.PasswordHash, currentPassword) {
|
||||
return ErrReauthFailed
|
||||
}
|
||||
hash, err := s.passwords.Hash(newPassword)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
now := time.Now()
|
||||
return s.db.Transaction(func(tx *gorm.DB) error {
|
||||
if err := tx.Model(&admin).Update("password_hash", hash).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Model(&model.AdminRefreshToken{}).
|
||||
Where("admin_id = ? AND revoked_at IS NULL", adminID).
|
||||
Update("revoked_at", now).Error
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Auth) Audit(admin *model.AdminUser, action, resource, resourceID, reason, ip, traceID string, detail any) error {
|
||||
return WriteAudit(s.db, admin, action, resource, resourceID, reason, ip, traceID, detail)
|
||||
}
|
||||
|
||||
func (s *Auth) issue(admin *model.AdminUser, updateLogin bool) (*TokenPair, error) {
|
||||
access, expiresAt, err := s.tokens.NewAccessToken(admin.ID.String(), admin.Username)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
refresh, hash, refreshExpires, err := s.tokens.NewRefreshToken()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
record := model.AdminRefreshToken{AdminID: admin.ID, TokenHash: hash, ExpiresAt: refreshExpires}
|
||||
if err := s.db.Transaction(func(tx *gorm.DB) error {
|
||||
if err := tx.Create(&record).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if updateLogin {
|
||||
return tx.Model(admin).Update("last_login_at", time.Now()).Error
|
||||
}
|
||||
return nil
|
||||
}); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &TokenPair{AccessToken: access, RefreshToken: refresh, ExpiresAt: expiresAt, Admin: map[string]any{"id": admin.ID, "username": admin.Username}}, nil
|
||||
}
|
||||
|
||||
func WriteAudit(db *gorm.DB, admin *model.AdminUser, action, resource, resourceID, reason, ip, traceID string, detail any) error {
|
||||
data, _ := json.Marshal(detail)
|
||||
log := model.AdminAuditLog{AdminID: &admin.ID, AdminUsername: admin.Username, Action: action, ResourceType: resource, ResourceID: resourceID, Reason: reason, Detail: string(data), IPAddress: ip, TraceID: traceID}
|
||||
return db.Create(&log).Error
|
||||
}
|
||||
@@ -0,0 +1,78 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"juhe-factory/api/internal/provider/apimart"
|
||||
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
type balanceChannel struct {
|
||||
ID uuid.UUID
|
||||
Name string
|
||||
BaseURL string
|
||||
APIKeyCiphertext string
|
||||
}
|
||||
|
||||
func (s *AdminData) triggerChannelBalanceSync() {
|
||||
if !s.balanceSyncing.CompareAndSwap(false, true) {
|
||||
return
|
||||
}
|
||||
go func() {
|
||||
defer s.balanceSyncing.Store(false)
|
||||
s.syncChannelBalances(context.Background())
|
||||
}()
|
||||
}
|
||||
|
||||
func (s *AdminData) syncChannelBalances(ctx context.Context) {
|
||||
if s.Encryptor == nil {
|
||||
return
|
||||
}
|
||||
var channels []balanceChannel
|
||||
if err := s.DB.Table("channels").Select("id,name,base_url,api_key_ciphertext").
|
||||
Where("deleted_at IS NULL AND api_key_ciphertext IS NOT NULL AND api_key_ciphertext<>''").
|
||||
Find(&channels).Error; err != nil {
|
||||
slog.Warn("读取待同步渠道余额失败", "error", err)
|
||||
return
|
||||
}
|
||||
|
||||
syncCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
|
||||
defer cancel()
|
||||
client := s.HTTPClient
|
||||
if client == nil {
|
||||
client = http.DefaultClient
|
||||
}
|
||||
|
||||
var wait sync.WaitGroup
|
||||
for _, channel := range channels {
|
||||
channel := channel
|
||||
wait.Add(1)
|
||||
go func() {
|
||||
defer wait.Done()
|
||||
apiKey, err := s.Encryptor.Decrypt(channel.APIKeyCiphertext)
|
||||
if err != nil {
|
||||
slog.Warn("解密渠道 API Key 失败", "channel", channel.Name, "error", err)
|
||||
return
|
||||
}
|
||||
balance, err := apimart.NewClient(client).Balance(syncCtx, channel.BaseURL, apiKey)
|
||||
if err != nil {
|
||||
slog.Warn("同步渠道余额失败", "channel", channel.Name, "error", err)
|
||||
return
|
||||
}
|
||||
remainBalance := balance.RemainBalance
|
||||
if balance.UnlimitedQuota {
|
||||
remainBalance = -1
|
||||
}
|
||||
if err := s.DB.Exec(`INSERT INTO channel_balance_snapshots(channel_id,balance,currency,synced_at)
|
||||
VALUES(?,?,?,?)`, channel.ID, remainBalance, "USD", time.Now()).Error; err != nil {
|
||||
slog.Warn("保存渠道余额快照失败", "channel", channel.Name, "error", err)
|
||||
}
|
||||
}()
|
||||
}
|
||||
wait.Wait()
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,22 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
"gorm.io/gorm/schema"
|
||||
)
|
||||
|
||||
func TestAnalysisModelConfigIgnoresComputedPricing(t *testing.T) {
|
||||
_, err := schema.Parse(&analysisModelConfig{}, &sync.Map{}, schema.NamingStrategy{})
|
||||
if err != nil {
|
||||
t.Fatalf("parse analysis model config: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestScriptAnalysisModelConfigIgnoresComputedPricing(t *testing.T) {
|
||||
_, err := schema.Parse(&scriptAnalysisModelConfig{}, &sync.Map{}, schema.NamingStrategy{})
|
||||
if err != nil {
|
||||
t.Fatalf("parse script analysis model config: %v", err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,113 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
type deletionMediaRow struct {
|
||||
ID uuid.UUID
|
||||
ObjectKey string
|
||||
}
|
||||
|
||||
type deletionMediaSet struct {
|
||||
rows map[uuid.UUID]string
|
||||
}
|
||||
|
||||
func newDeletionMediaSet() *deletionMediaSet {
|
||||
return &deletionMediaSet{rows: make(map[uuid.UUID]string)}
|
||||
}
|
||||
|
||||
func (set *deletionMediaSet) addQuery(query *gorm.DB) error {
|
||||
var rows []deletionMediaRow
|
||||
if err := query.Scan(&rows).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
for _, row := range rows {
|
||||
if row.ID != uuid.Nil {
|
||||
set.rows[row.ID] = row.ObjectKey
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (set *deletionMediaSet) objectKeys() []string {
|
||||
keys := make([]string, 0, len(set.rows))
|
||||
seen := make(map[string]struct{}, len(set.rows))
|
||||
for _, key := range set.rows {
|
||||
if key == "" {
|
||||
continue
|
||||
}
|
||||
if _, exists := seen[key]; exists {
|
||||
continue
|
||||
}
|
||||
seen[key] = struct{}{}
|
||||
keys = append(keys, key)
|
||||
}
|
||||
return keys
|
||||
}
|
||||
|
||||
func (set *deletionMediaSet) deleteRows(tx *gorm.DB) error {
|
||||
if len(set.rows) == 0 {
|
||||
return nil
|
||||
}
|
||||
ids := make([]uuid.UUID, 0, len(set.rows))
|
||||
for id := range set.rows {
|
||||
ids = append(ids, id)
|
||||
}
|
||||
if err := tx.Exec("DELETE FROM channel_asset_cache WHERE media_asset_id IN ?", ids).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Exec("DELETE FROM media_assets WHERE id IN ?", ids).Error
|
||||
}
|
||||
|
||||
func collectMediaByObjectKey(tx *gorm.DB, set *deletionMediaSet, pattern string) error {
|
||||
return set.addQuery(tx.Table("media_assets").Select("id,object_key").Where("object_key LIKE ?", pattern))
|
||||
}
|
||||
|
||||
// removeEpisodeDerivedContent removes everything produced from an episode analysis.
|
||||
// It intentionally keeps the episode source/subtitle and all project assets.
|
||||
func removeEpisodeDerivedContent(tx *gorm.DB, projectID, episodeID uuid.UUID, media *deletionMediaSet) error {
|
||||
// 封面来自本剧集的首帧,删除反推结果前先清理外键引用,避免随后删除媒体对象时产生约束错误。
|
||||
if err := tx.Exec(`UPDATE creative_projects SET cover_asset_id=NULL
|
||||
WHERE id=? AND cover_asset_id=(SELECT cover_asset_id FROM project_episodes WHERE id=?)`, projectID, episodeID).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if err := tx.Exec("UPDATE project_episodes SET cover_asset_id=NULL WHERE id=?", episodeID).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
taskScope := `(task_type IN ('prompt_reverse','video_generation')
|
||||
OR (task_type='image_generation' AND input_data->>'target_type'='storyboard')) AND
|
||||
(episode_id=? OR storyboard_id IN (SELECT id FROM episode_storyboards WHERE episode_id=?))`
|
||||
if err := media.addQuery(tx.Table("media_assets media").Select("DISTINCT media.id,media.object_key").
|
||||
Joins("JOIN generation_outputs output ON output.media_asset_id=media.id").
|
||||
Joins("JOIN generation_tasks task ON task.id=output.task_id").
|
||||
Where(taskScope, episodeID, episodeID)); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := media.addQuery(tx.Table("media_assets media").Select("DISTINCT media.id,media.object_key").
|
||||
Joins("JOIN episode_storyboards storyboard ON storyboard.thumbnail_asset_id=media.id").
|
||||
Where("storyboard.episode_id=?", episodeID)); err != nil {
|
||||
return err
|
||||
}
|
||||
derivedPrefix := fmt.Sprintf("juyou_ran/video-redraw/projects/%s/episodes/%s/storyboards/%%", projectID, episodeID)
|
||||
if err := collectMediaByObjectKey(tx, media, derivedPrefix); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := tx.Exec("UPDATE episode_storyboards SET active_output_id=NULL WHERE episode_id=?", episodeID).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if err := tx.Exec("DELETE FROM generation_outputs WHERE task_id IN (SELECT id FROM generation_tasks WHERE "+taskScope+")", episodeID, episodeID).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if err := tx.Exec("DELETE FROM generation_tasks WHERE "+taskScope, episodeID, episodeID).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if err := tx.Exec(`UPDATE generation_tasks SET storyboard_id=NULL
|
||||
WHERE storyboard_id IN (SELECT id FROM episode_storyboards WHERE episode_id=?)`, episodeID).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Exec("DELETE FROM episode_storyboards WHERE episode_id=?", episodeID).Error
|
||||
}
|
||||
@@ -0,0 +1,152 @@
|
||||
// 创作媒体服务,封装媒体上传前的所有权检查、关联更新和旧媒体清理事务。
|
||||
package service
|
||||
|
||||
import (
|
||||
"strings"
|
||||
|
||||
"juhe-factory/api/internal/model"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
// MediaUploadTarget 汇总上传前需要的业务对象及待替换媒体信息。
|
||||
type MediaUploadTarget struct {
|
||||
Project *model.CreativeProject
|
||||
Episode *model.ProjectEpisode
|
||||
Asset *model.ProjectAsset
|
||||
OldMediaID *uuid.UUID
|
||||
OldObjectKey string
|
||||
}
|
||||
|
||||
// MediaObject 查询当前用户有权播放或下载的媒体对象键。
|
||||
func (s *Creative) MediaObject(userID, mediaID uuid.UUID) (string, error) {
|
||||
var objectKey string
|
||||
err := s.DB.Table("media_assets").
|
||||
Where("id=? AND owner_user_id=? AND deleted_at IS NULL", mediaID, userID).
|
||||
Pluck("object_key", &objectKey).Error
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if strings.TrimSpace(objectKey) == "" {
|
||||
return "", gorm.ErrRecordNotFound
|
||||
}
|
||||
return objectKey, nil
|
||||
}
|
||||
|
||||
// PrepareRedrawMediaUpload 校验重绘项目所有权,并返回指定媒体类型的替换信息。
|
||||
func (s *Creative) PrepareRedrawMediaUpload(userID, projectID uuid.UUID, kind string) (MediaUploadTarget, error) {
|
||||
var project model.CreativeProject
|
||||
if err := s.DB.Where("id=? AND user_id=? AND project_type='video_redraw' AND deleted_at IS NULL", projectID, userID).Take(&project).Error; err != nil {
|
||||
return MediaUploadTarget{}, err
|
||||
}
|
||||
oldMediaID := project.SourceVideoAssetID
|
||||
if kind == "subtitle" {
|
||||
oldMediaID = project.SubtitleAssetID
|
||||
}
|
||||
oldObjectKey, err := s.mediaObjectKey(oldMediaID)
|
||||
if err != nil {
|
||||
return MediaUploadTarget{}, err
|
||||
}
|
||||
return MediaUploadTarget{Project: &project, OldMediaID: oldMediaID, OldObjectKey: oldObjectKey}, nil
|
||||
}
|
||||
|
||||
// PersistRedrawMediaUpload 在事务中保存新媒体、更新项目关联并删除旧媒体记录。
|
||||
func (s *Creative) PersistRedrawMediaUpload(projectID uuid.UUID, kind, currentStatus string, media *model.MediaAsset, oldMediaID *uuid.UUID) error {
|
||||
column, status := "source_video_asset_id", "uploaded"
|
||||
if kind == "subtitle" {
|
||||
column, status = "subtitle_asset_id", currentStatus
|
||||
}
|
||||
return s.DB.Transaction(func(tx *gorm.DB) error {
|
||||
if err := tx.Create(media).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
updates := map[string]any{column: media.ID, "redraw_status": status}
|
||||
if err := tx.Table("creative_projects").Where("id=?", projectID).Updates(updates).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
return deleteReplacedMedia(tx, oldMediaID)
|
||||
})
|
||||
}
|
||||
|
||||
// PrepareEpisodeMediaUpload 校验剧集与项目所有权,并返回指定媒体类型的替换信息。
|
||||
func (s *Creative) PrepareEpisodeMediaUpload(userID, projectID, episodeID uuid.UUID, kind string) (MediaUploadTarget, error) {
|
||||
var episode model.ProjectEpisode
|
||||
err := s.DB.Table("project_episodes e").Joins("JOIN creative_projects p ON p.id=e.project_id").
|
||||
Where("e.id=? AND e.project_id=? AND e.deleted_at IS NULL AND p.user_id=? AND p.deleted_at IS NULL", episodeID, projectID, userID).
|
||||
Select("e.*").First(&episode).Error
|
||||
if err != nil {
|
||||
return MediaUploadTarget{}, err
|
||||
}
|
||||
oldMediaID := episode.SourceVideoAssetID
|
||||
if kind == "subtitle" {
|
||||
oldMediaID = episode.SubtitleAssetID
|
||||
}
|
||||
oldObjectKey, err := s.mediaObjectKey(oldMediaID)
|
||||
if err != nil {
|
||||
return MediaUploadTarget{}, err
|
||||
}
|
||||
return MediaUploadTarget{Episode: &episode, OldMediaID: oldMediaID, OldObjectKey: oldObjectKey}, nil
|
||||
}
|
||||
|
||||
// PersistEpisodeMediaUpload 在事务中保存新媒体、更新剧集关联并删除旧媒体记录。
|
||||
func (s *Creative) PersistEpisodeMediaUpload(episodeID uuid.UUID, kind string, status string, media *model.MediaAsset, oldMediaID *uuid.UUID) error {
|
||||
column := "source_video_asset_id"
|
||||
if kind == "subtitle" {
|
||||
column = "subtitle_asset_id"
|
||||
}
|
||||
return s.DB.Transaction(func(tx *gorm.DB) error {
|
||||
if err := tx.Create(media).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if err := tx.Table("project_episodes").Where("id=?", episodeID).Updates(map[string]any{column: media.ID, "status": status}).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
return deleteReplacedMedia(tx, oldMediaID)
|
||||
})
|
||||
}
|
||||
|
||||
// PrepareAssetMediaUpload 校验普通创作项目中的资产所有权并返回资产信息。
|
||||
func (s *Creative) PrepareAssetMediaUpload(userID, projectID, assetID uuid.UUID) (MediaUploadTarget, error) {
|
||||
var asset model.ProjectAsset
|
||||
err := s.DB.Table("project_assets a").Joins("JOIN creative_projects p ON p.id=a.project_id").
|
||||
Where("a.id=? AND a.project_id=? AND a.deleted_at IS NULL AND p.user_id=? AND p.project_type<>'video_redraw' AND p.deleted_at IS NULL", assetID, projectID, userID).
|
||||
Select("a.*").First(&asset).Error
|
||||
if err != nil {
|
||||
return MediaUploadTarget{}, err
|
||||
}
|
||||
return MediaUploadTarget{Asset: &asset}, nil
|
||||
}
|
||||
|
||||
// AttachUploadedAssetAudio 在事务中保存上传音频并更新资产关联。
|
||||
func (s *Creative) AttachUploadedAssetAudio(assetID uuid.UUID, media *model.MediaAsset) error {
|
||||
return s.DB.Transaction(func(tx *gorm.DB) error {
|
||||
if err := tx.Create(media).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Table("project_assets").Where("id=?", assetID).Update("audio_asset_id", media.ID).Error
|
||||
})
|
||||
}
|
||||
|
||||
// mediaObjectKey 查询待替换媒体的对象键;空媒体标识表示无需清理。
|
||||
func (s *Creative) mediaObjectKey(mediaID *uuid.UUID) (string, error) {
|
||||
if mediaID == nil {
|
||||
return "", nil
|
||||
}
|
||||
var objectKey string
|
||||
if err := s.DB.Table("media_assets").Where("id=?", *mediaID).Pluck("object_key", &objectKey).Error; err != nil {
|
||||
return "", err
|
||||
}
|
||||
return objectKey, nil
|
||||
}
|
||||
|
||||
// deleteReplacedMedia 删除旧媒体的渠道缓存和媒体记录,供上传事务复用。
|
||||
func deleteReplacedMedia(tx *gorm.DB, mediaID *uuid.UUID) error {
|
||||
if mediaID == nil {
|
||||
return nil
|
||||
}
|
||||
if err := tx.Exec("DELETE FROM channel_asset_cache WHERE media_asset_id=?", *mediaID).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Exec("DELETE FROM media_assets WHERE id=?", *mediaID).Error
|
||||
}
|
||||
@@ -0,0 +1,73 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"testing"
|
||||
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
func TestAssetMentionIndex(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
content string
|
||||
asset string
|
||||
want int
|
||||
}{
|
||||
{name: "space terminated mention", content: "镜头跟随 @丫丫 向前移动", asset: "丫丫", want: 13},
|
||||
{name: "punctuation terminated mention", content: "看向@客厅,停顿。", asset: "客厅", want: 6},
|
||||
{name: "end terminated mention", content: "参考@分镜 1 分镜图", asset: "分镜 1 分镜图", want: 6},
|
||||
{name: "reject shorter prefix", content: "@父亲 走进房间", asset: "父", want: -1},
|
||||
{name: "reject plain asset name", content: "父亲走进房间", asset: "父亲", want: -1},
|
||||
}
|
||||
for _, test := range tests {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
if got := assetMentionIndex(test.content, test.asset); got != test.want {
|
||||
t.Fatalf("assetMentionIndex(%q, %q) = %d, want %d", test.content, test.asset, got, test.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestRenameAssetMentions(t *testing.T) {
|
||||
tests := []struct {
|
||||
name, content, oldName, newName, want string
|
||||
protectedNames []string
|
||||
}{
|
||||
{name: "all complete mentions", content: "@父亲 走向@客厅,随后看向@父亲", oldName: "父亲", newName: "老周", want: "@老周 走向@客厅,随后看向@老周"},
|
||||
{name: "Chinese description without separator", content: "@林溪微微睁大眼睛", oldName: "林溪", newName: "林溪月", want: "@林溪月微微睁大眼睛"},
|
||||
{name: "reject known longer asset", content: "@父亲 走进房间", oldName: "父", newName: "老周", protectedNames: []string{"父亲"}, want: "@父亲 走进房间"},
|
||||
{name: "ignore plain name", content: "父亲走进房间", oldName: "父亲", newName: "老周", want: "父亲走进房间"},
|
||||
{name: "punctuation in name", content: "参考@A+B。", oldName: "A+B", newName: "组合", want: "参考@组合。"},
|
||||
}
|
||||
for _, test := range tests {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
if got := renameAssetMentions(test.content, test.oldName, test.newName, test.protectedNames); got != test.want {
|
||||
t.Fatalf("renameAssetMentions() = %q, want %q", got, test.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestRenameAssetRef(t *testing.T) {
|
||||
assetID := uuid.New()
|
||||
otherID := uuid.NewString()
|
||||
value := json.RawMessage(`[{"id":"` + assetID.String() + `","name":"父亲","type":"character"},{"id":"` + otherID + `","name":"客厅"}]`)
|
||||
got, changed, err := renameAssetRef(value, assetID, "老周")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var refs []map[string]any
|
||||
if unmarshalErr := json.Unmarshal(got, &refs); unmarshalErr != nil || !changed || refs[0]["name"] != "老周" || refs[1]["name"] != "客厅" || refs[1]["id"] != otherID {
|
||||
t.Fatalf("renameAssetRef() = %s, changed=%v, err=%v", got, changed, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAssetRefNamesKeepsHistoricalMentionName(t *testing.T) {
|
||||
assetID := uuid.New()
|
||||
value := json.RawMessage(`[{"id":"` + assetID.String() + `","name":"改名前名称","type":"character"}]`)
|
||||
names, err := assetRefNames(value, assetID)
|
||||
if err != nil || len(names) != 1 || names[0] != "改名前名称" {
|
||||
t.Fatalf("assetRefNames() = %#v, err=%v", names, err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,54 @@
|
||||
package service
|
||||
|
||||
import "testing"
|
||||
|
||||
func TestShortDramaProjectDoesNotRequireLocalization(t *testing.T) {
|
||||
input := normalizeCreateProjectInput(ProjectInput{
|
||||
Name: "测试短剧",
|
||||
StyleID: "00000000-0000-0000-0000-000000000001",
|
||||
EraType: "modern_city",
|
||||
AspectRatio: "9:16",
|
||||
})
|
||||
if input.ProjectType != "premium_drama" {
|
||||
t.Fatalf("expected request without localization to be treated as premium_drama, got %q", input.ProjectType)
|
||||
}
|
||||
if err := validateProjectInput(input); err != nil {
|
||||
t.Fatalf("short drama project should not require localization: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestVideoRedrawScriptUsesCreationDefaults(t *testing.T) {
|
||||
followerCount := "百万粉丝"
|
||||
input := normalizeCreateProjectInput(ProjectInput{
|
||||
ProjectType: "video_redraw",
|
||||
Name: "测试剧本",
|
||||
StyleID: "00000000-0000-0000-0000-000000000001",
|
||||
ShortDramaType: "都市情感",
|
||||
PlayCount: "全网播放破亿",
|
||||
AudienceProfile: "25 至 35 岁女性",
|
||||
Producer: "测试出品方",
|
||||
CastMembers: []ProjectCastMemberInput{{
|
||||
Name: "测试主演", FollowerCount: &followerCount,
|
||||
}},
|
||||
})
|
||||
if input.EraType != "modern_city" || input.AspectRatio != "9:16" {
|
||||
t.Fatalf("unexpected redraw creation defaults: era=%q ratio=%q", input.EraType, input.AspectRatio)
|
||||
}
|
||||
if input.Localization == nil || *input.Localization != "china" {
|
||||
t.Fatalf("expected default localization china, got %v", input.Localization)
|
||||
}
|
||||
if err := validateProjectInput(input); err != nil {
|
||||
t.Fatalf("redraw project metadata should be valid after defaults: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestVideoRedrawProjectOnlyRequiresName(t *testing.T) {
|
||||
input := normalizeCreateProjectInput(ProjectInput{
|
||||
ProjectType: "video_redraw",
|
||||
Name: "测试剧本",
|
||||
StyleID: "00000000-0000-0000-0000-000000000001",
|
||||
})
|
||||
if err := validateProjectInput(input); err != nil {
|
||||
t.Fatalf("redraw project optional metadata should be accepted: %v", err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,735 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"juhe-factory/api/internal/billing"
|
||||
"juhe-factory/api/internal/model"
|
||||
queuepkg "juhe-factory/api/internal/queue"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/clause"
|
||||
)
|
||||
|
||||
var activeTaskStatuses = []string{"pending_submission", "submitting", "submitted", "processing", "result_ready", "downloading", "cancel_requested"}
|
||||
|
||||
func (s *Creative) InsertStoryboard(userID, projectID, storyboardID uuid.UUID, position string) (*model.EpisodeStoryboard, error) {
|
||||
if position != "before" && position != "after" {
|
||||
return nil, errors.New("插入位置无效")
|
||||
}
|
||||
created := &model.EpisodeStoryboard{}
|
||||
err := s.DB.Transaction(func(tx *gorm.DB) error {
|
||||
var target model.EpisodeStoryboard
|
||||
if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Table("episode_storyboards sb").Select("sb.*").
|
||||
Joins(`JOIN creative_projects p ON p.id=? AND p.user_id=? AND p.deleted_at IS NULL AND
|
||||
(sb.project_id=p.id OR EXISTS (SELECT 1 FROM project_episodes e WHERE e.id=sb.episode_id AND e.project_id=p.id AND e.deleted_at IS NULL))`, projectID, userID).
|
||||
Where("sb.id=? AND sb.deleted_at IS NULL", storyboardID).Take(&target).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
parentColumn, parentID := storyboardParent(target)
|
||||
sequence, startMS := target.SequenceNo, target.StartMS
|
||||
if position == "after" {
|
||||
sequence, startMS = target.SequenceNo+1, target.EndMS
|
||||
}
|
||||
if err := tx.Exec(`UPDATE episode_storyboards SET sequence_no=sequence_no+100000
|
||||
WHERE `+parentColumn+`=? AND sequence_no>=? AND deleted_at IS NULL`, parentID, sequence).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if err := tx.Exec(`UPDATE episode_storyboards SET sequence_no=sequence_no-99999,start_ms=start_ms+5000,end_ms=end_ms+5000
|
||||
WHERE `+parentColumn+`=? AND sequence_no>=? AND deleted_at IS NULL`, parentID, sequence+100000).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
created = &model.EpisodeStoryboard{
|
||||
EpisodeID: target.EpisodeID, ProjectID: target.ProjectID, SequenceNo: sequence, StableKey: "manual-" + uuid.NewString(),
|
||||
StartMS: startMS, EndMS: startMS + 5000, DurationSeconds: 5, Title: fmt.Sprintf("分镜 %d", sequence),
|
||||
Dialogue: json.RawMessage(`[]`), AssetRefs: json.RawMessage(`[]`), Status: "idle",
|
||||
}
|
||||
return tx.Create(created).Error
|
||||
})
|
||||
return created, err
|
||||
}
|
||||
|
||||
func (s *Creative) DeleteStoryboard(userID, projectID, storyboardID uuid.UUID) ([]string, error) {
|
||||
media := newDeletionMediaSet()
|
||||
err := s.DB.Transaction(func(tx *gorm.DB) error {
|
||||
var storyboard model.EpisodeStoryboard
|
||||
if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Table("episode_storyboards sb").Select("sb.*").
|
||||
Joins(`JOIN creative_projects p ON p.id=? AND p.user_id=? AND p.deleted_at IS NULL AND
|
||||
(sb.project_id=p.id OR EXISTS (SELECT 1 FROM project_episodes e WHERE e.id=sb.episode_id AND e.project_id=p.id AND e.deleted_at IS NULL))`, projectID, userID).
|
||||
Where("sb.id=? AND sb.deleted_at IS NULL", storyboardID).Take(&storyboard).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
parentColumn, parentID := storyboardParent(storyboard)
|
||||
var active int64
|
||||
if err := tx.Model(&model.GenerationTask{}).Where("storyboard_id=? AND status IN ?", storyboardID, activeTaskStatuses).Count(&active).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if active > 0 {
|
||||
return errors.New("当前分镜存在进行中的任务,暂时不能删除")
|
||||
}
|
||||
if err := media.addQuery(tx.Table("media_assets media").Select("DISTINCT media.id,media.object_key").
|
||||
Joins("JOIN generation_outputs output ON output.media_asset_id=media.id").
|
||||
Joins("JOIN generation_tasks gt ON gt.id=output.task_id").
|
||||
Where("gt.storyboard_id=? AND (gt.task_type IN ('video_generation','prompt_reverse') OR (gt.task_type='image_generation' AND gt.input_data->>'target_type'='storyboard'))", storyboardID)); err != nil {
|
||||
return err
|
||||
}
|
||||
if storyboard.ThumbnailAssetID != nil {
|
||||
if err := media.addQuery(tx.Table("media_assets").Select("id,object_key").Where("id=?", *storyboard.ThumbnailAssetID)); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
storyboardMediaPattern := fmt.Sprintf("juyou_ran/video-redraw/projects/%s/storyboards/%%-%s/%%", projectID, storyboardID)
|
||||
if storyboard.EpisodeID != nil {
|
||||
storyboardMediaPattern = fmt.Sprintf("juyou_ran/video-redraw/projects/%s/episodes/%s/storyboards/%%-%s/%%", projectID, *storyboard.EpisodeID, storyboardID)
|
||||
}
|
||||
if err := collectMediaByObjectKey(tx, media, storyboardMediaPattern); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := tx.Model(&storyboard).Update("active_output_id", nil).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if err := tx.Exec(`DELETE FROM generation_outputs WHERE task_id IN (
|
||||
SELECT id FROM generation_tasks WHERE storyboard_id=? AND (task_type IN ('video_generation','prompt_reverse')
|
||||
OR (task_type='image_generation' AND input_data->>'target_type'='storyboard')))`, storyboardID).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if err := tx.Exec(`DELETE FROM generation_tasks WHERE storyboard_id=? AND (task_type IN ('video_generation','prompt_reverse')
|
||||
OR (task_type='image_generation' AND input_data->>'target_type'='storyboard'))`, storyboardID).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if err := tx.Model(&model.GenerationTask{}).Where("storyboard_id=?", storyboardID).Update("storyboard_id", nil).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if err := tx.Exec("DELETE FROM episode_storyboards WHERE id=?", storyboardID).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if err := tx.Exec(`UPDATE episode_storyboards SET sequence_no=sequence_no+100000
|
||||
WHERE `+parentColumn+`=? AND sequence_no>? AND deleted_at IS NULL`, parentID, storyboard.SequenceNo).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
shiftMS := int64(storyboard.DurationSeconds * 1000)
|
||||
if err := tx.Exec(`UPDATE episode_storyboards SET sequence_no=sequence_no-100001,
|
||||
start_ms=GREATEST(0,start_ms-?),end_ms=GREATEST(1,end_ms-?)
|
||||
WHERE `+parentColumn+`=? AND sequence_no>? AND deleted_at IS NULL`, shiftMS, shiftMS, parentID, storyboard.SequenceNo+100000).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
return media.deleteRows(tx)
|
||||
})
|
||||
return media.objectKeys(), err
|
||||
}
|
||||
|
||||
func storyboardParent(storyboard model.EpisodeStoryboard) (string, uuid.UUID) {
|
||||
if storyboard.ProjectID != nil {
|
||||
return "project_id", *storyboard.ProjectID
|
||||
}
|
||||
return "episode_id", *storyboard.EpisodeID
|
||||
}
|
||||
|
||||
func (s *Creative) DeleteEpisode(userID, projectID, episodeID uuid.UUID) ([]string, error) {
|
||||
media := newDeletionMediaSet()
|
||||
err := s.DB.Transaction(func(tx *gorm.DB) error {
|
||||
var episode model.ProjectEpisode
|
||||
if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Table("project_episodes e").Select("e.*").
|
||||
Joins("JOIN creative_projects p ON p.id=e.project_id AND p.user_id=? AND p.deleted_at IS NULL", userID).
|
||||
Where("e.id=? AND e.project_id=? AND e.deleted_at IS NULL", episodeID, projectID).Take(&episode).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
var active int64
|
||||
if err := tx.Model(&model.GenerationTask{}).Where("episode_id=? AND status IN ?", episodeID, activeTaskStatuses).Count(&active).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if active > 0 {
|
||||
return errors.New("剧集存在排队中或处理中的任务,暂时不能删除")
|
||||
}
|
||||
if err := tx.Model(&model.DramaParseTask{}).Where("episode_id=? AND status IN ?", episodeID, []string{"queued", "running", "retry_wait", "cancel_requested"}).Count(&active).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if active > 0 {
|
||||
return errors.New("剧集存在进行中的剧本解析任务,暂时不能删除")
|
||||
}
|
||||
if err := removeEpisodeDerivedContent(tx, projectID, episodeID, media); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := media.addQuery(tx.Table("media_assets media").Select("DISTINCT media.id,media.object_key").
|
||||
Joins(`JOIN project_episodes episode ON media.id=episode.cover_asset_id
|
||||
OR media.id=episode.source_video_asset_id OR media.id=episode.subtitle_asset_id`).
|
||||
Where("episode.id=?", episodeID)); err != nil {
|
||||
return err
|
||||
}
|
||||
// 如果删除的是第一集,将项目封面切换到新的第一集封面。
|
||||
if episode.CoverAssetID != nil {
|
||||
var nextEpisode struct{ CoverAssetID *uuid.UUID }
|
||||
if err := tx.Table("project_episodes").Select("cover_asset_id").Where("project_id=? AND id<>? AND deleted_at IS NULL AND cover_asset_id IS NOT NULL", projectID, episodeID).
|
||||
Order("episode_no").Limit(1).Find(&nextEpisode).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if err := tx.Model(&model.CreativeProject{}).Where("id=?", projectID).Update("cover_asset_id", nextEpisode.CoverAssetID).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
episodeMediaPattern := fmt.Sprintf("juyou_ran/video-redraw/projects/%s/episodes/%%%s/%%", projectID, episodeID)
|
||||
if err := collectMediaByObjectKey(tx, media, episodeMediaPattern); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := tx.Exec("UPDATE generation_tasks SET episode_id=NULL WHERE episode_id=?", episodeID).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if err := tx.Exec("UPDATE project_assets SET source_episode_id=NULL WHERE source_episode_id=?", episodeID).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if err := tx.Exec("UPDATE drama_parse_batches SET current_episode_id=NULL WHERE current_episode_id=?", episodeID).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if err := tx.Exec("DELETE FROM project_episodes WHERE id=?", episodeID).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
return media.deleteRows(tx)
|
||||
})
|
||||
return media.objectKeys(), err
|
||||
}
|
||||
|
||||
// ResetEpisodeAnalysis permanently removes generated storyboards, prompts and their task history.
|
||||
// Project assets are intentionally left untouched so they can be reused for a new analysis.
|
||||
func (s *Creative) ResetEpisodeAnalysis(userID, projectID, episodeID uuid.UUID) ([]string, error) {
|
||||
media := newDeletionMediaSet()
|
||||
err := s.DB.Transaction(func(tx *gorm.DB) error {
|
||||
var episode model.ProjectEpisode
|
||||
if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Table("project_episodes e").Select("e.*").
|
||||
Joins("JOIN creative_projects p ON p.id=e.project_id AND p.user_id=? AND p.deleted_at IS NULL", userID).
|
||||
Where("e.id=? AND e.project_id=? AND e.deleted_at IS NULL", episodeID, projectID).Take(&episode).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
var active int64
|
||||
if err := tx.Model(&model.GenerationTask{}).Where("episode_id=? AND status IN ?", episodeID, activeTaskStatuses).Count(&active).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if active > 0 {
|
||||
return errors.New("当前剧集存在进行中的生成任务,暂时不能重新解析")
|
||||
}
|
||||
if err := tx.Model(&model.DramaParseTask{}).Where("episode_id=? AND status IN ?", episodeID, []string{"queued", "running", "retry_wait", "cancel_requested"}).Count(&active).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if active > 0 {
|
||||
return errors.New("当前剧集存在进行中的剧本解析任务,暂时不能重新解析")
|
||||
}
|
||||
if err := removeEpisodeDerivedContent(tx, projectID, episodeID, media); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := media.deleteRows(tx); err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Model(&episode).Updates(map[string]any{"status": "uploaded", "analysis_message": nil, "redraw_script": nil}).Error
|
||||
})
|
||||
return media.objectKeys(), err
|
||||
}
|
||||
|
||||
func (s *Creative) ListTasks(userID, projectID uuid.UUID, episodeID *uuid.UUID) ([]map[string]any, error) {
|
||||
items := make([]map[string]any, 0)
|
||||
query := s.DB.Table("generation_tasks gt").
|
||||
Select(`gt.id,gt.request_id,gt.task_type,gt.status,gt.episode_id,gt.storyboard_id,gt.input_data->>'mode' AS mode,gt.input_data->>'phase_started_at' AS phase_started_at,gt.estimated_points::text AS estimated_points,
|
||||
gt.actual_points::text AS actual_points,gt.error_code,gt.error_message,gt.created_at,gt.submitted_at,gt.finished_at,
|
||||
coalesce(sb.sequence_no,(gt.input_data->>'current_sequence')::int,0) AS sequence_no,gt.input_data->>'asset_id' AS asset_id,media.public_url AS result_url,
|
||||
(SELECT max(progress.created_at) FROM generation_outputs progress WHERE progress.task_id=gt.id) AS last_result_at`).
|
||||
Joins("JOIN creative_projects p ON p.id=gt.project_id AND p.user_id=? AND p.deleted_at IS NULL", userID).
|
||||
Joins("LEFT JOIN episode_storyboards sb ON sb.id=gt.storyboard_id").
|
||||
Joins("LEFT JOIN generation_outputs output ON output.task_id=gt.id AND output.sequence_no=1").
|
||||
Joins("LEFT JOIN media_assets media ON media.id=output.media_asset_id AND media.deleted_at IS NULL").
|
||||
Where("gt.project_id=?", projectID)
|
||||
if episodeID != nil {
|
||||
query = query.Where("gt.episode_id=?", *episodeID)
|
||||
}
|
||||
return items, query.Order("gt.created_at DESC").Limit(500).Find(&items).Error
|
||||
}
|
||||
|
||||
func (s *Creative) CancelTask(userID, projectID, taskID uuid.UUID) error {
|
||||
var channelID *uuid.UUID
|
||||
err := s.DB.Transaction(func(tx *gorm.DB) error {
|
||||
var task model.GenerationTask
|
||||
if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Where("id=? AND project_id=? AND user_id=?", taskID, projectID, userID).First(&task).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
var projectType string
|
||||
if err := tx.Table("creative_projects").Where("id=?", projectID).Pluck("project_type", &projectType).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
channelID = task.ChannelID
|
||||
if task.TaskType == "prompt_reverse" && (task.Status == "pending_submission" || task.Status == "submitted" || task.Status == "processing" || task.Status == "cancel_requested") {
|
||||
return cancelPromptReverseTask(tx, &task)
|
||||
}
|
||||
switch task.Status {
|
||||
case "pending_submission":
|
||||
refunded := false
|
||||
chargeCancellation := task.TaskType == "video_generation" || projectType == "premium_drama"
|
||||
if !chargeCancellation {
|
||||
var err error
|
||||
refunded, err = billing.RefundGenerationTask(tx, &task, "视频生成失败返还")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
actual := "0.00"
|
||||
if chargeCancellation {
|
||||
actual = task.PrepaidPoints
|
||||
}
|
||||
updates := map[string]any{"status": "cancelled", "actual_points": actual, "finished_at": time.Now(), "error_code": nil, "error_message": "用户取消"}
|
||||
if refunded {
|
||||
updates["cost_refunded"] = true
|
||||
}
|
||||
if err := tx.Model(&task).Updates(updates).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
return resetRelatedStatus(tx, task)
|
||||
case "submitted", "processing":
|
||||
if task.TaskType != "prompt_reverse" && strings.TrimSpace(task.UpstreamTaskID) == "" {
|
||||
return errors.New("上游任务标识尚未就绪,请稍后重试取消")
|
||||
}
|
||||
return tx.Model(&task).Updates(map[string]any{"status": "cancel_requested", "error_code": nil, "error_message": "等待当前处理安全结束"}).Error
|
||||
case "submitting":
|
||||
return errors.New("任务正在提交,请稍后重试取消")
|
||||
case "cancel_requested":
|
||||
return nil
|
||||
case "result_ready", "downloading", "succeeded":
|
||||
return errors.New("生成已经完成,不能取消")
|
||||
case "failed", "cancelled":
|
||||
return nil
|
||||
default:
|
||||
return fmt.Errorf("任务状态 %s 不允许取消", task.Status)
|
||||
}
|
||||
})
|
||||
if err == nil && channelID != nil && s.Queue != nil {
|
||||
_ = queuepkg.EnqueueID(s.Queue, queuepkg.TypeDispatchChannel, *channelID, 0)
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
||||
func cancelPromptReverseTask(tx *gorm.DB, task *model.GenerationTask) error {
|
||||
var input struct {
|
||||
Mode string `json:"mode"`
|
||||
}
|
||||
_ = json.Unmarshal(task.InputData, &input)
|
||||
updates := map[string]any{
|
||||
"status": "cancelled", "actual_points": task.PrepaidPoints, "finished_at": time.Now(),
|
||||
"error_code": nil, "error_message": "用户取消,已扣积分不退",
|
||||
}
|
||||
if err := tx.Model(task).Updates(updates).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if task.StoryboardID != nil {
|
||||
return nil
|
||||
}
|
||||
if task.EpisodeID == nil && task.ProjectID != nil {
|
||||
if input.Mode == "script" {
|
||||
return tx.Model(&model.CreativeProject{}).Where("id=?", *task.ProjectID).
|
||||
Updates(map[string]any{"redraw_status": "review", "analysis_message": "剧本反推已取消"}).Error
|
||||
}
|
||||
var storyboardCount int64
|
||||
if err := tx.Model(&model.EpisodeStoryboard{}).Where("project_id=? AND deleted_at IS NULL", *task.ProjectID).Count(&storyboardCount).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
status := "uploaded"
|
||||
if storyboardCount > 0 {
|
||||
status = "review"
|
||||
}
|
||||
return tx.Model(&model.CreativeProject{}).Where("id=?", *task.ProjectID).Updates(map[string]any{"redraw_status": status, "analysis_message": "视频分析已取消"}).Error
|
||||
}
|
||||
if task.EpisodeID == nil {
|
||||
return nil
|
||||
}
|
||||
var storyboardCount int64
|
||||
if err := tx.Model(&model.EpisodeStoryboard{}).Where("episode_id=? AND deleted_at IS NULL", *task.EpisodeID).Count(&storyboardCount).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
status := "uploaded"
|
||||
if storyboardCount > 0 {
|
||||
status = "review"
|
||||
}
|
||||
return tx.Model(&model.ProjectEpisode{}).Where("id=?", *task.EpisodeID).Updates(map[string]any{"status": status, "analysis_message": "视频分析已取消"}).Error
|
||||
}
|
||||
|
||||
func (s *Creative) failQueuedTask(taskID uuid.UUID, code, message string) error {
|
||||
return s.failQueuedTasks([]uuid.UUID{taskID}, code, message)
|
||||
}
|
||||
|
||||
func (s *Creative) failQueuedTasks(taskIDs []uuid.UUID, code, message string) error {
|
||||
return s.DB.Transaction(func(tx *gorm.DB) error {
|
||||
for _, taskID := range taskIDs {
|
||||
var task model.GenerationTask
|
||||
if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Where("id=?", taskID).First(&task).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if task.Status == "succeeded" || task.Status == "failed" || task.Status == "cancelled" {
|
||||
continue
|
||||
}
|
||||
refundRemark := map[string]string{"prompt_reverse": "视频反推失败返还", "script_analysis": "剧本分析失败返还", "image_generation": "图片生成失败返还", "video_generation": "视频生成失败返还"}[task.TaskType]
|
||||
var refunded bool
|
||||
var err error
|
||||
if task.TaskType == "prompt_reverse" || task.TaskType == "script_analysis" {
|
||||
refunded, err = billing.RefundTextGenerationTask(tx, &task, refundRemark)
|
||||
} else {
|
||||
refunded, err = billing.RefundGenerationTask(tx, &task, refundRemark)
|
||||
}
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
updates := map[string]any{
|
||||
"status": "failed", "actual_points": "0.00", "finished_at": time.Now(),
|
||||
"error_code": code, "error_message": message,
|
||||
}
|
||||
if refunded {
|
||||
updates["cost_refunded"] = true
|
||||
}
|
||||
if err := tx.Model(&task).Updates(updates).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if task.TaskType == "video_generation" && task.StoryboardID != nil {
|
||||
if err := tx.Model(&model.EpisodeStoryboard{}).Where("id=?", *task.StoryboardID).Update("status", "idle").Error; err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
if task.TaskType == "prompt_reverse" && task.StoryboardID == nil && task.EpisodeID != nil {
|
||||
var storyboardCount int64
|
||||
if err := tx.Model(&model.EpisodeStoryboard{}).Where("episode_id=? AND deleted_at IS NULL", *task.EpisodeID).Count(&storyboardCount).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
status := "uploaded"
|
||||
if storyboardCount > 0 {
|
||||
status = "review"
|
||||
}
|
||||
if err := tx.Model(&model.ProjectEpisode{}).Where("id=?", *task.EpisodeID).Updates(map[string]any{"status": status, "analysis_message": message}).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
} else if task.TaskType == "prompt_reverse" && task.StoryboardID == nil && task.ProjectID != nil {
|
||||
var storyboardCount int64
|
||||
if err := tx.Model(&model.EpisodeStoryboard{}).Where("project_id=? AND deleted_at IS NULL", *task.ProjectID).Count(&storyboardCount).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
status := "uploaded"
|
||||
if storyboardCount > 0 {
|
||||
status = "review"
|
||||
}
|
||||
if err := tx.Model(&model.CreativeProject{}).Where("id=?", *task.ProjectID).Updates(map[string]any{"redraw_status": status, "analysis_message": message}).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
})
|
||||
}
|
||||
|
||||
func resetRelatedStatus(tx *gorm.DB, task model.GenerationTask) error {
|
||||
if task.TaskType == "video_generation" && task.StoryboardID != nil {
|
||||
var active int64
|
||||
if err := tx.Model(&model.GenerationTask{}).Where("storyboard_id=? AND task_type='video_generation' AND status IN ?", *task.StoryboardID, activeTaskStatuses).Count(&active).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if active == 0 {
|
||||
return tx.Model(&model.EpisodeStoryboard{}).Where("id=?", *task.StoryboardID).Update("status", "idle").Error
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *Creative) ListStoryboardOutputs(userID, projectID, storyboardID uuid.UUID) ([]map[string]any, error) {
|
||||
items := make([]map[string]any, 0)
|
||||
err := s.DB.Table("generation_outputs output").
|
||||
Select(`output.id,output.media_asset_id,output.output_type,output.sequence_no,output.metadata,output.created_at,output.task_id,
|
||||
media.public_url,media.mime_type,media.size_bytes,gt.estimated_points::text AS estimated_points,
|
||||
gt.actual_points::text AS actual_points,CASE WHEN gt.task_type='video_generation' THEN sb.active_output_id=output.id ELSE sb.thumbnail_asset_id=output.media_asset_id END AS active,
|
||||
CASE WHEN gt.task_type='video_generation' THEN sb.active_output_id=output.id OR coalesce(output.metadata->>'candidate','false')='true' ELSE false END AS candidate`).
|
||||
Joins("JOIN generation_tasks gt ON gt.id=output.task_id AND gt.storyboard_id=? AND (gt.task_type='video_generation' OR (gt.task_type='image_generation' AND gt.input_data->>'target_type'='storyboard'))", storyboardID).
|
||||
Joins("JOIN media_assets media ON media.id=output.media_asset_id AND media.deleted_at IS NULL").
|
||||
Joins("JOIN episode_storyboards sb ON sb.id=gt.storyboard_id AND sb.deleted_at IS NULL").
|
||||
Joins(`JOIN creative_projects p ON p.id=? AND p.user_id=? AND p.deleted_at IS NULL AND
|
||||
(sb.project_id=p.id OR EXISTS (SELECT 1 FROM project_episodes e WHERE e.id=sb.episode_id AND e.project_id=p.id AND e.deleted_at IS NULL))`, projectID, userID).
|
||||
Order("output.created_at DESC").Find(&items).Error
|
||||
return items, err
|
||||
}
|
||||
|
||||
// RemoveStoryboardImage 清空分镜主图引用,保留生成结果以便在历史记录中恢复。
|
||||
func (s *Creative) RemoveStoryboardImage(userID, projectID, storyboardID uuid.UUID) error {
|
||||
ownedStoryboard := `(project_id=? AND EXISTS (SELECT 1 FROM creative_projects p WHERE p.id=? AND p.user_id=? AND p.deleted_at IS NULL)) OR
|
||||
(episode_id IN (SELECT e.id FROM project_episodes e JOIN creative_projects p ON p.id=e.project_id
|
||||
WHERE e.project_id=? AND p.user_id=? AND e.deleted_at IS NULL AND p.deleted_at IS NULL))`
|
||||
result := s.DB.Model(&model.EpisodeStoryboard{}).
|
||||
Where("id=? AND thumbnail_asset_id IS NOT NULL AND deleted_at IS NULL AND ("+ownedStoryboard+")", storyboardID, projectID, projectID, userID, projectID, userID).
|
||||
Updates(map[string]any{"thumbnail_asset_id": nil, "status": "idle", "user_edited": true})
|
||||
if result.Error != nil {
|
||||
return result.Error
|
||||
}
|
||||
if result.RowsAffected == 0 {
|
||||
var exists int64
|
||||
if err := s.DB.Model(&model.EpisodeStoryboard{}).
|
||||
Where("id=? AND deleted_at IS NULL AND ("+ownedStoryboard+")", storyboardID, projectID, projectID, userID, projectID, userID).
|
||||
Count(&exists).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if exists == 0 {
|
||||
return gorm.ErrRecordNotFound
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *Creative) AddStoryboardCandidate(userID, projectID, storyboardID, outputID uuid.UUID) error {
|
||||
return s.DB.Transaction(func(tx *gorm.DB) error {
|
||||
var row struct{ Candidate bool }
|
||||
if err := tx.Table("generation_outputs output").
|
||||
Select("sb.active_output_id=output.id OR coalesce(output.metadata->>'candidate','false')='true' AS candidate").
|
||||
Joins("JOIN generation_tasks gt ON gt.id=output.task_id AND gt.storyboard_id=? AND gt.task_type='video_generation'", storyboardID).
|
||||
Joins("JOIN episode_storyboards sb ON sb.id=gt.storyboard_id AND sb.deleted_at IS NULL").
|
||||
Joins("JOIN project_episodes e ON e.id=sb.episode_id AND e.deleted_at IS NULL").
|
||||
Joins("JOIN creative_projects p ON p.id=e.project_id AND p.id=? AND p.user_id=? AND p.project_type='premium_drama' AND p.deleted_at IS NULL", projectID, userID).
|
||||
Where("output.id=?", outputID).Take(&row).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if row.Candidate {
|
||||
return nil
|
||||
}
|
||||
var count int64
|
||||
if err := tx.Table("generation_outputs output").
|
||||
Joins("JOIN generation_tasks gt ON gt.id=output.task_id AND gt.storyboard_id=? AND gt.task_type='video_generation'", storyboardID).
|
||||
Joins("JOIN episode_storyboards sb ON sb.id=gt.storyboard_id").
|
||||
Where("sb.active_output_id=output.id OR coalesce(output.metadata->>'candidate','false')='true'").Count(&count).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if count >= 3 {
|
||||
return errors.New("当前分镜最多保留 3 个备选视频,请先删除一个备选视频")
|
||||
}
|
||||
var activeTasks int64
|
||||
if err := tx.Model(&model.GenerationTask{}).
|
||||
Where("storyboard_id=? AND task_type='video_generation' AND status IN ?", storyboardID, activeTaskStatuses).
|
||||
Count(&activeTasks).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if count+activeTasks >= 3 {
|
||||
return errors.New("生成中的视频已占用备选位置,请等待生成完成或先删除一个备选视频")
|
||||
}
|
||||
return tx.Model(&model.GenerationOutput{}).Where("id=?", outputID).
|
||||
Update("metadata", gorm.Expr("jsonb_set(coalesce(metadata,'{}'::jsonb),'{candidate}','true'::jsonb,true)")).Error
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Creative) RemoveStoryboardCandidate(userID, projectID, storyboardID, outputID uuid.UUID) error {
|
||||
return s.DB.Transaction(func(tx *gorm.DB) error {
|
||||
var row struct {
|
||||
Active bool
|
||||
Candidate bool
|
||||
}
|
||||
if err := tx.Table("generation_outputs output").
|
||||
Select("sb.active_output_id=output.id AS active,sb.active_output_id=output.id OR coalesce(output.metadata->>'candidate','false')='true' AS candidate").
|
||||
Joins("JOIN generation_tasks gt ON gt.id=output.task_id AND gt.storyboard_id=? AND gt.task_type='video_generation'", storyboardID).
|
||||
Joins("JOIN episode_storyboards sb ON sb.id=gt.storyboard_id AND sb.deleted_at IS NULL").
|
||||
Joins("JOIN project_episodes e ON e.id=sb.episode_id AND e.deleted_at IS NULL").
|
||||
Joins("JOIN creative_projects p ON p.id=e.project_id AND p.id=? AND p.user_id=? AND p.project_type='premium_drama' AND p.deleted_at IS NULL", projectID, userID).
|
||||
Where("output.id=?", outputID).Take(&row).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if !row.Candidate {
|
||||
return errors.New("该视频不在备选列表中")
|
||||
}
|
||||
if row.Active {
|
||||
var replacement struct{ ID uuid.UUID }
|
||||
if err := tx.Table("generation_outputs output").Select("output.id").
|
||||
Joins("JOIN generation_tasks gt ON gt.id=output.task_id AND gt.storyboard_id=? AND gt.task_type='video_generation'", storyboardID).
|
||||
Where("output.id<>? AND coalesce(output.metadata->>'candidate','false')='true'", outputID).
|
||||
Order("output.created_at DESC").Limit(1).Scan(&replacement).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
var replacementID *uuid.UUID
|
||||
if replacement.ID != uuid.Nil {
|
||||
replacementID = &replacement.ID
|
||||
}
|
||||
status := "idle"
|
||||
if replacementID != nil {
|
||||
status = "completed"
|
||||
}
|
||||
if err := tx.Model(&model.EpisodeStoryboard{}).Where("id=?", storyboardID).
|
||||
Updates(map[string]any{"active_output_id": replacementID, "status": status}).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return tx.Model(&model.GenerationOutput{}).Where("id=?", outputID).
|
||||
Update("metadata", gorm.Expr("coalesce(metadata,'{}'::jsonb)-'candidate'")).Error
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Creative) ActivateStoryboardOutput(userID, projectID, storyboardID, outputID uuid.UUID) error {
|
||||
var output struct {
|
||||
MediaAssetID uuid.UUID
|
||||
TaskType string
|
||||
}
|
||||
if err := s.DB.Table("generation_outputs output").Select("output.media_asset_id,gt.task_type").
|
||||
Joins("JOIN generation_tasks gt ON gt.id=output.task_id AND gt.storyboard_id=? AND (gt.task_type='video_generation' OR (gt.task_type='image_generation' AND gt.input_data->>'target_type'='storyboard'))", storyboardID).
|
||||
Joins("JOIN episode_storyboards sb ON sb.id=gt.storyboard_id AND sb.deleted_at IS NULL").
|
||||
Joins(`JOIN creative_projects p ON p.id=? AND p.user_id=? AND p.deleted_at IS NULL AND
|
||||
(sb.project_id=p.id OR EXISTS (SELECT 1 FROM project_episodes e WHERE e.id=sb.episode_id AND e.project_id=p.id AND e.deleted_at IS NULL))`, projectID, userID).
|
||||
Where("output.id=?", outputID).Take(&output).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
updates := map[string]any{"status": "completed", "updated_at": time.Now()}
|
||||
if output.TaskType == "video_generation" {
|
||||
updates["active_output_id"] = outputID
|
||||
} else {
|
||||
updates["thumbnail_asset_id"] = output.MediaAssetID
|
||||
}
|
||||
result := s.DB.Model(&model.EpisodeStoryboard{}).Where("id=? AND deleted_at IS NULL", storyboardID).Updates(updates)
|
||||
if result.Error != nil {
|
||||
return result.Error
|
||||
}
|
||||
if result.RowsAffected == 0 {
|
||||
return gorm.ErrRecordNotFound
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *Creative) DeleteStoryboardOutput(userID, projectID, storyboardID, outputID uuid.UUID) (string, error) {
|
||||
var objectKey string
|
||||
err := s.DB.Transaction(func(tx *gorm.DB) error {
|
||||
var row struct {
|
||||
TaskID uuid.UUID
|
||||
MediaID uuid.UUID
|
||||
ObjectKey string
|
||||
Status string
|
||||
Candidate bool
|
||||
}
|
||||
if err := tx.Table("generation_outputs output").Select("output.task_id,output.media_asset_id AS media_id,media.object_key,gt.status,sb.active_output_id=output.id OR coalesce(output.metadata->>'candidate','false')='true' AS candidate").
|
||||
Joins("JOIN generation_tasks gt ON gt.id=output.task_id AND gt.storyboard_id=? AND (gt.task_type='video_generation' OR (gt.task_type='image_generation' AND gt.input_data->>'target_type'='storyboard'))", storyboardID).
|
||||
Joins("JOIN media_assets media ON media.id=output.media_asset_id").
|
||||
Joins("JOIN episode_storyboards sb ON sb.id=gt.storyboard_id").
|
||||
Joins(`JOIN creative_projects p ON p.id=? AND p.user_id=? AND p.deleted_at IS NULL AND
|
||||
(sb.project_id=p.id OR EXISTS (SELECT 1 FROM project_episodes e WHERE e.id=sb.episode_id AND e.project_id=p.id AND e.deleted_at IS NULL))`, projectID, userID).
|
||||
Where("output.id=?", outputID).Take(&row).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if row.Status == "pending_submission" || row.Status == "submitting" || row.Status == "submitted" || row.Status == "processing" || row.Status == "cancel_requested" {
|
||||
return errors.New("视频仍在生成中,暂时不能删除历史记录")
|
||||
}
|
||||
if row.Candidate {
|
||||
return errors.New("备选视频不能直接删除,请先移入历史记录")
|
||||
}
|
||||
if err := tx.Exec("UPDATE episode_storyboards SET active_output_id=NULL,status='idle' WHERE id=? AND active_output_id=?", storyboardID, outputID).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if err := tx.Exec("UPDATE episode_storyboards SET thumbnail_asset_id=NULL,status='idle' WHERE id=? AND thumbnail_asset_id=?", storyboardID, row.MediaID).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if err := tx.Exec("DELETE FROM generation_outputs WHERE id=?", outputID).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if err := tx.Exec("DELETE FROM generation_tasks WHERE id=?", row.TaskID).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if err := tx.Exec("DELETE FROM channel_asset_cache WHERE media_asset_id=?", row.MediaID).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if err := tx.Exec("DELETE FROM media_assets WHERE id=?", row.MediaID).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
objectKey = row.ObjectKey
|
||||
return nil
|
||||
})
|
||||
return objectKey, err
|
||||
}
|
||||
|
||||
func (s *Creative) ListAssetOutputs(userID, projectID, assetID uuid.UUID) ([]map[string]any, error) {
|
||||
items := make([]map[string]any, 0)
|
||||
err := s.DB.Table("generation_outputs output").
|
||||
Select(`output.id,output.media_asset_id,output.output_type,output.sequence_no,output.metadata,output.created_at,output.task_id,
|
||||
media.public_url,media.mime_type,media.size_bytes,gt.estimated_points::text AS estimated_points,
|
||||
gt.actual_points::text AS actual_points,(a.image_asset_id=output.media_asset_id) AS active`).
|
||||
Joins("JOIN generation_tasks gt ON gt.id=output.task_id AND gt.task_type='image_generation' AND gt.input_data->>'asset_id'=?", assetID.String()).
|
||||
Joins("JOIN media_assets media ON media.id=output.media_asset_id AND media.deleted_at IS NULL").
|
||||
Joins("JOIN project_assets a ON a.id=? AND a.project_id=? AND a.deleted_at IS NULL", assetID, projectID).
|
||||
Joins("JOIN creative_projects p ON p.id=a.project_id AND p.user_id=? AND p.deleted_at IS NULL", userID).
|
||||
Order("output.created_at DESC").Find(&items).Error
|
||||
return items, err
|
||||
}
|
||||
|
||||
func (s *Creative) ActivateAssetOutput(userID, projectID, assetID, outputID uuid.UUID) error {
|
||||
result := s.DB.Exec(`UPDATE project_assets a SET image_asset_id=(SELECT media_asset_id FROM generation_outputs WHERE id=?),updated_at=CURRENT_TIMESTAMP
|
||||
WHERE a.id=? AND a.project_id=? AND a.deleted_at IS NULL AND EXISTS (
|
||||
SELECT 1 FROM generation_outputs output JOIN generation_tasks gt ON gt.id=output.task_id
|
||||
JOIN creative_projects p ON p.id=a.project_id
|
||||
WHERE output.id=? AND gt.task_type='image_generation' AND gt.input_data->>'asset_id'=? AND p.user_id=? AND p.deleted_at IS NULL)`,
|
||||
outputID, assetID, projectID, outputID, assetID.String(), userID)
|
||||
if result.Error != nil {
|
||||
return result.Error
|
||||
}
|
||||
if result.RowsAffected == 0 {
|
||||
return gorm.ErrRecordNotFound
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *Creative) DeleteAssetOutput(userID, projectID, assetID, outputID uuid.UUID) (string, error) {
|
||||
var objectKey string
|
||||
err := s.DB.Transaction(func(tx *gorm.DB) error {
|
||||
var row struct {
|
||||
TaskID uuid.UUID
|
||||
MediaID uuid.UUID
|
||||
ObjectKey string
|
||||
Status string
|
||||
}
|
||||
if err := tx.Table("generation_outputs output").Select("output.task_id,output.media_asset_id AS media_id,media.object_key,gt.status").
|
||||
Joins("JOIN generation_tasks gt ON gt.id=output.task_id AND gt.task_type='image_generation' AND gt.input_data->>'asset_id'=?", assetID.String()).
|
||||
Joins("JOIN media_assets media ON media.id=output.media_asset_id").
|
||||
Joins("JOIN project_assets a ON a.id=? AND a.project_id=? AND a.deleted_at IS NULL", assetID, projectID).
|
||||
Joins("JOIN creative_projects p ON p.id=a.project_id AND p.user_id=? AND p.deleted_at IS NULL", userID).
|
||||
Where("output.id=?", outputID).Take(&row).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
for _, status := range activeTaskStatuses {
|
||||
if status == row.Status {
|
||||
return errors.New("鍥剧墖浠嶅湪鐢熸垚涓紝鏆傛椂涓嶈兘鍒犻櫎鍘嗗彶璁板綍")
|
||||
}
|
||||
}
|
||||
if err := tx.Exec("UPDATE project_assets SET image_asset_id=NULL WHERE id=? AND image_asset_id=?", assetID, row.MediaID).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if err := tx.Exec("DELETE FROM generation_outputs WHERE id=?", outputID).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if err := tx.Exec("DELETE FROM generation_tasks WHERE id=?", row.TaskID).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if err := tx.Exec("DELETE FROM channel_asset_cache WHERE media_asset_id=?", row.MediaID).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if err := tx.Exec("DELETE FROM media_assets WHERE id=?", row.MediaID).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
objectKey = row.ObjectKey
|
||||
return nil
|
||||
})
|
||||
return objectKey, err
|
||||
}
|
||||
|
||||
func (s *Creative) UpdateEpisodeAnalysisSettings(userID, projectID, episodeID uuid.UUID, audioSource, sourceLanguage string) error {
|
||||
if audioSource != "video_audio" && audioSource != "subtitle_file" {
|
||||
return errors.New("音频来源无效")
|
||||
}
|
||||
sourceLanguage = strings.TrimSpace(sourceLanguage)
|
||||
if len(sourceLanguage) > 16 {
|
||||
return errors.New("源语言设置无效")
|
||||
}
|
||||
return s.DB.Transaction(func(tx *gorm.DB) error {
|
||||
var episode model.ProjectEpisode
|
||||
if err := tx.Table("project_episodes e").Select("e.*").
|
||||
Joins("JOIN creative_projects p ON p.id=e.project_id AND p.user_id=? AND p.deleted_at IS NULL", userID).
|
||||
Where("e.id=? AND e.project_id=? AND e.deleted_at IS NULL", episodeID, projectID).Take(&episode).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
// Selecting subtitle mode is persisted even before the file is uploaded.
|
||||
// The workbench disables analysis until a subtitle asset is present.
|
||||
updates := map[string]any{"audio_source": audioSource, "source_language": nil}
|
||||
if sourceLanguage != "" {
|
||||
updates["source_language"] = sourceLanguage
|
||||
}
|
||||
return tx.Model(&episode).Updates(updates).Error
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,108 @@
|
||||
// 本文件负责按创作项目类型筛选当前剧集可下载的视频,并生成 ZIP 文件清单。
|
||||
package service
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"path"
|
||||
"strings"
|
||||
|
||||
"juhe-factory/api/internal/videoarchive"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
// episodeVideoArchiveRow 描述生成视频查询结果及其在分镜中的业务状态。
|
||||
type episodeVideoArchiveRow struct {
|
||||
ObjectKey string
|
||||
MimeType string
|
||||
SequenceNo int
|
||||
Active bool
|
||||
}
|
||||
|
||||
// EpisodeVideoArchive 查询当前剧集视频:短剧创作仅包含主视频和备选视频,视频转绘包含主视频和历史视频。
|
||||
func (s *Creative) EpisodeVideoArchive(userID, projectID, episodeID uuid.UUID) (videoarchive.Manifest, error) {
|
||||
// 校验项目归属并获取项目类型,当前剧集视频 ZIP 仅对短剧创作和视频转绘开放。
|
||||
var project struct {
|
||||
ProjectType string
|
||||
}
|
||||
if err := s.DB.Table("creative_projects").Select("project_type").
|
||||
Where("id=? AND user_id=? AND deleted_at IS NULL", projectID, userID).
|
||||
Take(&project).Error; err != nil {
|
||||
return videoarchive.Manifest{}, err
|
||||
}
|
||||
if project.ProjectType != "premium_drama" && project.ProjectType != "video_redraw" {
|
||||
return videoarchive.Manifest{}, gorm.ErrRecordNotFound
|
||||
}
|
||||
var episode struct {
|
||||
ProjectName string
|
||||
EpisodeName string
|
||||
}
|
||||
if err := s.DB.Table("project_episodes episode").
|
||||
Select("project.name AS project_name,episode.name AS episode_name").
|
||||
Joins("JOIN creative_projects project ON project.id=episode.project_id AND project.deleted_at IS NULL").
|
||||
Where("episode.id=? AND episode.project_id=? AND episode.deleted_at IS NULL", episodeID, projectID).
|
||||
Take(&episode).Error; err != nil {
|
||||
return videoarchive.Manifest{}, err
|
||||
}
|
||||
|
||||
rows := make([]episodeVideoArchiveRow, 0)
|
||||
query := s.DB.Table("generation_outputs output").
|
||||
Select(`media.object_key,media.mime_type,storyboard.sequence_no,
|
||||
storyboard.active_output_id=output.id AS active`).
|
||||
Joins("JOIN generation_tasks task ON task.id=output.task_id AND task.task_type='video_generation'").
|
||||
Joins("JOIN episode_storyboards storyboard ON storyboard.id=task.storyboard_id AND storyboard.deleted_at IS NULL").
|
||||
Joins("JOIN media_assets media ON media.id=output.media_asset_id AND media.deleted_at IS NULL").
|
||||
Where("storyboard.episode_id=?", episodeID)
|
||||
if project.ProjectType == "premium_drama" {
|
||||
query = query.Where("storyboard.active_output_id=output.id OR coalesce(output.metadata->>'candidate','false')='true'")
|
||||
}
|
||||
if err := query.Order("storyboard.sequence_no ASC,(storyboard.active_output_id=output.id) DESC,output.created_at DESC").Find(&rows).Error; err != nil {
|
||||
return videoarchive.Manifest{}, err
|
||||
}
|
||||
if len(rows) == 0 {
|
||||
return videoarchive.Manifest{}, errors.New("当前剧集暂无可下载视频")
|
||||
}
|
||||
|
||||
entries := make([]videoarchive.Entry, 0, len(rows))
|
||||
versions := make(map[int]int)
|
||||
for _, row := range rows {
|
||||
label := "主视频"
|
||||
if !row.Active {
|
||||
versions[row.SequenceNo]++
|
||||
if project.ProjectType == "premium_drama" {
|
||||
label = fmt.Sprintf("备选视频%d", versions[row.SequenceNo])
|
||||
} else {
|
||||
label = fmt.Sprintf("历史视频%d", versions[row.SequenceNo])
|
||||
}
|
||||
}
|
||||
name := fmt.Sprintf("分镜%03d-%s%s", row.SequenceNo, label, videoArchiveExtension(row.ObjectKey, row.MimeType))
|
||||
entries = append(entries, videoarchive.Entry{ObjectKey: row.ObjectKey, Name: videoarchive.SafeFilename(name, "视频.mp4")})
|
||||
}
|
||||
archiveName := videoarchive.SafeFilename(episode.ProjectName+episode.EpisodeName, "剧集视频") + ".zip"
|
||||
return videoarchive.Manifest{Name: archiveName, Entries: entries}, nil
|
||||
}
|
||||
|
||||
// videoArchiveExtension 根据已保存的媒体类型和对象键生成稳定的视频扩展名。
|
||||
func videoArchiveExtension(objectKey, mimeType string) string {
|
||||
switch strings.ToLower(strings.TrimSpace(strings.Split(mimeType, ";")[0])) {
|
||||
case "video/webm":
|
||||
return ".webm"
|
||||
case "video/quicktime":
|
||||
return ".mov"
|
||||
case "video/x-m4v":
|
||||
return ".m4v"
|
||||
case "video/x-matroska":
|
||||
return ".mkv"
|
||||
case "video/x-msvideo":
|
||||
return ".avi"
|
||||
case "video/mp4":
|
||||
return ".mp4"
|
||||
}
|
||||
extension := strings.ToLower(path.Ext(objectKey))
|
||||
if map[string]bool{".mp4": true, ".webm": true, ".mov": true, ".m4v": true, ".mkv": true, ".avi": true}[extension] {
|
||||
return extension
|
||||
}
|
||||
return ".mp4"
|
||||
}
|
||||
@@ -0,0 +1,544 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"archive/zip"
|
||||
"bytes"
|
||||
"crypto/sha256"
|
||||
"encoding/binary"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"encoding/xml"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"path/filepath"
|
||||
"regexp"
|
||||
"sort"
|
||||
"strings"
|
||||
"time"
|
||||
"unicode"
|
||||
"unicode/utf16"
|
||||
"unicode/utf8"
|
||||
|
||||
dramapkg "juhe-factory/api/internal/drama"
|
||||
"juhe-factory/api/internal/model"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/richardlehane/mscfb"
|
||||
"golang.org/x/text/encoding/simplifiedchinese"
|
||||
"golang.org/x/text/transform"
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/clause"
|
||||
)
|
||||
|
||||
const maxDramaImportBytes = 3 * 1024 * 1024
|
||||
|
||||
var chapterTitlePattern = regexp.MustCompile(`(?im)^\s*(?:#+\s*|[*+-]\s*)?((?:第[零〇一二三四五六七八九十百千万两\d]+[章节回卷篇部集幕].*)|(?:chapter\s+\d+.*))\s*$`)
|
||||
|
||||
type DramaChapter struct {
|
||||
Title string `json:"title"`
|
||||
Content string `json:"content,omitempty"`
|
||||
CharCount int `json:"char_count"`
|
||||
}
|
||||
|
||||
type DramaImportPreview struct {
|
||||
ImportToken uuid.UUID `json:"import_token"`
|
||||
HasChapters bool `json:"has_chapters"`
|
||||
TotalCharacters int `json:"total_characters"`
|
||||
Chapters []DramaChapter `json:"chapters"`
|
||||
SingleEpisodeConfirmationNeeded bool `json:"single_episode_confirmation_required"`
|
||||
}
|
||||
|
||||
func (s *Creative) PreviewDramaImport(userID, projectID uuid.UUID, filename string, data []byte, rawText string) (*DramaImportPreview, error) {
|
||||
if err := s.requirePremiumProject(userID, projectID); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if len(data) > maxDramaImportBytes || len([]byte(rawText)) > maxDramaImportBytes {
|
||||
return nil, errors.New("小说文件或文本不能超过3MB")
|
||||
}
|
||||
text, sourceType, err := parseDramaSource(filename, data, rawText)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
text = normalizeDramaText(text)
|
||||
if text == "" {
|
||||
return nil, errors.New("未读取到有效小说正文")
|
||||
}
|
||||
chapters := splitDramaChapters(text)
|
||||
metadata := make([]DramaChapter, 0, len(chapters))
|
||||
for _, chapter := range chapters {
|
||||
metadata = append(metadata, DramaChapter{Title: chapter.Title, CharCount: len([]rune(chapter.Content))})
|
||||
}
|
||||
previewJSON, _ := json.Marshal(map[string]any{"has_chapters": len(chapters) > 0, "chapters": metadata})
|
||||
hash := sha256.Sum256([]byte(text))
|
||||
session := model.DramaImportSession{
|
||||
ID: uuid.New(), UserID: userID, ProjectID: projectID, SourceType: sourceType,
|
||||
SourceFilename: cleanDramaFilename(filename), RawContent: text, ContentSHA256: hex.EncodeToString(hash[:]),
|
||||
PreviewData: previewJSON, ExpiresAt: time.Now().Add(30 * time.Minute),
|
||||
}
|
||||
if err := s.DB.Create(&session).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &DramaImportPreview{
|
||||
ImportToken: session.ID, HasChapters: len(chapters) > 0, TotalCharacters: len([]rune(text)),
|
||||
Chapters: metadata, SingleEpisodeConfirmationNeeded: len(chapters) == 0,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (s *Creative) ConfirmDramaImport(userID, projectID, token uuid.UUID, chaptersPerEpisode, startEpisodeNo int, treatAsSingle bool) ([]model.ProjectEpisode, error) {
|
||||
if chaptersPerEpisode < 1 {
|
||||
chaptersPerEpisode = 1
|
||||
}
|
||||
if startEpisodeNo < 1 {
|
||||
return nil, errors.New("起始集数无效")
|
||||
}
|
||||
created := make([]model.ProjectEpisode, 0)
|
||||
err := s.DB.Transaction(func(tx *gorm.DB) error {
|
||||
var project model.CreativeProject
|
||||
if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Select("id").Where("id=? AND user_id=? AND deleted_at IS NULL", projectID, userID).Take(&project).Error; err != nil {
|
||||
return gorm.ErrRecordNotFound
|
||||
}
|
||||
var session model.DramaImportSession
|
||||
if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Where("id=? AND user_id=? AND project_id=?", token, userID, projectID).Take(&session).Error; err != nil {
|
||||
return gorm.ErrRecordNotFound
|
||||
}
|
||||
if session.ConsumedAt != nil || time.Now().After(session.ExpiresAt) {
|
||||
return errors.New("导入预览已失效,请重新上传")
|
||||
}
|
||||
chapters := splitDramaChapters(session.RawContent)
|
||||
if len(chapters) == 0 {
|
||||
if !treatAsSingle {
|
||||
return errors.New("需要确认将全部内容作为一集导入")
|
||||
}
|
||||
chapters = []DramaChapter{{Title: fmt.Sprintf("第%d集", startEpisodeNo), Content: session.RawContent, CharCount: len([]rune(session.RawContent))}}
|
||||
chaptersPerEpisode = 1
|
||||
}
|
||||
for offset, index := 0, 0; index < len(chapters); offset, index = offset+1, index+chaptersPerEpisode {
|
||||
end := index + chaptersPerEpisode
|
||||
if end > len(chapters) {
|
||||
end = len(chapters)
|
||||
}
|
||||
group := chapters[index:end]
|
||||
episodeNo := startEpisodeNo + offset
|
||||
name := fmt.Sprintf("第%d集", episodeNo)
|
||||
if len(group) == 1 && strings.TrimSpace(group[0].Title) != "" {
|
||||
name = group[0].Title
|
||||
}
|
||||
name = truncateRunes(strings.TrimSpace(name), 160)
|
||||
exists, err := episodeNameExists(tx, projectID, nil, name)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if exists {
|
||||
return errors.New("同一项目下不可有同名剧集")
|
||||
}
|
||||
parts := make([]string, 0, len(group))
|
||||
for _, chapter := range group {
|
||||
parts = append(parts, strings.TrimSpace(chapter.Content))
|
||||
}
|
||||
content := strings.TrimSpace(strings.Join(parts, "\n\n"))
|
||||
if err := dramapkg.ValidateEpisodeContent(content); err != nil {
|
||||
return fmt.Errorf("第%d集:%w", episodeNo, err)
|
||||
}
|
||||
episode := model.ProjectEpisode{ID: uuid.New(), ProjectID: projectID, EpisodeNo: episodeNo, Name: name, AudioSource: "video_audio", Status: "draft"}
|
||||
if err := tx.Create(&episode).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
hash := sha256.Sum256([]byte(content))
|
||||
source := model.EpisodeSource{ID: uuid.New(), EpisodeID: episode.ID, Title: episode.Name, RawContent: content, CharCount: len([]rune(content)), SourceType: session.SourceType, SourceFilename: session.SourceFilename, ContentSHA256: hex.EncodeToString(hash[:])}
|
||||
if err := tx.Create(&source).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
created = append(created, episode)
|
||||
}
|
||||
now := time.Now()
|
||||
return tx.Model(&session).Updates(map[string]any{"consumed_at": now, "raw_content": ""}).Error
|
||||
})
|
||||
return created, err
|
||||
}
|
||||
|
||||
func (s *Creative) EpisodeSource(userID, projectID, episodeID uuid.UUID) (*model.EpisodeSource, error) {
|
||||
var source model.EpisodeSource
|
||||
result := s.DB.Table("episode_sources source").Joins("JOIN project_episodes episode ON episode.id=source.episode_id AND episode.deleted_at IS NULL").Joins("JOIN creative_projects project ON project.id=episode.project_id AND project.deleted_at IS NULL").Where("source.episode_id=? AND episode.project_id=? AND project.user_id=?", episodeID, projectID, userID).Limit(1).Find(&source)
|
||||
if result.Error != nil {
|
||||
return nil, result.Error
|
||||
}
|
||||
if result.RowsAffected == 0 {
|
||||
return nil, gorm.ErrRecordNotFound
|
||||
}
|
||||
return &source, nil
|
||||
}
|
||||
|
||||
func (s *Creative) SaveEpisodeSource(userID, projectID, episodeID uuid.UUID, content, title string) (*model.EpisodeSource, error) {
|
||||
content = normalizeDramaText(content)
|
||||
if err := dramapkg.ValidateEpisodeContent(content); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := s.requirePremiumEpisode(userID, projectID, episodeID); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
hash := sha256.Sum256([]byte(content))
|
||||
source := model.EpisodeSource{ID: uuid.New(), EpisodeID: episodeID, Title: truncateRunes(strings.TrimSpace(title), 200), RawContent: content, CharCount: len([]rune(content)), SourceType: "pasted_text", ContentSHA256: hex.EncodeToString(hash[:])}
|
||||
err := s.DB.Clauses(clause.OnConflict{Columns: []clause.Column{{Name: "episode_id"}}, DoUpdates: clause.AssignmentColumns([]string{"title", "raw_content", "char_count", "source_type", "source_filename", "content_sha256", "updated_at"})}).Create(&source).Error
|
||||
return &source, err
|
||||
}
|
||||
|
||||
func (s *Creative) requirePremiumProject(userID, projectID uuid.UUID) error {
|
||||
var count int64
|
||||
err := s.DB.Table("creative_projects").Where("id=? AND user_id=? AND project_type='premium_drama' AND deleted_at IS NULL", projectID, userID).Count(&count).Error
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if count == 0 {
|
||||
return gorm.ErrRecordNotFound
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *Creative) requirePremiumEpisode(userID, projectID, episodeID uuid.UUID) error {
|
||||
var count int64
|
||||
err := s.DB.Table("project_episodes episode").Joins("JOIN creative_projects project ON project.id=episode.project_id").Where("episode.id=? AND episode.project_id=? AND episode.deleted_at IS NULL AND project.user_id=? AND project.project_type='premium_drama' AND project.deleted_at IS NULL", episodeID, projectID, userID).Count(&count).Error
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if count == 0 {
|
||||
return gorm.ErrRecordNotFound
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func parseDramaSource(filename string, data []byte, rawText string) (string, string, error) {
|
||||
if strings.TrimSpace(rawText) != "" {
|
||||
if len(data) > 0 {
|
||||
return "", "", errors.New("文件和粘贴文本只能选择一种")
|
||||
}
|
||||
if looksLikeBinary([]byte(rawText)) {
|
||||
return "", "", errors.New("粘贴内容包含二进制控制字符")
|
||||
}
|
||||
if looksLikeHTMLOrScript([]byte(rawText)) {
|
||||
return "", "", errors.New("不支持HTML或脚本内容")
|
||||
}
|
||||
return rawText, "pasted_text", nil
|
||||
}
|
||||
if len(data) == 0 {
|
||||
return "", "", errors.New("请选择小说文件或粘贴文本")
|
||||
}
|
||||
ext := strings.ToLower(filepath.Ext(filename))
|
||||
switch ext {
|
||||
case ".txt":
|
||||
text, err := decodePlainText(data)
|
||||
return text, "txt", err
|
||||
case ".docx":
|
||||
text, err := extractDOCX(data)
|
||||
return text, "docx", err
|
||||
case ".doc":
|
||||
text, err := extractLegacyDOC(data)
|
||||
return text, "doc", err
|
||||
default:
|
||||
return "", "", errors.New("仅支持TXT、DOC和DOCX文件")
|
||||
}
|
||||
}
|
||||
|
||||
func decodePlainText(data []byte) (string, error) {
|
||||
if hasDisallowedBinarySignature(data) || bytes.IndexByte(data, 0) >= 0 || looksLikeBinary(data) {
|
||||
return "", errors.New("TXT文件包含二进制内容")
|
||||
}
|
||||
if looksLikeHTMLOrScript(data) {
|
||||
return "", errors.New("不支持HTML或脚本文件")
|
||||
}
|
||||
data = bytes.TrimPrefix(data, []byte{0xEF, 0xBB, 0xBF})
|
||||
if utf8.Valid(data) {
|
||||
return string(data), nil
|
||||
}
|
||||
decoded, err := io.ReadAll(transform.NewReader(bytes.NewReader(data), simplifiedchinese.GB18030.NewDecoder()))
|
||||
if err != nil || !utf8.Valid(decoded) {
|
||||
return "", errors.New("TXT文件编码不支持,请使用UTF-8或GB18030")
|
||||
}
|
||||
return string(decoded), nil
|
||||
}
|
||||
|
||||
func extractDOCX(data []byte) (string, error) {
|
||||
if len(data) < 4 || !bytes.Equal(data[:4], []byte{'P', 'K', 3, 4}) {
|
||||
return "", errors.New("DOCX文件结构无效")
|
||||
}
|
||||
reader, err := zip.NewReader(bytes.NewReader(data), int64(len(data)))
|
||||
if err != nil {
|
||||
return "", errors.New("DOCX文件结构无效")
|
||||
}
|
||||
if len(reader.File) > 1000 {
|
||||
return "", errors.New("DOCX文件包含过多内容")
|
||||
}
|
||||
var document, contentTypes []byte
|
||||
var expanded uint64
|
||||
for _, file := range reader.File {
|
||||
expanded += file.UncompressedSize64
|
||||
if expanded > 20*1024*1024 {
|
||||
return "", errors.New("DOCX展开后内容过大")
|
||||
}
|
||||
name := strings.ToLower(strings.ReplaceAll(file.Name, "\\", "/"))
|
||||
if strings.Contains(name, "vbaproject") || strings.Contains(name, "activex") || strings.Contains(name, "embeddings/") {
|
||||
return "", errors.New("DOCX包含宏、ActiveX或嵌入对象")
|
||||
}
|
||||
if name == "[content_types].xml" {
|
||||
stream, openErr := file.Open()
|
||||
if openErr != nil {
|
||||
return "", openErr
|
||||
}
|
||||
contentTypes, err = io.ReadAll(io.LimitReader(stream, 1024*1024))
|
||||
stream.Close()
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
}
|
||||
if name == "word/document.xml" {
|
||||
stream, openErr := file.Open()
|
||||
if openErr != nil {
|
||||
return "", openErr
|
||||
}
|
||||
document, err = io.ReadAll(io.LimitReader(stream, 12*1024*1024))
|
||||
stream.Close()
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
}
|
||||
}
|
||||
contentTypeText := strings.ToLower(string(contentTypes))
|
||||
if len(document) == 0 || !strings.Contains(contentTypeText, "wordprocessingml.document.main+xml") || strings.Contains(contentTypeText, "macroenabled") {
|
||||
return "", errors.New("DOCX缺少Word正文结构")
|
||||
}
|
||||
decoder := xml.NewDecoder(bytes.NewReader(document))
|
||||
var output strings.Builder
|
||||
for {
|
||||
token, tokenErr := decoder.Token()
|
||||
if tokenErr == io.EOF {
|
||||
break
|
||||
}
|
||||
if tokenErr != nil {
|
||||
return "", errors.New("DOCX正文XML无效")
|
||||
}
|
||||
switch value := token.(type) {
|
||||
case xml.CharData:
|
||||
output.Write([]byte(value))
|
||||
case xml.EndElement:
|
||||
if value.Name.Local == "p" {
|
||||
output.WriteString("\n")
|
||||
} else if value.Name.Local == "tab" {
|
||||
output.WriteString("\t")
|
||||
}
|
||||
}
|
||||
}
|
||||
return output.String(), nil
|
||||
}
|
||||
|
||||
func extractLegacyDOC(data []byte) (string, error) {
|
||||
magic := []byte{0xD0, 0xCF, 0x11, 0xE0, 0xA1, 0xB1, 0x1A, 0xE1}
|
||||
if len(data) < len(magic) || !bytes.Equal(data[:len(magic)], magic) {
|
||||
return "", errors.New("DOC文件结构无效")
|
||||
}
|
||||
reader, err := mscfb.New(bytes.NewReader(data))
|
||||
if err != nil {
|
||||
return "", errors.New("DOC文件结构无效")
|
||||
}
|
||||
streams := map[string][]byte{}
|
||||
for entry, nextErr := reader.Next(); nextErr == nil; entry, nextErr = reader.Next() {
|
||||
name := strings.ToLower(entry.Name)
|
||||
if strings.Contains(name, "vba") || strings.Contains(name, "macros") || strings.Contains(name, "objectpool") {
|
||||
return "", errors.New("DOC包含宏或嵌入对象")
|
||||
}
|
||||
if name == "worddocument" || name == "0table" || name == "1table" {
|
||||
content, readErr := io.ReadAll(io.LimitReader(entry, 12*1024*1024))
|
||||
if readErr != nil {
|
||||
return "", readErr
|
||||
}
|
||||
streams[name] = content
|
||||
}
|
||||
}
|
||||
word := streams["worddocument"]
|
||||
if len(word) < 0x1AA || binary.LittleEndian.Uint16(word[:2]) != 0xA5EC {
|
||||
return "", errors.New("DOC缺少有效WordDocument流")
|
||||
}
|
||||
flags := binary.LittleEndian.Uint16(word[0x0A:0x0C])
|
||||
tableName := "0table"
|
||||
if flags&0x0200 != 0 {
|
||||
tableName = "1table"
|
||||
}
|
||||
table := streams[tableName]
|
||||
fcClx := int(binary.LittleEndian.Uint32(word[0x1A2:0x1A6]))
|
||||
lcbClx := int(binary.LittleEndian.Uint32(word[0x1A6:0x1AA]))
|
||||
text, parseErr := extractDOCPieces(word, table, fcClx, lcbClx)
|
||||
if parseErr != nil || strings.TrimSpace(text) == "" {
|
||||
return "", errors.New("DOC正文无法安全解析")
|
||||
}
|
||||
return text, nil
|
||||
}
|
||||
|
||||
func extractDOCPieces(word, table []byte, offset, length int) (string, error) {
|
||||
if offset < 0 || length <= 0 || offset+length > len(table) {
|
||||
return "", errors.New("invalid CLX")
|
||||
}
|
||||
clx := table[offset : offset+length]
|
||||
pos := 0
|
||||
for pos < len(clx) && clx[pos] == 0x01 {
|
||||
if pos+3 > len(clx) {
|
||||
return "", io.ErrUnexpectedEOF
|
||||
}
|
||||
size := int(binary.LittleEndian.Uint16(clx[pos+1 : pos+3]))
|
||||
pos += 3 + size
|
||||
}
|
||||
if pos+5 > len(clx) || clx[pos] != 0x02 {
|
||||
return "", errors.New("missing Pcdt")
|
||||
}
|
||||
plcSize := int(binary.LittleEndian.Uint32(clx[pos+1 : pos+5]))
|
||||
pos += 5
|
||||
if plcSize < 4 || pos+plcSize > len(clx) {
|
||||
return "", errors.New("invalid PlcPcd")
|
||||
}
|
||||
plc := clx[pos : pos+plcSize]
|
||||
pieces := (plcSize - 4) / 12
|
||||
if pieces <= 0 || pieces > 100000 {
|
||||
return "", errors.New("invalid piece count")
|
||||
}
|
||||
cpBytes := (pieces + 1) * 4
|
||||
if cpBytes+pieces*8 > len(plc) {
|
||||
return "", io.ErrUnexpectedEOF
|
||||
}
|
||||
var output strings.Builder
|
||||
for i := 0; i < pieces; i++ {
|
||||
cpStart := int(binary.LittleEndian.Uint32(plc[i*4 : i*4+4]))
|
||||
cpEnd := int(binary.LittleEndian.Uint32(plc[(i+1)*4 : (i+1)*4+4]))
|
||||
if cpEnd <= cpStart {
|
||||
continue
|
||||
}
|
||||
pcd := cpBytes + i*8
|
||||
rawFC := binary.LittleEndian.Uint32(plc[pcd+2 : pcd+6])
|
||||
compressed := rawFC&0x40000000 != 0
|
||||
fc := int(rawFC & 0x3FFFFFFF)
|
||||
chars := cpEnd - cpStart
|
||||
if compressed {
|
||||
fc /= 2
|
||||
if fc < 0 || fc+chars > len(word) {
|
||||
continue
|
||||
}
|
||||
decoded, _ := io.ReadAll(transform.NewReader(bytes.NewReader(word[fc:fc+chars]), simplifiedchinese.GB18030.NewDecoder()))
|
||||
output.Write(decoded)
|
||||
} else {
|
||||
byteLen := chars * 2
|
||||
if fc < 0 || fc+byteLen > len(word) {
|
||||
continue
|
||||
}
|
||||
units := make([]uint16, chars)
|
||||
for j := 0; j < chars; j++ {
|
||||
units[j] = binary.LittleEndian.Uint16(word[fc+j*2 : fc+j*2+2])
|
||||
}
|
||||
output.WriteString(string(utf16.Decode(units)))
|
||||
}
|
||||
}
|
||||
return output.String(), nil
|
||||
}
|
||||
|
||||
func splitDramaChapters(content string) []DramaChapter {
|
||||
matches := chapterTitlePattern.FindAllStringSubmatchIndex(content, -1)
|
||||
if len(matches) == 0 {
|
||||
return nil
|
||||
}
|
||||
chapters := make([]DramaChapter, 0, len(matches))
|
||||
for i, match := range matches {
|
||||
start := match[0]
|
||||
if i == 0 {
|
||||
start = 0
|
||||
}
|
||||
end := len(content)
|
||||
if i+1 < len(matches) {
|
||||
end = matches[i+1][0]
|
||||
}
|
||||
body := strings.TrimSpace(content[start:end])
|
||||
if body == "" {
|
||||
continue
|
||||
}
|
||||
title := strings.TrimSpace(content[match[2]:match[3]])
|
||||
chapters = append(chapters, DramaChapter{Title: title, Content: body, CharCount: len([]rune(body))})
|
||||
}
|
||||
return chapters
|
||||
}
|
||||
|
||||
func normalizeDramaText(value string) string {
|
||||
value = strings.ReplaceAll(value, "\x00", "")
|
||||
value = strings.ReplaceAll(strings.ReplaceAll(value, "\r\n", "\n"), "\r", "\n")
|
||||
lines := strings.Split(value, "\n")
|
||||
cleaned := make([]string, 0, len(lines))
|
||||
blank := false
|
||||
for _, line := range lines {
|
||||
line = strings.TrimRightFunc(line, unicode.IsSpace)
|
||||
if strings.TrimSpace(line) == "" {
|
||||
if !blank {
|
||||
cleaned = append(cleaned, "")
|
||||
blank = true
|
||||
}
|
||||
continue
|
||||
}
|
||||
cleaned = append(cleaned, line)
|
||||
blank = false
|
||||
}
|
||||
return strings.TrimSpace(strings.Join(cleaned, "\n"))
|
||||
}
|
||||
|
||||
func looksLikeBinary(data []byte) bool {
|
||||
if len(data) == 0 {
|
||||
return false
|
||||
}
|
||||
controls := 0
|
||||
for _, b := range data {
|
||||
if b < 0x09 || (b > 0x0D && b < 0x20) {
|
||||
controls++
|
||||
}
|
||||
}
|
||||
return controls*100/len(data) > 2
|
||||
}
|
||||
func looksLikeHTMLOrScript(data []byte) bool {
|
||||
sample := strings.ToLower(string(data))
|
||||
markers := []string{"<!doctype html", "<html", "<head", "<body", "<script", "</script", "<iframe", "<object", "<embed", "javascript:", "vbscript:", "onerror=", "onload="}
|
||||
for _, marker := range markers {
|
||||
if strings.Contains(sample, marker) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
func hasDisallowedBinarySignature(data []byte) bool {
|
||||
signatures := [][]byte{
|
||||
{0x25, 0x50, 0x44, 0x46, 0x2D},
|
||||
{0x50, 0x4B, 0x03, 0x04},
|
||||
{0xD0, 0xCF, 0x11, 0xE0, 0xA1, 0xB1, 0x1A, 0xE1},
|
||||
{0x89, 0x50, 0x4E, 0x47},
|
||||
{0xFF, 0xD8, 0xFF},
|
||||
{0x47, 0x49, 0x46, 0x38},
|
||||
{0x4D, 0x5A},
|
||||
{0x7F, 0x45, 0x4C, 0x46},
|
||||
{0x52, 0x61, 0x72, 0x21},
|
||||
{0x37, 0x7A, 0xBC, 0xAF, 0x27, 0x1C},
|
||||
{0x1F, 0x8B},
|
||||
}
|
||||
data = bytes.TrimSpace(data)
|
||||
for _, signature := range signatures {
|
||||
if bytes.HasPrefix(data, signature) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
func cleanDramaFilename(value string) string {
|
||||
value = filepath.Base(strings.TrimSpace(value))
|
||||
return truncateRunes(value, 255)
|
||||
}
|
||||
func truncateRunes(value string, limit int) string {
|
||||
runes := []rune(value)
|
||||
if len(runes) > limit {
|
||||
return string(runes[:limit])
|
||||
}
|
||||
return value
|
||||
}
|
||||
|
||||
// Stable chapter ordering is useful in tests and when import previews are reconstructed.
|
||||
func sortEpisodesByNumber(items []model.ProjectEpisode) {
|
||||
sort.SliceStable(items, func(i, j int) bool { return items[i].EpisodeNo < items[j].EpisodeNo })
|
||||
}
|
||||
@@ -0,0 +1,95 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"archive/zip"
|
||||
"bytes"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestParseDramaSourceRejectsDisguisedAndExecutableContent(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
filename string
|
||||
data []byte
|
||||
raw string
|
||||
}{
|
||||
{name: "pdf renamed txt", filename: "novel.txt", data: []byte("%PDF-1.7\ncontent")},
|
||||
{name: "zip renamed txt", filename: "novel.txt", data: []byte{'P', 'K', 3, 4, 1, 2, 3}},
|
||||
{name: "html after long prefix", filename: "novel.txt", data: []byte(strings.Repeat("正文", 3000) + "<script>alert(1)</script>")},
|
||||
{name: "script pasted", raw: "小说正文\njavascript:alert(1)"},
|
||||
{name: "binary pasted", raw: "正文\x00内容"},
|
||||
}
|
||||
for _, test := range cases {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
if _, _, err := parseDramaSource(test.filename, test.data, test.raw); err == nil {
|
||||
t.Fatal("expected unsafe import to be rejected")
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestExtractDOCXReadsParagraphsAndValidatesWordPackage(t *testing.T) {
|
||||
valid := makeTestDOCX(t, `<?xml version="1.0"?><Types xmlns="http://schemas.openxmlformats.org/package/2006/content-types"><Override PartName="/word/document.xml" ContentType="application/vnd.openxmlformats-officedocument.wordprocessingml.document.main+xml"/></Types>`, nil)
|
||||
text, err := extractDOCX(valid)
|
||||
if err != nil {
|
||||
t.Fatalf("extractDOCX returned error: %v", err)
|
||||
}
|
||||
if !strings.Contains(text, "第一章 开始") || !strings.Contains(text, "这是正文") {
|
||||
t.Fatalf("unexpected extracted text: %q", text)
|
||||
}
|
||||
|
||||
fake := makeTestDOCX(t, `<?xml version="1.0"?><Types/>`, nil)
|
||||
if _, err := extractDOCX(fake); err == nil {
|
||||
t.Fatal("expected ZIP with fake DOCX structure to be rejected")
|
||||
}
|
||||
|
||||
macro := makeTestDOCX(t, `<?xml version="1.0"?><Types xmlns="http://schemas.openxmlformats.org/package/2006/content-types"><Override PartName="/word/document.xml" ContentType="application/vnd.openxmlformats-officedocument.wordprocessingml.document.main+xml"/></Types>`, map[string]string{"word/vbaProject.bin": "macro"})
|
||||
if _, err := extractDOCX(macro); err == nil {
|
||||
t.Fatal("expected macro-enabled package to be rejected")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSplitDramaChaptersPreservesOriginalText(t *testing.T) {
|
||||
source := "作品说明\n\n第一章 初见\n正文一\n\n第二章 重逢\n正文二"
|
||||
chapters := splitDramaChapters(source)
|
||||
if len(chapters) != 2 {
|
||||
t.Fatalf("expected 2 chapters, got %d", len(chapters))
|
||||
}
|
||||
joined := strings.TrimSpace(chapters[0].Content + "\n\n" + chapters[1].Content)
|
||||
if joined != source {
|
||||
t.Fatalf("chapter splitting lost or duplicated text:\nwant: %q\n got: %q", source, joined)
|
||||
}
|
||||
if chapters[0].Title != "第一章 初见" || chapters[1].Title != "第二章 重逢" {
|
||||
t.Fatalf("unexpected chapter titles: %#v", chapters)
|
||||
}
|
||||
if chapters := splitDramaChapters("没有章节标题的完整正文"); chapters != nil {
|
||||
t.Fatalf("text without chapter titles must require single-episode confirmation: %#v", chapters)
|
||||
}
|
||||
}
|
||||
|
||||
func makeTestDOCX(t *testing.T, contentTypes string, extras map[string]string) []byte {
|
||||
t.Helper()
|
||||
var buffer bytes.Buffer
|
||||
writer := zip.NewWriter(&buffer)
|
||||
files := map[string]string{
|
||||
"[Content_Types].xml": contentTypes,
|
||||
"word/document.xml": `<?xml version="1.0"?><w:document xmlns:w="http://schemas.openxmlformats.org/wordprocessingml/2006/main"><w:body><w:p><w:r><w:t>第一章 开始</w:t></w:r></w:p><w:p><w:r><w:t>这是正文</w:t></w:r></w:p></w:body></w:document>`,
|
||||
}
|
||||
for name, content := range extras {
|
||||
files[name] = content
|
||||
}
|
||||
for name, content := range files {
|
||||
entry, err := writer.Create(name)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err = entry.Write([]byte(content)); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
if err := writer.Close(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return buffer.Bytes()
|
||||
}
|
||||
@@ -0,0 +1,232 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"time"
|
||||
|
||||
"juhe-factory/api/internal/billing"
|
||||
dramapkg "juhe-factory/api/internal/drama"
|
||||
"juhe-factory/api/internal/model"
|
||||
queuepkg "juhe-factory/api/internal/queue"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
type dramaTextModel struct {
|
||||
ModelID uuid.UUID
|
||||
ChannelID uuid.UUID
|
||||
ModelName string
|
||||
ChannelName string
|
||||
Pricing billing.TextPricing `gorm:"-"`
|
||||
}
|
||||
|
||||
func (s *Creative) QueueDramaParse(userID, projectID, episodeID uuid.UUID) (*model.DramaParseTask, error) {
|
||||
if s.Queue == nil {
|
||||
return nil, errors.New("解析任务队列不可用")
|
||||
}
|
||||
if err := s.requirePremiumEpisode(userID, projectID, episodeID); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var source model.EpisodeSource
|
||||
if err := s.DB.Where("episode_id=?", episodeID).Take(&source).Error; err != nil {
|
||||
return nil, errors.New("请先导入或填写剧集原文")
|
||||
}
|
||||
if err := dramapkg.ValidateEpisodeContent(source.RawContent); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
configured, err := s.dramaTextModel(projectID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
prompt := dramaParsePrompt()
|
||||
snapshot, _ := json.Marshal(map[string]any{"model_id": configured.ModelID, "model_name": configured.ModelName, "channel_id": configured.ChannelID, "channel_name": configured.ChannelName})
|
||||
task := &model.DramaParseTask{
|
||||
ID: uuid.New(), RequestID: uuid.NewString(), UserID: userID, ProjectID: projectID,
|
||||
EpisodeID: episodeID, ChannelID: configured.ChannelID, ModelID: configured.ModelID,
|
||||
SourceSHA256: source.ContentSHA256, Status: "queued", ModelSnapshot: snapshot, PromptSnapshot: prompt,
|
||||
ContextSnapshot: json.RawMessage("{}"),
|
||||
BillingSnapshot: configured.Pricing.MarshalSnapshot(),
|
||||
EstimatedPoints: "0.00", PrepaidPoints: "0.00",
|
||||
}
|
||||
if err := s.DB.Transaction(func(tx *gorm.DB) error {
|
||||
var active int64
|
||||
if err := tx.Model(&model.GenerationTask{}).Where("episode_id=? AND status IN ?", episodeID, activeTaskStatuses).Count(&active).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if active > 0 {
|
||||
return errors.New("当前剧集存在进行中的生成任务,暂时不能重新解析")
|
||||
}
|
||||
if err := tx.Model(&model.DramaParseTask{}).Where("episode_id=? AND status IN ?", episodeID, []string{"queued", "running", "retry_wait", "cancel_requested"}).Count(&active).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if active > 0 {
|
||||
return errors.New("当前剧集已有解析任务")
|
||||
}
|
||||
if configured.Pricing.Mode == billing.TextBillingPerRequest {
|
||||
amount, err := billing.ChargeTextRequest(tx, userID, task.ID.String(), configured.Pricing, "剧本解析")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
task.EstimatedPoints = amount
|
||||
task.PrepaidPoints = amount
|
||||
}
|
||||
if err := tx.Omit("TokenCountSource").Create(task).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Model(&model.ProjectEpisode{}).Where("id=?", episodeID).Updates(map[string]any{"status": "analyzing", "analysis_message": "剧本解析排队中"}).Error
|
||||
}); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := queuepkg.EnqueueID(s.Queue, queuepkg.TypeParseDramaEpisode, task.ID, 0); err != nil {
|
||||
s.DB.Model(task).Updates(map[string]any{"status": "failed", "error_code": "queue_unavailable", "error_message": err.Error(), "finished_at": time.Now()})
|
||||
s.DB.Model(&model.ProjectEpisode{}).Where("id=?", episodeID).Updates(map[string]any{"status": "failed", "analysis_message": "解析任务入队失败"})
|
||||
return nil, errors.New("解析任务入队失败")
|
||||
}
|
||||
return task, nil
|
||||
}
|
||||
|
||||
func (s *Creative) CancelDramaParse(userID, projectID, taskID uuid.UUID) error {
|
||||
result := s.DB.Model(&model.DramaParseTask{}).Where("id=? AND user_id=? AND project_id=? AND status IN ?", taskID, userID, projectID, []string{"queued", "running", "retry_wait"}).Updates(map[string]any{"status": "cancel_requested", "cancel_requested_at": time.Now(), "error_code": "user_cancelled", "error_message": "用户主动取消,费用不退"})
|
||||
if result.Error != nil {
|
||||
return result.Error
|
||||
}
|
||||
if result.RowsAffected == 0 {
|
||||
return gorm.ErrRecordNotFound
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *Creative) ListDramaParseTasks(userID, projectID uuid.UUID, episodeID *uuid.UUID) ([]model.DramaParseTask, error) {
|
||||
items := make([]model.DramaParseTask, 0)
|
||||
query := s.DB.Where("user_id=? AND project_id=?", userID, projectID)
|
||||
if episodeID != nil {
|
||||
query = query.Where("episode_id=?", *episodeID)
|
||||
}
|
||||
err := query.Order("created_at DESC").Limit(100).Find(&items).Error
|
||||
return items, err
|
||||
}
|
||||
|
||||
func (s *Creative) dramaTextModel(projectID uuid.UUID) (dramaTextModel, error) {
|
||||
var item dramaTextModel
|
||||
err := s.DB.Table("project_model_configs config").Select("model.id AS model_id,model.channel_id,model.name AS model_name,channel.name AS channel_name").Joins("JOIN models model ON model.id=config.model_id AND model.model_type='text' AND model.enabled=true AND model.deleted_at IS NULL").Joins("JOIN channels channel ON channel.id=model.channel_id AND channel.enabled=true AND channel.deleted_at IS NULL").Where("config.project_id=? AND config.model_type='text'", projectID).Take(&item).Error
|
||||
if err != nil {
|
||||
return item, errors.New("请先配置短剧创作文本模型")
|
||||
}
|
||||
item.Pricing, err = billing.LoadTextPricing(s.DB, item.ModelID)
|
||||
if err != nil {
|
||||
return item, errors.New("短剧创作文本模型计费配置无效")
|
||||
}
|
||||
return item, nil
|
||||
}
|
||||
|
||||
func dramaParsePrompt() string {
|
||||
return "【剧本解析规则】\n" + dramaStoryboardRules + "\n\n【角色、场景、道具解析规则】\n" + dramaEntityRules + "\n\n" + dramaOutputContract
|
||||
}
|
||||
|
||||
const dramaStoryboardRules = `你是一名资深影视剧本分析师和分镜设计师。请结合已有角色、场景、道具资料和当前剧集原文完成短剧分镜分析。
|
||||
|
||||
【层级定义】
|
||||
1. 一个 storyboard 是工作台中的一个“分镜”,对应一段完整的 5 至 15 秒视频。
|
||||
2. 一个 storyboard 内部可以包含多个摄影镜头。摄影镜头是该段视频中的机位、景别或画面切换,不得把每个摄影镜头分别输出为独立 storyboard。
|
||||
3. duration_seconds 表示整个 storyboard 的总时长,不是内部某个镜头的时长,必须是 5 至 15 之间的整数。
|
||||
4. script_content 必须按时间顺序列出该分镜内的全部摄影镜头。各镜头时间必须连续,不得重叠或断档,时长之和必须等于 duration_seconds。
|
||||
|
||||
【基本要求】
|
||||
1. 只处理本次提供的原文或分段,不得自行补写、删改剧情,绝对禁止修改角色台词。
|
||||
2. 按叙事顺序完整覆盖原文,不得遗漏关键动作、对白和剧情转折。
|
||||
3. 根据完整动作阶段和叙事意义划分 storyboard,不得为了凑时长合并无关剧情,也不得把一个连续动作机械拆成多个 storyboard。
|
||||
4. 分镜中的角色、场景和道具必须使用实体列表中的正式名称,不得使用别名、泛称或未定义名称。
|
||||
5. script_content 是工作台直接展示的完整分镜正文,必须使用下方“分镜正文固定格式”,不得只写剧情概述。
|
||||
6. prompt_content 是与整个 storyboard 对应、可直接用于视频生成的完整视觉提示词,也必须使用下方“分镜正文固定格式”,不得省略第一帧、站位或镜头细节。
|
||||
7. dialogue 提取本分镜原文中实际存在的对白;没有对白时返回空数组。
|
||||
8. asset_names 只列出本分镜实际出现的实体正式名称,并按 characters、scenes、props 分类。
|
||||
9. source_excerpt 必须引用本分镜所依据的当前原文片段。
|
||||
|
||||
【分镜正文固定格式】
|
||||
script_content 和 prompt_content 必须分别输出完整内容,并严格使用以下结构:
|
||||
场景=场景资产正式名称
|
||||
【第一帧】
|
||||
- 画面:明确景别、构图、角色外观、服装、表情、环境和光线色调。
|
||||
- 站位:明确每个角色及关键道具在画面中的前后左右关系、朝向、视线和接触关系。
|
||||
【画面内容】
|
||||
【镜头 1】(0-X秒)
|
||||
- 镜头:明确景别、机位角度、构图和运镜方式。
|
||||
- 画面:写清角色站位、连续动作、表情与视线变化、环境变化、光影色调、关键道具交互,以及镜头结束时的画面。
|
||||
- 声音/台词:写入该时间段实际发生的环境声、动作声和原文对白;没有则写“无”。
|
||||
【镜头 2】(X-duration_seconds秒)
|
||||
- 按相同项目继续描述,直至覆盖整个分镜。
|
||||
N 使用该分镜在返回数组中的顺序,从 1 开始;所有镜头时间连续、不重叠、不留空,最后一个镜头结束时间必须等于 duration_seconds。台词必须写入实际发生的镜头段落,禁止在末尾集中罗列。
|
||||
禁止只写“保持统一风格”“自然运镜”“角色互动”“气氛紧张”等空泛描述;必须提供具体、可拍摄、可执行的视觉和声音细节。
|
||||
|
||||
【内部镜头规则】
|
||||
1. 每个内部镜头按照“分镜正文固定格式”独立分段,并填写镜头、画面、声音/台词三项。
|
||||
2. 每个镜头必须描述开始时的主体状态、镜头内连续动作、人物视线或表情变化,以及结束时画面停留的位置。
|
||||
3. 关键动作、关键道具、信息揭示和人物反应应通过不同镜头清楚呈现;明显改变景别、机位、视角或主体时,应划分新的内部镜头。
|
||||
4. 5 至 7 秒通常包含 2 至 3 个镜头,8 至 12 秒通常包含 3 至 5 个镜头,13 至 15 秒通常包含 4 至 6 个镜头。以清楚呈现原文为准,不得无意义切镜。
|
||||
5. 原文明确适合长镜头时可以只有一个镜头,但必须写清持续动作、运镜过程和画面变化,不能只写静态概述。
|
||||
6. 只能描述观众能够看到或听到的内容。不得用“人物震惊”“气氛紧张”“陷入回忆”等抽象概括代替可见的表情、动作、声音和画面变化。
|
||||
7. 保持人物位置、服装、视线方向、道具状态、时间和环境连续;任何变化都必须有原文依据或在镜头内交代。
|
||||
8. 内部镜头之间按需标明硬切、视线匹配、动作匹配、声音先入等衔接关系。
|
||||
|
||||
【输出前自检】
|
||||
1. 每个 storyboard 是否为 5 至 15 秒。
|
||||
2. 内部镜头是否按时间连续排列,最后一个镜头结束时间是否等于 duration_seconds。
|
||||
3. 是否把内部摄影镜头错误拆成多个 storyboard,或把无关剧情错误合并为一个 storyboard。
|
||||
4. 每个内部镜头是否具有时间段、景别、机位、运镜和可见动作。
|
||||
5. 关键动作、线索特写、人物反应和原文对白是否完整覆盖。
|
||||
6. 是否存在抽象心理、笼统气氛或无法拍摄的概括;如有,改成可见或可听内容。
|
||||
7. 角色、场景、道具名称和连续性是否正确。
|
||||
完成自检后,按平台固定协议返回结果,不输出自检过程。`
|
||||
|
||||
const dramaEntityRules = `必须按以下规则识别和更新角色、场景、道具。
|
||||
|
||||
【通用规则】
|
||||
1. 先与用户提供的已有实体比对。相同实体必须复用已有正式名称,不得重复创建。
|
||||
2. 原文中的昵称、职位称呼、简称、化名和不同写法应合并到同一实体的 aliases;不得把别名当成新实体。
|
||||
3. 仅提取原文明示或可直接、客观推断的实体。旁白、叙述者、无明确身份的群演不作为角色资产。
|
||||
4. 已有实体原则上保留;当前原文明确提供新增或变化信息时才更新。不得用推测覆盖用户已经确认的资料。
|
||||
5. 同一角色只有在原文明示明显不同的时代、年龄阶段或造型,且需要独立视觉资产时,才拆成带限定词的独立正式名称。
|
||||
6. canonical_name 使用简洁、稳定的中文正式名称;aliases 去重且不得包含 canonical_name。
|
||||
7. description 写适合资产列表展示的一至两句客观视觉摘要。详细信息必须写入 attributes,避免只放在 description 中。
|
||||
|
||||
【角色 characters】
|
||||
attributes 尽量使用以下键:age、gender、identity、appearance、hairstyle、costume。外貌需描述可见且相对稳定的体型、脸型、肤色和五官特征;发型需描述长度、颜色和样式;服装需描述上装、下装、鞋履的颜色、款式和材质。禁止把情绪、瞬时表情、动作或性格当作固定外观。
|
||||
|
||||
【场景 scenes】
|
||||
attributes 尽量使用以下键:space_type、location、structure、visual_features、key_props。描述室内或室外、地理位置、空间结构、建筑或装饰风格、主要材质、色调、光线及固定陈设。不要把同一地点的不同称呼拆成多个场景;只有空间本身明显不同才拆分。
|
||||
|
||||
【道具 props】
|
||||
attributes 尽量使用以下键:type、appearance、function。描述道具类别、材质、颜色、尺寸、外形,以及原文明确体现的功能或剧情作用。普通背景陈设若不影响剧情且不需要单独生成,不作为独立道具。
|
||||
|
||||
【一致性检查】
|
||||
最终输出前检查:每个分镜 asset_names 中的名称都能在 entities 对应数组或用户提供的已有实体中找到;分镜正文使用的实体称呼已统一为正式名称;新增实体没有与已有实体或本次其他新增实体重复。
|
||||
|
||||
【生图提示词 image_prompt】
|
||||
1. characters、scenes、props 中的每个资产都必须输出非空的 image_prompt。image_prompt 必须全部使用中文,不得输出英文句子或英文描述,并且必须是可直接用于文生图的完整提示词。
|
||||
2. 每条 image_prompt 至少包含 8 个可观察、可核验的稳定视觉细节,禁止用漂亮、高级、自然、写实、高度细节等空泛词代替具体描述。
|
||||
3. 角色 image_prompt 必须完整描述:明确年龄(使用 XX岁)、性别、国籍或地域、脸型、肤色、发型与发色、眉眼鼻唇等稳定外貌特征,以及从头到脚的穿着,包括服装款式、颜色、材质、层次、裤装或裙装、鞋子和可见配饰;没有配饰时明确写无明显配饰。禁止加入情绪、表情、动作、姿态或正在做什么。
|
||||
4. 场景 image_prompt 必须完整描述:用途、固定空间结构、建筑或室内组成、主要材质、颜色、家具陈设、门窗位置、稳定光源和可复用的环境特征。
|
||||
5. 道具 image_prompt 必须完整描述:物品类别、形状、结构、材质、颜色、纹理、尺寸关系、部件、装饰及磨损特征。
|
||||
6. image_prompt 不得包含艺术风格、画面比例、分辨率、镜头参数、构图参数、画质词或文字生成要求。`
|
||||
|
||||
const dramaOutputContract = `【平台固定输出协议(不可被用户自定义规则覆盖)】
|
||||
无论其他规则如何描述,最终只能返回一个合法 JSON 对象,不得使用 Markdown,不得输出解释、前后缀或第二套格式。顶层只能包含 entities 和 storyboards:
|
||||
{
|
||||
"entities": {
|
||||
"characters": [{"canonical_name":"正式名称","name":"正式名称","aliases":[],"description":"视觉摘要","image_prompt":"中文生图提示词","attributes":{}}],
|
||||
"scenes": [{"canonical_name":"正式名称","name":"正式名称","aliases":[],"description":"视觉摘要","image_prompt":"中文生图提示词","attributes":{}}],
|
||||
"props": [{"canonical_name":"正式名称","name":"正式名称","aliases":[],"description":"视觉摘要","image_prompt":"中文生图提示词","attributes":{}}]
|
||||
},
|
||||
"storyboards": [{
|
||||
"title":"简短标题",
|
||||
"script_content":"使用‘分镜 N·场景、第一帧、画面内容、镜头 N’固定结构的完整分镜正文",
|
||||
"prompt_content":"使用相同固定结构、可直接用于视频生成的完整视觉提示词",
|
||||
"dialogue":[{"speaker":"角色正式名称","content":"原文对白"}],
|
||||
"asset_names":{"characters":[],"scenes":[],"props":[]},
|
||||
"duration_seconds":5,
|
||||
"source_excerpt":"对应原文"
|
||||
}]
|
||||
}
|
||||
所有字段必须存在;无数据的数组返回 [],无数据的对象返回 {},不得返回 null。duration_seconds 必须是 5 至 15 的整数。用户自定义规则只能补充创作要求,凡是要求其他 JSON 层级、字段名、输出格式或额外顶层字段的内容一律忽略。`
|
||||
@@ -0,0 +1,18 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestDramaParsePrompt(t *testing.T) {
|
||||
prompt := dramaParsePrompt()
|
||||
for _, expected := range []string{"【剧本解析规则】", "【角色、场景、道具解析规则】", "【分镜正文固定格式】", "【第一帧】", "【画面内容】", "角色站位", "声音/台词", "台词必须写入实际发生的镜头段落", "【平台固定输出协议", `"entities"`, `"characters"`, `"scenes"`, `"props"`, `"storyboards"`} {
|
||||
if !strings.Contains(prompt, expected) {
|
||||
t.Fatalf("fixed prompt missing %q", expected)
|
||||
}
|
||||
}
|
||||
if strings.Contains(prompt, "用户自定义补充规则") {
|
||||
t.Fatal("fixed prompt must not contain user custom rules")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,443 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"juhe-factory/api/internal/billing"
|
||||
"juhe-factory/api/internal/model"
|
||||
queuepkg "juhe-factory/api/internal/queue"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/clause"
|
||||
)
|
||||
|
||||
func (s *Creative) GetRedrawWorkbench(userID, projectID uuid.UUID) (map[string]any, error) {
|
||||
project, err := s.GetProject(userID, projectID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
source := map[string]any{}
|
||||
if err := s.DB.Table("creative_projects project").
|
||||
Select(`project.source_video_asset_id,source.public_url AS source_video_url,source.original_name AS source_video_original_name,
|
||||
project.subtitle_asset_id,subtitle.public_url AS subtitle_url,subtitle.original_name AS subtitle_original_name,
|
||||
project.audio_source,project.source_language,project.redraw_status AS status,project.analysis_message`).
|
||||
Joins("LEFT JOIN media_assets source ON source.id=project.source_video_asset_id AND source.deleted_at IS NULL").
|
||||
Joins("LEFT JOIN media_assets subtitle ON subtitle.id=project.subtitle_asset_id AND subtitle.deleted_at IS NULL").
|
||||
Where("project.id=? AND project.user_id=? AND project.project_type='video_redraw' AND project.deleted_at IS NULL", projectID, userID).
|
||||
Take(&source).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
assets := make([]map[string]any, 0)
|
||||
if err := s.DB.Table("project_assets asset").
|
||||
Select("asset.id,asset.project_id,asset.asset_type,asset.name,asset.user_edited").
|
||||
Where("asset.project_id=? AND asset.deleted_at IS NULL", projectID).
|
||||
Order("asset.asset_type,asset.created_at").Find(&assets).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
storyboards := make([]map[string]any, 0)
|
||||
if err := s.DB.Table("episode_storyboards storyboard").
|
||||
Select(`storyboard.*,thumbnail.public_url AS thumbnail_url,0 AS history_count`).
|
||||
Joins("LEFT JOIN media_assets thumbnail ON thumbnail.id=storyboard.thumbnail_asset_id").
|
||||
Where("storyboard.project_id=? AND storyboard.deleted_at IS NULL", projectID).
|
||||
Order("storyboard.sequence_no").Find(&storyboards).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var scriptRow struct {
|
||||
RedrawScript string
|
||||
}
|
||||
if err := s.DB.Table("creative_projects").Select("coalesce(redraw_script,'') AS redraw_script").
|
||||
Where("id=? AND user_id=? AND project_type='video_redraw' AND deleted_at IS NULL", projectID, userID).
|
||||
Take(&scriptRow).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return map[string]any{"project": project, "source": source, "assets": assets, "storyboards": storyboards, "script_content": scriptRow.RedrawScript}, nil
|
||||
}
|
||||
|
||||
func (s *Creative) GetRedrawEpisodeWorkbench(userID, projectID, episodeID uuid.UUID) (map[string]any, error) {
|
||||
data, err := s.GetWorkbench(userID, projectID, episodeID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
episode, ok := data["episode"].(map[string]any)
|
||||
if !ok {
|
||||
return nil, errors.New("剧集数据无效")
|
||||
}
|
||||
var script string
|
||||
if err := s.DB.Table("project_episodes episode").
|
||||
Select("coalesce(episode.redraw_script,'')").
|
||||
Joins("JOIN creative_projects project ON project.id=episode.project_id AND project.user_id=? AND project.project_type='video_redraw' AND project.deleted_at IS NULL", userID).
|
||||
Where("episode.id=? AND episode.project_id=? AND episode.deleted_at IS NULL", episodeID, projectID).
|
||||
Scan(&script).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return map[string]any{
|
||||
"project": data["project"], "episode": episode, "source": episode,
|
||||
"assets": data["assets"], "storyboards": data["storyboards"], "script_content": script,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// SaveRedrawEpisodeScript 校验项目归属和类型后保存用户编辑的剧集反推剧本。
|
||||
func (s *Creative) SaveRedrawEpisodeScript(userID, projectID, episodeID uuid.UUID, content string) error {
|
||||
return s.DB.Transaction(func(tx *gorm.DB) error {
|
||||
var episode model.ProjectEpisode
|
||||
if err := tx.Table("project_episodes episode").Select("episode.id").
|
||||
Joins("JOIN creative_projects project ON project.id=episode.project_id AND project.user_id=? AND project.project_type='video_redraw' AND project.deleted_at IS NULL", userID).
|
||||
Where("episode.id=? AND episode.project_id=? AND episode.deleted_at IS NULL", episodeID, projectID).
|
||||
Take(&episode).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Model(&model.ProjectEpisode{}).Where("id=?", episode.ID).Update("redraw_script", content).Error
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Creative) QueueRedrawEpisodeScript(userID, projectID, episodeID uuid.UUID) (*model.GenerationTask, error) {
|
||||
if s.Queue == nil {
|
||||
return nil, errors.New("反推任务队列不可用")
|
||||
}
|
||||
config, err := s.analysisModel(projectID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
task := &model.GenerationTask{
|
||||
RequestID: "script_reverse_" + uuid.NewString(), UserID: userID, ChannelID: &config.ChannelID,
|
||||
ModelID: &config.ModelID, ProjectID: &projectID, EpisodeID: &episodeID, TaskType: "prompt_reverse", Status: "submitted",
|
||||
InputData: json.RawMessage(`{"mode":"script"}`), EstimatedPoints: "0.00", PrepaidPoints: "0.00",
|
||||
}
|
||||
if err := s.DB.Transaction(func(tx *gorm.DB) error {
|
||||
scriptPrompt, err := selectedPromptContent(tx, userID, "剧本反推")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
task.InputData = mustJSON(map[string]string{"mode": "script", "prompt": scriptPrompt})
|
||||
var episode model.ProjectEpisode
|
||||
if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Table("project_episodes episode").Select("episode.*").
|
||||
Joins("JOIN creative_projects project ON project.id=episode.project_id AND project.user_id=? AND project.project_type='video_redraw' AND project.deleted_at IS NULL", userID).
|
||||
Where("episode.id=? AND episode.project_id=? AND episode.deleted_at IS NULL", episodeID, projectID).Take(&episode).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if episode.Status != "review" {
|
||||
return errors.New("请等待全部分镜提示词反推完成")
|
||||
}
|
||||
var total, incomplete int64
|
||||
if err := tx.Model(&model.EpisodeStoryboard{}).Where("episode_id=? AND deleted_at IS NULL", episodeID).Count(&total).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if total == 0 {
|
||||
return errors.New("暂无可用于反推剧本的分镜提示词")
|
||||
}
|
||||
if err := tx.Model(&model.EpisodeStoryboard{}).
|
||||
Where("episode_id=? AND deleted_at IS NULL AND trim(coalesce(prompt_content,''))=''", episodeID).
|
||||
Count(&incomplete).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if incomplete > 0 {
|
||||
return errors.New("请等待所有分镜提示词反推完成")
|
||||
}
|
||||
var prompts []struct {
|
||||
SequenceNo int `json:"sequence_no"`
|
||||
StartMS int64 `json:"start_ms"`
|
||||
EndMS int64 `json:"end_ms"`
|
||||
PromptContent string `json:"prompt_content"`
|
||||
}
|
||||
if err := tx.Model(&model.EpisodeStoryboard{}).Select("sequence_no,start_ms,end_ms,prompt_content").
|
||||
Where("episode_id=? AND deleted_at IS NULL", episodeID).Order("sequence_no").Find(&prompts).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
task.PromptSnapshot = string(mustJSON(prompts))
|
||||
var active int64
|
||||
if err := tx.Model(&model.GenerationTask{}).Where("episode_id=? AND task_type='prompt_reverse' AND status IN ?", episodeID, activeTaskStatuses).Count(&active).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if active > 0 {
|
||||
return errors.New("当前剧集已有反推任务正在执行")
|
||||
}
|
||||
if err := billing.CreateTextGenerationTask(tx, task, config.Pricing, "剧本反推"); err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Model(&episode).Updates(map[string]any{"status": "generating", "analysis_message": "正在反推完整剧本"}).Error
|
||||
}); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := queuepkg.EnqueueID(s.Queue, queuepkg.TypeAnalyzeEpisode, task.ID, 0); err != nil {
|
||||
if settleErr := s.failQueuedTask(task.ID, "queue_unavailable", "剧本反推任务入队失败"); settleErr != nil {
|
||||
return nil, fmt.Errorf("剧本反推任务入队失败且预扣返还失败: %w", settleErr)
|
||||
}
|
||||
_ = s.DB.Model(&model.ProjectEpisode{}).Where("id=?", episodeID).Updates(map[string]any{"status": "review", "analysis_message": "剧本反推任务入队失败"}).Error
|
||||
return nil, errors.New("反推任务队列暂时不可用,请稍后重试")
|
||||
}
|
||||
return task, nil
|
||||
}
|
||||
|
||||
func mustJSON(value any) []byte {
|
||||
result, _ := json.Marshal(value)
|
||||
return result
|
||||
}
|
||||
|
||||
// selectedPromptContent 按用户偏好选取指定类型的提示词内容,未命中用户偏好时兜底取最新系统提示词。
|
||||
// 用户未设置偏好或系统未配置均属于正常业务情况,使用 Limit(1).Find 避免触发 GORM 的 record not found 日志。
|
||||
func selectedPromptContent(tx *gorm.DB, userID uuid.UUID, promptType string) (string, error) {
|
||||
var selected struct{ Content string }
|
||||
if err := tx.Table("prompts p").Select("p.content").
|
||||
Joins("JOIN user_prompt_preferences pref ON pref.prompt_id=p.id AND pref.user_id=? AND pref.prompt_type=?", userID, promptType).
|
||||
Where("p.type=? AND p.deleted_at IS NULL AND (p.scope='system' OR (p.scope='user' AND p.owner_user_id=?))", promptType, userID).
|
||||
Limit(1).Find(&selected).Error; err != nil {
|
||||
return "", err
|
||||
}
|
||||
if strings.TrimSpace(selected.Content) == "" {
|
||||
if err := tx.Table("prompts").Select("content").
|
||||
Where("type=? AND scope='system' AND deleted_at IS NULL", promptType).
|
||||
Order("updated_at DESC").Limit(1).Find(&selected).Error; err != nil {
|
||||
return "", err
|
||||
}
|
||||
if strings.TrimSpace(selected.Content) == "" {
|
||||
return "", fmt.Errorf("请先在提示词管理中配置%s提示词", promptType)
|
||||
}
|
||||
}
|
||||
return strings.TrimSpace(selected.Content), nil
|
||||
}
|
||||
|
||||
func (s *Creative) UpdateRedrawAnalysisSettings(userID, projectID uuid.UUID, audioSource, sourceLanguage string) error {
|
||||
if audioSource != "video_audio" && audioSource != "subtitle_file" {
|
||||
return errors.New("台词来源无效")
|
||||
}
|
||||
sourceLanguage = strings.TrimSpace(sourceLanguage)
|
||||
var language any
|
||||
if sourceLanguage != "" {
|
||||
language = sourceLanguage
|
||||
}
|
||||
result := s.DB.Table("creative_projects").
|
||||
Where("id=? AND user_id=? AND project_type='video_redraw' AND deleted_at IS NULL", projectID, userID).
|
||||
Updates(map[string]any{"audio_source": audioSource, "source_language": language})
|
||||
if result.Error != nil {
|
||||
return result.Error
|
||||
}
|
||||
if result.RowsAffected == 0 {
|
||||
return gorm.ErrRecordNotFound
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *Creative) QueueRedrawAnalysis(userID, projectID uuid.UUID) (*model.GenerationTask, error) {
|
||||
return s.queueRedrawAnalysis(userID, projectID, nil)
|
||||
}
|
||||
|
||||
func (s *Creative) QueueRedrawStoryboardAnalysis(userID, projectID, storyboardID uuid.UUID) (*model.GenerationTask, error) {
|
||||
return s.queueRedrawAnalysis(userID, projectID, &storyboardID)
|
||||
}
|
||||
|
||||
func (s *Creative) QueueRedrawScript(userID, projectID uuid.UUID) (*model.GenerationTask, error) {
|
||||
if s.Queue == nil {
|
||||
return nil, errors.New("反推任务队列不可用")
|
||||
}
|
||||
config, err := s.analysisModel(projectID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
task := &model.GenerationTask{
|
||||
RequestID: "script_reverse_" + uuid.NewString(), UserID: userID, ChannelID: &config.ChannelID,
|
||||
ModelID: &config.ModelID, ProjectID: &projectID, TaskType: "prompt_reverse", Status: "submitted",
|
||||
InputData: json.RawMessage(`{"mode":"script"}`), EstimatedPoints: "0.00", PrepaidPoints: "0.00",
|
||||
}
|
||||
if err := s.DB.Transaction(func(tx *gorm.DB) error {
|
||||
scriptPrompt, err := selectedPromptContent(tx, userID, "剧本反推")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
task.InputData = mustJSON(map[string]string{"mode": "script", "prompt": scriptPrompt})
|
||||
var project model.CreativeProject
|
||||
if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).
|
||||
Where("id=? AND user_id=? AND project_type='video_redraw' AND deleted_at IS NULL", projectID, userID).
|
||||
Take(&project).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if project.RedrawStatus != "review" {
|
||||
return errors.New("请等待全部分镜提示词反推完成")
|
||||
}
|
||||
var total, incomplete int64
|
||||
if err := tx.Model(&model.EpisodeStoryboard{}).Where("project_id=? AND deleted_at IS NULL", projectID).Count(&total).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if total == 0 {
|
||||
return errors.New("暂无可用于反推剧本的分镜提示词")
|
||||
}
|
||||
if err := tx.Model(&model.EpisodeStoryboard{}).
|
||||
Where("project_id=? AND deleted_at IS NULL AND trim(coalesce(prompt_content,''))=''", projectID).
|
||||
Count(&incomplete).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if incomplete > 0 {
|
||||
return errors.New("请等待所有分镜提示词反推完成")
|
||||
}
|
||||
var storyboardPrompts []struct {
|
||||
SequenceNo int `json:"sequence_no"`
|
||||
StartMS int64 `json:"start_ms"`
|
||||
EndMS int64 `json:"end_ms"`
|
||||
PromptContent string `json:"prompt_content"`
|
||||
}
|
||||
if err := tx.Model(&model.EpisodeStoryboard{}).Select("sequence_no,start_ms,end_ms,prompt_content").
|
||||
Where("project_id=? AND deleted_at IS NULL", projectID).Order("sequence_no").Find(&storyboardPrompts).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
promptSnapshot, err := json.Marshal(storyboardPrompts)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
task.PromptSnapshot = string(promptSnapshot)
|
||||
var active int64
|
||||
if err := tx.Model(&model.GenerationTask{}).Where("project_id=? AND task_type='prompt_reverse' AND status IN ?", projectID, activeTaskStatuses).Count(&active).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if active > 0 {
|
||||
return errors.New("当前剧本已有反推任务正在执行")
|
||||
}
|
||||
if err := billing.CreateTextGenerationTask(tx, task, config.Pricing, "剧本反推"); err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Model(&project).Updates(map[string]any{"redraw_status": "generating", "analysis_message": "正在反推完整剧本"}).Error
|
||||
}); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := queuepkg.EnqueueID(s.Queue, queuepkg.TypeAnalyzeEpisode, task.ID, 0); err != nil {
|
||||
if settleErr := s.failQueuedTask(task.ID, "queue_unavailable", "剧本反推任务入队失败"); settleErr != nil {
|
||||
return nil, fmt.Errorf("剧本反推任务入队失败且预扣返还失败: %w", settleErr)
|
||||
}
|
||||
_ = s.DB.Model(&model.CreativeProject{}).Where("id=?", projectID).Updates(map[string]any{"redraw_status": "review", "analysis_message": "剧本反推任务入队失败"}).Error
|
||||
return nil, errors.New("反推任务队列暂时不可用,请稍后重试")
|
||||
}
|
||||
return task, nil
|
||||
}
|
||||
|
||||
func (s *Creative) queueRedrawAnalysis(userID, projectID uuid.UUID, storyboardID *uuid.UUID) (*model.GenerationTask, error) {
|
||||
if s.Queue == nil {
|
||||
return nil, errors.New("反推任务队列不可用")
|
||||
}
|
||||
config, err := s.analysisModel(projectID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
requestPrefix := "analysis_"
|
||||
input := json.RawMessage(`{}`)
|
||||
if storyboardID != nil {
|
||||
requestPrefix = "storyboard_analysis_"
|
||||
input = json.RawMessage(`{"mode":"storyboard"}`)
|
||||
}
|
||||
task := &model.GenerationTask{
|
||||
RequestID: requestPrefix + uuid.NewString(), UserID: userID, ChannelID: &config.ChannelID,
|
||||
ModelID: &config.ModelID, ProjectID: &projectID, StoryboardID: storyboardID,
|
||||
TaskType: "prompt_reverse", Status: "submitted", InputData: input, EstimatedPoints: "0.00", PrepaidPoints: "0.00",
|
||||
}
|
||||
if err := s.DB.Transaction(func(tx *gorm.DB) error {
|
||||
var project model.CreativeProject
|
||||
if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).
|
||||
Where("id=? AND user_id=? AND project_type='video_redraw' AND deleted_at IS NULL", projectID, userID).
|
||||
Take(&project).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if project.SourceVideoAssetID == nil {
|
||||
return errors.New("请先上传原视频")
|
||||
}
|
||||
if project.AudioSource == "subtitle_file" && project.SubtitleAssetID == nil {
|
||||
return errors.New("当前选择字幕文件识别,请先上传字幕")
|
||||
}
|
||||
if storyboardID != nil {
|
||||
var storyboard model.EpisodeStoryboard
|
||||
if err := tx.Where("id=? AND project_id=? AND deleted_at IS NULL", *storyboardID, projectID).Take(&storyboard).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if storyboard.Locked {
|
||||
return errors.New("当前分镜已保护,无法重新反推")
|
||||
}
|
||||
}
|
||||
var active int64
|
||||
if err := tx.Model(&model.GenerationTask{}).
|
||||
Where("project_id=? AND task_type='prompt_reverse' AND status IN ?", projectID, activeTaskStatuses).
|
||||
Count(&active).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if active > 0 {
|
||||
return errors.New("当前剧本已有反推任务正在执行")
|
||||
}
|
||||
if err := billing.CreateTextGenerationTask(tx, task, config.Pricing, "视频反推"); err != nil {
|
||||
return err
|
||||
}
|
||||
if storyboardID == nil {
|
||||
return tx.Model(&model.CreativeProject{}).Where("id=?", projectID).
|
||||
Updates(map[string]any{"redraw_status": "analyzing", "analysis_message": "等待视频分析"}).Error
|
||||
}
|
||||
return nil
|
||||
}); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := queuepkg.EnqueueID(s.Queue, queuepkg.TypeAnalyzeEpisode, task.ID, 0); err != nil {
|
||||
if settleErr := s.failQueuedTask(task.ID, "queue_unavailable", "反推任务入队失败"); settleErr != nil {
|
||||
return nil, fmt.Errorf("反推任务入队失败且预扣返还失败: %w", settleErr)
|
||||
}
|
||||
return nil, errors.New("反推任务队列暂时不可用,请稍后重试")
|
||||
}
|
||||
return task, nil
|
||||
}
|
||||
|
||||
func (s *Creative) ResetRedrawAnalysis(userID, projectID uuid.UUID) ([]string, error) {
|
||||
media := newDeletionMediaSet()
|
||||
err := s.DB.Transaction(func(tx *gorm.DB) error {
|
||||
var project model.CreativeProject
|
||||
if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).
|
||||
Where("id=? AND user_id=? AND project_type='video_redraw' AND deleted_at IS NULL", projectID, userID).
|
||||
Take(&project).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
var active int64
|
||||
if err := tx.Model(&model.GenerationTask{}).Where("project_id=? AND status IN ?", projectID, activeTaskStatuses).Count(&active).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if active > 0 {
|
||||
return errors.New("当前剧本存在进行中的任务,暂时不能重新反推")
|
||||
}
|
||||
if err := removeRedrawDerivedContent(tx, projectID, media); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := media.deleteRows(tx); err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Model(&project).Updates(map[string]any{"redraw_status": "uploaded", "analysis_message": nil, "redraw_script": nil}).Error
|
||||
})
|
||||
return media.objectKeys(), err
|
||||
}
|
||||
|
||||
func removeRedrawDerivedContent(tx *gorm.DB, projectID uuid.UUID, media *deletionMediaSet) error {
|
||||
taskScope := `project_id=? AND (task_type IN ('prompt_reverse','video_generation')
|
||||
OR (task_type='image_generation' AND input_data->>'target_type'='storyboard'))`
|
||||
if err := media.addQuery(tx.Table("media_assets media").Select("DISTINCT media.id,media.object_key").
|
||||
Joins("JOIN generation_outputs output ON output.media_asset_id=media.id").
|
||||
Joins("JOIN generation_tasks task ON task.id=output.task_id").Where(taskScope, projectID)); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := media.addQuery(tx.Table("media_assets media").Select("DISTINCT media.id,media.object_key").
|
||||
Joins("JOIN episode_storyboards storyboard ON storyboard.thumbnail_asset_id=media.id").
|
||||
Where("storyboard.project_id=?", projectID)); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := collectMediaByObjectKey(tx, media, fmt.Sprintf("juyou_ran/video-redraw/projects/%s/storyboards/%%", projectID)); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := collectMediaByObjectKey(tx, media, fmt.Sprintf("juyou_ran/video-redraw/projects/%s/episodes/%%/storyboards/%%", projectID)); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := tx.Model(&model.CreativeProject{}).Where("id=?", projectID).Update("cover_asset_id", nil).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if err := tx.Exec("UPDATE episode_storyboards SET active_output_id=NULL WHERE project_id=?", projectID).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if err := tx.Exec("DELETE FROM generation_outputs WHERE task_id IN (SELECT id FROM generation_tasks WHERE "+taskScope+")", projectID).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if err := tx.Exec("DELETE FROM generation_tasks WHERE "+taskScope, projectID).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Exec("DELETE FROM episode_storyboards WHERE project_id=?", projectID).Error
|
||||
}
|
||||
@@ -0,0 +1,298 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"juhe-factory/api/internal/billing"
|
||||
"juhe-factory/api/internal/model"
|
||||
queuepkg "juhe-factory/api/internal/queue"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/clause"
|
||||
)
|
||||
|
||||
type ScriptAnalysisInput struct {
|
||||
Name string `json:"name"`
|
||||
SourceContent string `json:"source_content"`
|
||||
ResultContent string `json:"result_content"`
|
||||
}
|
||||
|
||||
type ScriptAnalysisImportProject struct {
|
||||
ID uuid.UUID `json:"id"`
|
||||
Name string `json:"name"`
|
||||
EpisodeCount int `json:"episode_count"`
|
||||
CharCount int `json:"char_count"`
|
||||
UpdatedAt time.Time `json:"updated_at"`
|
||||
}
|
||||
|
||||
func (s *Creative) ListScriptAnalyses(userID uuid.UUID) ([]model.ScriptAnalysis, error) {
|
||||
items := make([]model.ScriptAnalysis, 0)
|
||||
err := s.DB.Where("user_id=? AND deleted_at IS NULL", userID).Order("updated_at DESC").Find(&items).Error
|
||||
return items, err
|
||||
}
|
||||
|
||||
func (s *Creative) CreateScriptAnalysis(userID uuid.UUID) (*model.ScriptAnalysis, error) {
|
||||
item := &model.ScriptAnalysis{UserID: userID, Name: "未命名剧本", AnalysisResult: json.RawMessage(`{}`), AnalysisStatus: "idle"}
|
||||
if err := s.DB.Create(item).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return item, nil
|
||||
}
|
||||
|
||||
func (s *Creative) ListScriptAnalysisImportProjects(userID uuid.UUID) ([]ScriptAnalysisImportProject, error) {
|
||||
items := make([]ScriptAnalysisImportProject, 0)
|
||||
err := s.DB.Table("creative_projects project").
|
||||
Select(`project.id,project.name,project.updated_at,
|
||||
count(episode.id) FILTER (WHERE episode.deleted_at IS NULL AND episode.status='review' AND btrim(coalesce(episode.redraw_script,'')) <> '') AS episode_count,
|
||||
coalesce(sum(char_length(episode.redraw_script)) FILTER (WHERE episode.deleted_at IS NULL AND episode.status='review' AND btrim(coalesce(episode.redraw_script,'')) <> ''),0) AS char_count`).
|
||||
Joins("LEFT JOIN project_episodes episode ON episode.project_id=project.id").
|
||||
Where(`project.user_id=? AND project.project_type='video_redraw' AND project.deleted_at IS NULL AND
|
||||
EXISTS (
|
||||
SELECT 1 FROM project_episodes source_episode
|
||||
WHERE source_episode.project_id=project.id AND source_episode.deleted_at IS NULL
|
||||
AND source_episode.status='review' AND btrim(coalesce(source_episode.redraw_script,'')) <> ''
|
||||
)`, userID).
|
||||
Group("project.id").Order("project.updated_at DESC").Scan(&items).Error
|
||||
return items, err
|
||||
}
|
||||
|
||||
func (s *Creative) ImportScriptAnalysisProject(userID, analysisID, projectID uuid.UUID) (*model.ScriptAnalysis, error) {
|
||||
err := s.DB.Transaction(func(tx *gorm.DB) error {
|
||||
var project model.CreativeProject
|
||||
if err := tx.Select("id", "name").
|
||||
Where("id=? AND user_id=? AND project_type='video_redraw' AND deleted_at IS NULL", projectID, userID).
|
||||
Take(&project).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
var episodes []model.ProjectEpisode
|
||||
if err := tx.Select("episode_no", "redraw_script").
|
||||
Where("project_id=? AND deleted_at IS NULL AND status='review' AND btrim(coalesce(redraw_script,'')) <> ''", projectID).
|
||||
Order("episode_no").Find(&episodes).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if len(episodes) == 0 {
|
||||
return errors.New("该项目暂无可导入的剧本")
|
||||
}
|
||||
parts := make([]string, 0, len(episodes))
|
||||
for _, episode := range episodes {
|
||||
parts = append(parts, fmt.Sprintf("【第%d集】\n%s", episode.EpisodeNo, strings.TrimSpace(episode.RedrawScript)))
|
||||
}
|
||||
name := project.Name
|
||||
return replaceScriptAnalysisSource(tx, userID, analysisID, name, strings.Join(parts, "\n\n"))
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return s.GetScriptAnalysis(userID, analysisID)
|
||||
}
|
||||
|
||||
func (s *Creative) ImportScriptAnalysisFile(userID, analysisID uuid.UUID, filename string, data []byte) (*model.ScriptAnalysis, error) {
|
||||
extension := strings.ToLower(filepath.Ext(filename))
|
||||
if extension != ".txt" && extension != ".doc" {
|
||||
return nil, errors.New("仅支持TXT和DOC文件")
|
||||
}
|
||||
if len(data) > maxDramaImportBytes {
|
||||
return nil, errors.New("剧本文件不能超过3MB")
|
||||
}
|
||||
content, _, err := parseDramaSource(filename, data, "")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
content = normalizeDramaText(content)
|
||||
if content == "" {
|
||||
return nil, errors.New("未读取到有效剧本正文")
|
||||
}
|
||||
name := strings.TrimSpace(strings.TrimSuffix(filepath.Base(filename), filepath.Ext(filename)))
|
||||
if name == "" {
|
||||
name = "未命名剧本"
|
||||
}
|
||||
if err := s.DB.Transaction(func(tx *gorm.DB) error {
|
||||
return replaceScriptAnalysisSource(tx, userID, analysisID, name, content)
|
||||
}); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return s.GetScriptAnalysis(userID, analysisID)
|
||||
}
|
||||
|
||||
func replaceScriptAnalysisSource(tx *gorm.DB, userID, analysisID uuid.UUID, name, content string) error {
|
||||
var item model.ScriptAnalysis
|
||||
if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).
|
||||
Where("id=? AND user_id=? AND deleted_at IS NULL", analysisID, userID).Take(&item).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if item.AnalysisStatus == "queued" || item.AnalysisStatus == "running" {
|
||||
return errors.New("剧本正在分析,暂时不能导入")
|
||||
}
|
||||
if err := tx.Where("script_analysis_id=?", analysisID).Delete(&model.ScriptAnalysisCharacter{}).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Model(&item).Updates(map[string]any{
|
||||
"name": strings.TrimSpace(name),
|
||||
"source_content": content, "result_content": "",
|
||||
"analysis_result": json.RawMessage(`{}`), "analysis_status": "idle", "analysis_message": "",
|
||||
}).Error
|
||||
}
|
||||
|
||||
func (s *Creative) GetScriptAnalysis(userID, id uuid.UUID) (*model.ScriptAnalysis, error) {
|
||||
var item model.ScriptAnalysis
|
||||
if err := s.DB.Preload("Characters", func(db *gorm.DB) *gorm.DB { return db.Order("sort_order,name") }).
|
||||
Where("id=? AND user_id=? AND deleted_at IS NULL", id, userID).Take(&item).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &item, nil
|
||||
}
|
||||
|
||||
type scriptAnalysisModelConfig struct {
|
||||
ModelID uuid.UUID
|
||||
ChannelID uuid.UUID
|
||||
ModelName string
|
||||
BaseURL string
|
||||
APIKeyCiphertext string
|
||||
Pricing billing.TextPricing `gorm:"-"`
|
||||
}
|
||||
|
||||
func (s *Creative) scriptAnalysisModel(userID uuid.UUID) (scriptAnalysisModelConfig, error) {
|
||||
var config scriptAnalysisModelConfig
|
||||
result := s.DB.Raw(`SELECT model.id AS model_id,channel.id AS channel_id,model.name AS model_name,
|
||||
channel.base_url,channel.api_key_ciphertext
|
||||
FROM user_model_configs preference
|
||||
JOIN models model ON model.id=preference.model_id AND model.model_type='text' AND model.enabled=true AND model.deleted_at IS NULL
|
||||
JOIN channels channel ON channel.id=model.channel_id AND channel.enabled=true AND channel.deleted_at IS NULL
|
||||
WHERE preference.user_id=? AND preference.project_type='script_analysis' AND preference.model_type='text'`, userID).Scan(&config)
|
||||
if result.Error != nil {
|
||||
return config, result.Error
|
||||
}
|
||||
if result.RowsAffected == 0 {
|
||||
return config, errors.New("请先配置剧本分析文本模型")
|
||||
}
|
||||
pricing, err := billing.LoadTextPricing(s.DB, config.ModelID)
|
||||
if err != nil {
|
||||
return config, errors.New("剧本分析文本模型计费配置无效")
|
||||
}
|
||||
config.Pricing = pricing
|
||||
return config, nil
|
||||
}
|
||||
|
||||
func (s *Creative) QueueScriptAnalysis(userID, analysisID uuid.UUID) (*model.GenerationTask, error) {
|
||||
if s.Queue == nil {
|
||||
return nil, errors.New("剧本分析任务队列不可用")
|
||||
}
|
||||
config, err := s.scriptAnalysisModel(userID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
task := &model.GenerationTask{
|
||||
RequestID: "script_analysis_" + uuid.NewString(), UserID: userID, ChannelID: &config.ChannelID,
|
||||
ModelID: &config.ModelID, ScriptAnalysisID: &analysisID, TaskType: "script_analysis", Status: "submitted",
|
||||
EstimatedPoints: "0.00", PrepaidPoints: "0.00",
|
||||
}
|
||||
err = s.DB.Transaction(func(tx *gorm.DB) error {
|
||||
var item model.ScriptAnalysis
|
||||
if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).
|
||||
Where("id=? AND user_id=? AND deleted_at IS NULL", analysisID, userID).Take(&item).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if strings.TrimSpace(item.SourceContent) == "" {
|
||||
return errors.New("请先填写剧本原文")
|
||||
}
|
||||
var active int64
|
||||
if err := tx.Model(&model.GenerationTask{}).
|
||||
Where("script_analysis_id=? AND status IN ?", analysisID, activeTaskStatuses).Count(&active).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if active > 0 {
|
||||
return errors.New("当前剧本正在分析,请勿重复提交")
|
||||
}
|
||||
prompt, err := selectedPromptContent(tx, userID, "剧本分析")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
task.PromptSnapshot = prompt
|
||||
task.InputData = mustJSON(map[string]string{"source_content": item.SourceContent})
|
||||
if err := billing.CreateTextGenerationTask(tx, task, config.Pricing, "剧本分析"); err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Model(&item).Updates(map[string]any{"analysis_status": "queued", "analysis_message": "剧本分析排队中"}).Error
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := queuepkg.EnqueueID(s.Queue, queuepkg.TypeAnalyzeScript, task.ID, 0); err != nil {
|
||||
if settleErr := s.failQueuedTask(task.ID, "queue_unavailable", "剧本分析任务入队失败"); settleErr != nil {
|
||||
return nil, fmt.Errorf("剧本分析任务入队失败且预扣返还失败: %w", settleErr)
|
||||
}
|
||||
_ = s.DB.Model(&model.ScriptAnalysis{}).Where("id=?", analysisID).
|
||||
Updates(map[string]any{"analysis_status": "failed", "analysis_message": "剧本分析任务入队失败"}).Error
|
||||
return nil, errors.New("剧本分析任务队列暂时不可用,请稍后重试")
|
||||
}
|
||||
return task, nil
|
||||
}
|
||||
|
||||
func (s *Creative) CancelScriptAnalysis(userID, analysisID uuid.UUID) error {
|
||||
return s.DB.Transaction(func(tx *gorm.DB) error {
|
||||
var task model.GenerationTask
|
||||
if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).
|
||||
Where("script_analysis_id=? AND user_id=? AND task_type='script_analysis' AND status IN ?", analysisID, userID,
|
||||
[]string{"submitted", "processing", "cancel_requested"}).
|
||||
Order("created_at DESC").Take(&task).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if task.Status == "cancel_requested" {
|
||||
return nil
|
||||
}
|
||||
if task.Status == "submitted" {
|
||||
if err := tx.Model(&task).Updates(map[string]any{
|
||||
"status": "cancelled", "actual_points": task.PrepaidPoints, "finished_at": time.Now(),
|
||||
"error_code": "user_cancelled", "error_message": "用户取消,已扣积分不退",
|
||||
}).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Model(&model.ScriptAnalysis{}).Where("id=? AND user_id=? AND deleted_at IS NULL", analysisID, userID).Updates(map[string]any{
|
||||
"analysis_status": "idle", "analysis_message": "剧本分析已取消,已扣积分不退",
|
||||
}).Error
|
||||
}
|
||||
if err := tx.Model(&task).Updates(map[string]any{
|
||||
"status": "cancel_requested", "error_code": "user_cancelled", "error_message": "正在取消分析",
|
||||
}).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Model(&model.ScriptAnalysis{}).Where("id=? AND user_id=? AND deleted_at IS NULL", analysisID, userID).
|
||||
Update("analysis_message", "正在取消分析").Error
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Creative) UpdateScriptAnalysis(userID, id uuid.UUID, input ScriptAnalysisInput) (*model.ScriptAnalysis, error) {
|
||||
name := strings.TrimSpace(input.Name)
|
||||
if name == "" {
|
||||
return nil, errors.New("剧本名称不能为空")
|
||||
}
|
||||
result := s.DB.Model(&model.ScriptAnalysis{}).
|
||||
Where("id=? AND user_id=? AND deleted_at IS NULL", id, userID).
|
||||
Updates(map[string]any{"name": name, "source_content": input.SourceContent, "result_content": input.ResultContent})
|
||||
if result.Error != nil {
|
||||
return nil, result.Error
|
||||
}
|
||||
if result.RowsAffected == 0 {
|
||||
return nil, gorm.ErrRecordNotFound
|
||||
}
|
||||
return s.GetScriptAnalysis(userID, id)
|
||||
}
|
||||
|
||||
func (s *Creative) DeleteScriptAnalysis(userID, id uuid.UUID) error {
|
||||
result := s.DB.Model(&model.ScriptAnalysis{}).
|
||||
Where("id=? AND user_id=? AND deleted_at IS NULL", id, userID).
|
||||
Update("deleted_at", gorm.Expr("CURRENT_TIMESTAMP"))
|
||||
if result.Error != nil {
|
||||
return result.Error
|
||||
}
|
||||
if result.RowsAffected == 0 {
|
||||
return gorm.ErrRecordNotFound
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,24 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
func TestImportScriptAnalysisFileRejectsUnsupportedFormat(t *testing.T) {
|
||||
service := &Creative{}
|
||||
_, err := service.ImportScriptAnalysisFile(uuid.New(), uuid.New(), "script.docx", []byte("content"))
|
||||
if err == nil || !strings.Contains(err.Error(), "仅支持TXT和DOC文件") {
|
||||
t.Fatalf("expected unsupported format error, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestImportScriptAnalysisFileRejectsOversizedFile(t *testing.T) {
|
||||
service := &Creative{}
|
||||
_, err := service.ImportScriptAnalysisFile(uuid.New(), uuid.New(), "script.txt", make([]byte, maxDramaImportBytes+1))
|
||||
if err == nil || !strings.Contains(err.Error(), "不能超过3MB") {
|
||||
t.Fatalf("expected file size error, got %v", err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,272 @@
|
||||
// WEB 用户服务,负责用户认证、账户资料、积分查询和兑换码业务。
|
||||
package service
|
||||
|
||||
import (
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"errors"
|
||||
"strings"
|
||||
"time"
|
||||
"unicode/utf8"
|
||||
|
||||
"juhe-factory/api/internal/billing"
|
||||
"juhe-factory/api/internal/model"
|
||||
"juhe-factory/api/internal/security"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/clause"
|
||||
)
|
||||
|
||||
type Web struct {
|
||||
DB *gorm.DB
|
||||
Passwords security.PasswordHasher
|
||||
Tokens security.WebTokenService
|
||||
RefreshTTL time.Duration
|
||||
}
|
||||
|
||||
var (
|
||||
ErrInvalidUsername = errors.New("用户名长度必须为 5~12 个字符")
|
||||
ErrUsernameExists = errors.New("用户名已被使用")
|
||||
ErrOriginalPassword = errors.New("原密码错误")
|
||||
ErrAccountDisabled = errors.New("账户已被禁用")
|
||||
)
|
||||
|
||||
type WebTokenPair struct {
|
||||
AccessToken string `json:"access_token"`
|
||||
RefreshToken string `json:"refresh_token"`
|
||||
ExpiresAt time.Time `json:"expires_at"`
|
||||
User map[string]any `json:"user"`
|
||||
}
|
||||
|
||||
func (s *Web) Login(username, password string) (*WebTokenPair, error) {
|
||||
var user model.WebUser
|
||||
if err := s.DB.Where("account = ? AND deleted_at IS NULL", strings.TrimSpace(username)).First(&user).Error; err != nil {
|
||||
return nil, ErrUnauthorized
|
||||
}
|
||||
if !user.Enabled {
|
||||
return nil, ErrAccountDisabled
|
||||
}
|
||||
if !s.Passwords.Verify(user.PasswordHash, password) {
|
||||
return nil, ErrUnauthorized
|
||||
}
|
||||
if user.LastOnlineAt == nil || user.LastOnlineAt.Before(time.Now().Add(-5*time.Minute)) {
|
||||
_ = s.DB.Model(&user).Update("last_online_at", time.Now()).Error
|
||||
}
|
||||
return s.issue(&user)
|
||||
}
|
||||
|
||||
func (s *Web) Refresh(refreshToken string) (*WebTokenPair, error) {
|
||||
var record model.WebRefreshToken
|
||||
if err := s.DB.Where("token_hash = ? AND revoked_at IS NULL AND expires_at > ?", security.HashToken(refreshToken), time.Now()).First(&record).Error; err != nil {
|
||||
return nil, ErrUnauthorized
|
||||
}
|
||||
var user model.WebUser
|
||||
if err := s.DB.Where("id = ? AND enabled = true AND deleted_at IS NULL", record.UserID).First(&user).Error; err != nil || user.SessionVersion != record.SessionVersion {
|
||||
return nil, ErrUnauthorized
|
||||
}
|
||||
now := time.Now()
|
||||
if err := s.DB.Model(&record).Update("revoked_at", now).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return s.issue(&user)
|
||||
}
|
||||
|
||||
func (s *Web) Logout(refreshToken string) error {
|
||||
if refreshToken == "" {
|
||||
return nil
|
||||
}
|
||||
return s.DB.Model(&model.WebRefreshToken{}).Where("token_hash = ? AND revoked_at IS NULL", security.HashToken(refreshToken)).Update("revoked_at", time.Now()).Error
|
||||
}
|
||||
|
||||
func (s *Web) Authenticate(raw string) (*model.WebUser, error) {
|
||||
claims, err := s.Tokens.ParseAccessToken(raw)
|
||||
if err != nil {
|
||||
return nil, ErrUnauthorized
|
||||
}
|
||||
id, err := uuid.Parse(claims.UserID)
|
||||
if err != nil {
|
||||
return nil, ErrUnauthorized
|
||||
}
|
||||
var user model.WebUser
|
||||
if err := s.DB.Where("id = ? AND enabled = true AND deleted_at IS NULL", id).First(&user).Error; err != nil || user.SessionVersion != claims.SessionVersion {
|
||||
return nil, ErrUnauthorized
|
||||
}
|
||||
return &user, nil
|
||||
}
|
||||
|
||||
func (s *Web) issue(user *model.WebUser) (*WebTokenPair, error) {
|
||||
access, expires, err := s.Tokens.NewAccessToken(user.ID.String(), user.UID, user.Account, user.SessionVersion)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
opaque := security.TokenService{RefreshTTL: s.RefreshTTL}
|
||||
refresh, hash, refreshExpires, err := opaque.NewRefreshToken()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
record := model.WebRefreshToken{UserID: user.ID, TokenHash: hash, SessionVersion: user.SessionVersion, ExpiresAt: refreshExpires}
|
||||
if err := s.DB.Create(&record).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &WebTokenPair{AccessToken: access, RefreshToken: refresh, ExpiresAt: expires, User: map[string]any{"id": user.ID, "uid": user.UID, "username": user.Username, "account": user.Account, "avatar_url": user.AvatarURL}}, nil
|
||||
}
|
||||
|
||||
// Account 返回用户资料、永久积分余额、每日限额和当日消耗数据。
|
||||
func (s *Web) Account(userID uuid.UUID) (map[string]any, error) {
|
||||
result := map[string]any{}
|
||||
err := s.DB.Raw(`SELECT id,uid,username,account,avatar_url,point_balance,daily_limit,last_online_at,
|
||||
coalesce((SELECT sum(-hold.change_amount) FROM point_ledger hold
|
||||
WHERE hold.user_id=web_users.id AND hold.business_type='generation_hold'
|
||||
AND hold.created_at>=CURRENT_DATE AND hold.created_at<CURRENT_DATE+INTERVAL '1 day'
|
||||
AND NOT EXISTS (SELECT 1 FROM point_ledger refund WHERE refund.idempotency_key='generation:refund:'||hold.business_id)),0) AS today_consumption
|
||||
FROM web_users WHERE id=? AND deleted_at IS NULL`, userID).Scan(&result).Error
|
||||
return result, err
|
||||
}
|
||||
|
||||
func (s *Web) UpdateProfile(userID uuid.UUID, username string) (map[string]any, error) {
|
||||
username = strings.TrimSpace(username)
|
||||
if count := utf8.RuneCountInString(username); count < 5 || count > 12 {
|
||||
return nil, ErrInvalidUsername
|
||||
}
|
||||
var duplicates int64
|
||||
if err := s.DB.Table("web_users").Where("username = ? AND id <> ?", username, userID).Count(&duplicates).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if duplicates > 0 {
|
||||
return nil, ErrUsernameExists
|
||||
}
|
||||
if err := s.DB.Table("web_users").Where("id = ? AND deleted_at IS NULL", userID).Update("username", username).Error; err != nil {
|
||||
var raced int64
|
||||
if s.DB.Table("web_users").Where("username = ? AND id <> ?", username, userID).Count(&raced).Error == nil && raced > 0 {
|
||||
return nil, ErrUsernameExists
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
return s.Account(userID)
|
||||
}
|
||||
|
||||
func (s *Web) ChangePassword(userID uuid.UUID, originalPassword, newPassword string) error {
|
||||
if len(newPassword) < 8 || len(newPassword) > 128 {
|
||||
return errors.New("新密码长度必须为 8~128 位")
|
||||
}
|
||||
var user model.WebUser
|
||||
if err := s.DB.Where("id = ? AND enabled = true AND deleted_at IS NULL", userID).First(&user).Error; err != nil {
|
||||
return ErrUnauthorized
|
||||
}
|
||||
if !s.Passwords.Verify(user.PasswordHash, originalPassword) {
|
||||
return ErrOriginalPassword
|
||||
}
|
||||
hash, err := s.Passwords.Hash(newPassword)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return s.DB.Model(&user).Update("password_hash", hash).Error
|
||||
}
|
||||
|
||||
func (s *Web) AvatarKey(userID uuid.UUID) (string, error) {
|
||||
var key string
|
||||
err := s.DB.Raw("SELECT coalesce(avatar_key, '') FROM web_users WHERE id = ? AND deleted_at IS NULL", userID).Scan(&key).Error
|
||||
return key, err
|
||||
}
|
||||
|
||||
func (s *Web) UpdateAvatar(userID uuid.UUID, url, key string) error {
|
||||
return s.DB.Table("web_users").Where("id = ? AND deleted_at IS NULL", userID).Updates(map[string]any{"avatar_url": url, "avatar_key": key}).Error
|
||||
}
|
||||
|
||||
func (s *Web) Usage30Days(userID uuid.UUID) ([]map[string]any, error) {
|
||||
items := make([]map[string]any, 0)
|
||||
err := s.DB.Raw(`WITH days AS (SELECT generate_series((CURRENT_DATE-INTERVAL '29 days')::date,CURRENT_DATE::date,'1 day')::date AS day),
|
||||
stats AS (SELECT hold.created_at::date AS day,coalesce(sum(-hold.change_amount),0) AS consumption
|
||||
FROM point_ledger hold WHERE hold.user_id=? AND hold.business_type='generation_hold'
|
||||
AND hold.created_at>=CURRENT_DATE-INTERVAL '29 days'
|
||||
AND NOT EXISTS (SELECT 1 FROM point_ledger refund WHERE refund.idempotency_key='generation:refund:'||hold.business_id)
|
||||
GROUP BY hold.created_at::date)
|
||||
SELECT days.day,coalesce(stats.consumption,0) AS consumption FROM days LEFT JOIN stats USING(day) ORDER BY days.day`, userID).Scan(&items).Error
|
||||
return items, err
|
||||
}
|
||||
|
||||
// ConsumptionRecordPage 描述用户积分流水的服务端分页结果。
|
||||
type ConsumptionRecordPage struct {
|
||||
Items []map[string]any `json:"items"`
|
||||
Total int64 `json:"total"`
|
||||
Page int `json:"page"`
|
||||
PageSize int `json:"page_size"`
|
||||
}
|
||||
|
||||
// normalizeConsumptionPagination 约束积分流水页码和每页数量,避免无效或过大的查询。
|
||||
func normalizeConsumptionPagination(page, pageSize int) (int, int) {
|
||||
if page < 1 {
|
||||
page = 1
|
||||
}
|
||||
if pageSize < 1 {
|
||||
pageSize = 10
|
||||
}
|
||||
if pageSize > 100 {
|
||||
pageSize = 100
|
||||
}
|
||||
return page, pageSize
|
||||
}
|
||||
|
||||
// ConsumptionRecords 返回用户消费、退款、兑换及后台发放形成的积分变动分页流水。
|
||||
func (s *Web) ConsumptionRecords(userID uuid.UUID, page, pageSize int) (ConsumptionRecordPage, error) {
|
||||
page, pageSize = normalizeConsumptionPagination(page, pageSize)
|
||||
result := ConsumptionRecordPage{Items: make([]map[string]any, 0), Page: page, PageSize: pageSize}
|
||||
ledger := s.DB.Table("point_ledger").
|
||||
Where("user_id=? AND business_type IN ?", userID, []string{"generation_hold", "generation_refund", "redemption", "admin_grant"})
|
||||
if err := ledger.Count(&result.Total).Error; err != nil {
|
||||
return result, err
|
||||
}
|
||||
offset := (page - 1) * pageSize
|
||||
err := s.DB.Raw(`SELECT id,abs(change_amount)::text AS amount,change_amount::text AS change_amount,
|
||||
balance_after::text AS balance_after,business_type,remark,created_at
|
||||
FROM point_ledger
|
||||
WHERE user_id=? AND business_type IN ('generation_hold','generation_refund','redemption','admin_grant')
|
||||
ORDER BY created_at DESC,id DESC LIMIT ? OFFSET ?`, userID, pageSize, offset).Scan(&result.Items).Error
|
||||
return result, err
|
||||
}
|
||||
|
||||
// Redeem 校验兑换码并将对应积分以永久有效批次计入用户账户。
|
||||
func (s *Web) Redeem(userID uuid.UUID, code string) (map[string]any, error) {
|
||||
normalized := strings.ReplaceAll(strings.TrimSpace(code), "-", "")
|
||||
if normalized == "" {
|
||||
return nil, errors.New("请输入兑换码")
|
||||
}
|
||||
sum := sha256.Sum256([]byte(normalized))
|
||||
hash := hex.EncodeToString(sum[:])
|
||||
result := map[string]any{}
|
||||
err := s.DB.Transaction(func(tx *gorm.DB) error {
|
||||
var record struct {
|
||||
ID uuid.UUID
|
||||
Points string
|
||||
Status string
|
||||
ExpiresAt *time.Time
|
||||
}
|
||||
if err := tx.Table("redemption_codes").Clauses(clause.Locking{Strength: "UPDATE"}).Select("id,points::text AS points,status,expires_at").Where("code_hash=?", hash).Take(&record).Error; err != nil {
|
||||
return errors.New("兑换码不存在")
|
||||
}
|
||||
if record.Status != "unused" {
|
||||
return errors.New("兑换码已使用或已过期")
|
||||
}
|
||||
if record.ExpiresAt != nil && !record.ExpiresAt.After(time.Now()) {
|
||||
if err := tx.Table("redemption_codes").Where("id=?", record.ID).Update("status", "expired").Error; err != nil {
|
||||
return err
|
||||
}
|
||||
return errors.New("兑换码已过期")
|
||||
}
|
||||
redeemedAt := time.Now()
|
||||
balance, credited, err := billing.CreditRedemptionPoints(tx, userID, record.ID, record.Points)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if !credited {
|
||||
return errors.New("兑换码积分已入账")
|
||||
}
|
||||
if err := tx.Table("redemption_codes").Where("id=?", record.ID).Updates(map[string]any{"status": "used", "redeemed_by": userID, "redeemed_at": redeemedAt}).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
result = map[string]any{"points": record.Points, "balance": balance}
|
||||
return nil
|
||||
})
|
||||
return result, err
|
||||
}
|
||||
@@ -0,0 +1,27 @@
|
||||
// 用户积分流水分页测试,验证无效参数修正和每页数量上限。
|
||||
package service
|
||||
|
||||
import "testing"
|
||||
|
||||
// TestNormalizeConsumptionPagination 验证积分流水分页参数始终落在允许范围内。
|
||||
func TestNormalizeConsumptionPagination(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
page int
|
||||
pageSize int
|
||||
wantPage int
|
||||
wantPageSize int
|
||||
}{
|
||||
{name: "defaults", page: 0, pageSize: 0, wantPage: 1, wantPageSize: 10},
|
||||
{name: "keeps valid values", page: 3, pageSize: 20, wantPage: 3, wantPageSize: 20},
|
||||
{name: "caps page size", page: 2, pageSize: 500, wantPage: 2, wantPageSize: 100},
|
||||
}
|
||||
for _, test := range tests {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
page, pageSize := normalizeConsumptionPagination(test.page, test.pageSize)
|
||||
if page != test.wantPage || pageSize != test.wantPageSize {
|
||||
t.Fatalf("pagination = (%d, %d), want (%d, %d)", page, pageSize, test.wantPage, test.wantPageSize)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,213 @@
|
||||
// Package storage 封装腾讯云 COS 对象存储访问,统一提供上传、读取、删除和公开地址生成能力。
|
||||
package storage
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"os"
|
||||
"path"
|
||||
"strings"
|
||||
|
||||
"juhe-factory/api/internal/config"
|
||||
|
||||
"github.com/aws/aws-sdk-go-v2/aws"
|
||||
awsconfig "github.com/aws/aws-sdk-go-v2/config"
|
||||
"github.com/aws/aws-sdk-go-v2/credentials"
|
||||
"github.com/aws/aws-sdk-go-v2/service/s3"
|
||||
)
|
||||
|
||||
const publicMediaCacheControl = "public, max-age=31536000, immutable"
|
||||
|
||||
// COS 保存腾讯云 COS 客户端、存储桶、公开域名和媒体体积限制。
|
||||
type COS struct {
|
||||
client *s3.Client
|
||||
bucket string
|
||||
publicBaseURL string
|
||||
maxImageBytes int64
|
||||
maxVideoBytes int64
|
||||
maxAudioBytes int64
|
||||
}
|
||||
|
||||
// NewCOS 根据服务端配置初始化腾讯云 COS;完全未配置时返回 nil,部分配置时返回错误。
|
||||
func NewCOS(ctx context.Context, cfg config.Config) (*COS, error) {
|
||||
configured := cfg.COSEndpoint != "" || cfg.COSSecretID != "" || cfg.COSSecretKey != "" || cfg.COSBucket != "" || cfg.COSRegion != ""
|
||||
if !configured {
|
||||
return nil, nil
|
||||
}
|
||||
if cfg.COSEndpoint == "" || cfg.COSSecretID == "" || cfg.COSSecretKey == "" || cfg.COSBucket == "" || cfg.COSRegion == "" {
|
||||
return nil, fmt.Errorf("COS 配置不完整,必须同时配置 Region、Endpoint、SecretId、SecretKey 和 Bucket")
|
||||
}
|
||||
if cfg.COSPublicBaseURL == "" {
|
||||
return nil, fmt.Errorf("视频转绘需要可供模型读取的 COS_PUBLIC_BASE_URL")
|
||||
}
|
||||
endpoint := normalizeHTTPURL(cfg.COSEndpoint)
|
||||
publicBaseURL := normalizeHTTPURL(cfg.COSPublicBaseURL)
|
||||
awsCfg, err := awsconfig.LoadDefaultConfig(ctx,
|
||||
awsconfig.WithRegion(cfg.COSRegion),
|
||||
awsconfig.WithCredentialsProvider(credentials.NewStaticCredentialsProvider(cfg.COSSecretID, cfg.COSSecretKey, "")),
|
||||
)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
client := s3.NewFromConfig(awsCfg, func(options *s3.Options) {
|
||||
options.BaseEndpoint = aws.String(endpoint)
|
||||
options.UsePathStyle = false
|
||||
})
|
||||
return &COS{
|
||||
client: client, bucket: cfg.COSBucket, publicBaseURL: publicBaseURL,
|
||||
maxImageBytes: int64(cfg.COSMaxImageSizeMB) * 1024 * 1024,
|
||||
maxVideoBytes: int64(cfg.COSMaxVideoSizeMB) * 1024 * 1024,
|
||||
maxAudioBytes: int64(cfg.COSMaxAudioSizeMB) * 1024 * 1024,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// normalizeHTTPURL 补全缺少协议的 HTTP 地址,并保留显式配置的协议。
|
||||
func normalizeHTTPURL(value string) string {
|
||||
value = strings.TrimSpace(value)
|
||||
if value != "" && !strings.Contains(value, "://") {
|
||||
return "https://" + value
|
||||
}
|
||||
return value
|
||||
}
|
||||
|
||||
// MaxVideoBytes 返回允许处理的视频最大字节数。
|
||||
func (storage *COS) MaxVideoBytes() int64 {
|
||||
if storage == nil {
|
||||
return 0
|
||||
}
|
||||
return storage.maxVideoBytes
|
||||
}
|
||||
|
||||
// MaxAudioBytes 返回允许处理的音频最大字节数。
|
||||
func (storage *COS) MaxAudioBytes() int64 {
|
||||
if storage == nil {
|
||||
return 0
|
||||
}
|
||||
return storage.maxAudioBytes
|
||||
}
|
||||
|
||||
// PublicURL 根据对象键生成无需签名的公开访问地址。
|
||||
func (storage *COS) PublicURL(key string) string {
|
||||
if storage == nil || storage.publicBaseURL == "" {
|
||||
return ""
|
||||
}
|
||||
return strings.TrimRight(storage.publicBaseURL, "/") + "/" + strings.TrimLeft(key, "/")
|
||||
}
|
||||
|
||||
// MaxImageBytes 返回允许处理的图片最大字节数。
|
||||
func (storage *COS) MaxImageBytes() int64 {
|
||||
if storage == nil {
|
||||
return 0
|
||||
}
|
||||
return storage.maxImageBytes
|
||||
}
|
||||
|
||||
// Put 上传对象并返回其公开访问地址。
|
||||
func (storage *COS) Put(ctx context.Context, key, contentType string, body io.Reader, size int64) (string, error) {
|
||||
if storage == nil {
|
||||
return "", fmt.Errorf("COS 未配置")
|
||||
}
|
||||
input := &s3.PutObjectInput{
|
||||
Bucket: &storage.bucket,
|
||||
Key: &key,
|
||||
Body: body,
|
||||
ContentType: &contentType,
|
||||
CacheControl: aws.String(publicMediaCacheControl),
|
||||
}
|
||||
if size >= 0 {
|
||||
input.ContentLength = aws.Int64(size)
|
||||
}
|
||||
if _, err := storage.client.PutObject(ctx, input); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return storage.PublicURL(key), nil
|
||||
}
|
||||
|
||||
// Open 打开对象读取流,并返回内容类型和内容长度;调用方负责关闭读取流。
|
||||
func (storage *COS) Open(ctx context.Context, key string) (io.ReadCloser, string, int64, error) {
|
||||
if storage == nil || strings.TrimSpace(key) == "" {
|
||||
return nil, "", 0, fmt.Errorf("COS 未配置或对象不存在")
|
||||
}
|
||||
result, err := storage.client.GetObject(ctx, &s3.GetObjectInput{Bucket: &storage.bucket, Key: &key})
|
||||
if err != nil {
|
||||
return nil, "", 0, err
|
||||
}
|
||||
return result.Body, aws.ToString(result.ContentType), aws.ToInt64(result.ContentLength), nil
|
||||
}
|
||||
|
||||
// Delete 删除指定对象;空对象键视为无需处理。
|
||||
func (storage *COS) Delete(ctx context.Context, key string) error {
|
||||
if storage == nil || strings.TrimSpace(key) == "" {
|
||||
return nil
|
||||
}
|
||||
_, err := storage.client.DeleteObject(ctx, &s3.DeleteObjectInput{Bucket: &storage.bucket, Key: &key})
|
||||
return err
|
||||
}
|
||||
|
||||
// PutFromURL 下载远端结果到临时文件,校验体积后上传到 COS。
|
||||
func (storage *COS) PutFromURL(ctx context.Context, client *http.Client, sourceURL, key, fallbackContentType string, maxBytes int64) (string, string, string, int64, error) {
|
||||
request, err := http.NewRequestWithContext(ctx, http.MethodGet, sourceURL, nil)
|
||||
if err != nil {
|
||||
return "", "", "", 0, err
|
||||
}
|
||||
response, err := client.Do(request)
|
||||
if err != nil {
|
||||
return "", "", "", 0, err
|
||||
}
|
||||
defer response.Body.Close()
|
||||
if response.StatusCode < 200 || response.StatusCode >= 300 {
|
||||
return "", "", "", 0, fmt.Errorf("下载生成结果失败: HTTP %d", response.StatusCode)
|
||||
}
|
||||
if maxBytes > 0 && response.ContentLength > maxBytes {
|
||||
return "", "", "", 0, fmt.Errorf("生成结果超过允许大小")
|
||||
}
|
||||
contentType := strings.TrimSpace(strings.Split(response.Header.Get("Content-Type"), ";")[0])
|
||||
if contentType == "" || contentType == "application/octet-stream" {
|
||||
contentType = fallbackContentType
|
||||
}
|
||||
key = keyWithContentType(key, contentType)
|
||||
temporary, err := os.CreateTemp("", "jcf-generated-*")
|
||||
if err != nil {
|
||||
return "", "", "", 0, err
|
||||
}
|
||||
temporaryName := temporary.Name()
|
||||
defer os.Remove(temporaryName)
|
||||
reader := io.Reader(response.Body)
|
||||
if maxBytes > 0 {
|
||||
reader = io.LimitReader(response.Body, maxBytes+1)
|
||||
}
|
||||
written, copyErr := io.Copy(temporary, reader)
|
||||
if closeErr := temporary.Close(); copyErr == nil {
|
||||
copyErr = closeErr
|
||||
}
|
||||
if copyErr != nil {
|
||||
return "", "", "", 0, copyErr
|
||||
}
|
||||
if maxBytes > 0 && written > maxBytes {
|
||||
return "", "", "", 0, fmt.Errorf("生成结果超过允许大小")
|
||||
}
|
||||
file, err := os.Open(temporaryName)
|
||||
if err != nil {
|
||||
return "", "", "", 0, err
|
||||
}
|
||||
defer file.Close()
|
||||
url, err := storage.Put(ctx, key, contentType, file, written)
|
||||
if err != nil {
|
||||
return "", "", "", 0, err
|
||||
}
|
||||
return url, key, contentType, written, nil
|
||||
}
|
||||
|
||||
// keyWithContentType 根据真实媒体类型修正对象扩展名。
|
||||
func keyWithContentType(key, contentType string) string {
|
||||
extension := map[string]string{
|
||||
"image/jpeg": ".jpg", "image/png": ".png", "image/webp": ".webp", "image/gif": ".gif",
|
||||
"video/mp4": ".mp4", "video/webm": ".webm", "video/quicktime": ".mov",
|
||||
}[strings.ToLower(contentType)]
|
||||
if extension == "" {
|
||||
return key
|
||||
}
|
||||
return strings.TrimSuffix(key, path.Ext(key)) + extension
|
||||
}
|
||||
@@ -0,0 +1,14 @@
|
||||
// Package storage 验证腾讯云 COS 对象键扩展名处理规则。
|
||||
package storage
|
||||
|
||||
import "testing"
|
||||
|
||||
// TestKeyWithContentTypeMatchesActualMedia 验证远端媒体类型会覆盖错误扩展名。
|
||||
func TestKeyWithContentTypeMatchesActualMedia(t *testing.T) {
|
||||
if got := keyWithContentType("folder/result.png", "image/webp"); got != "folder/result.webp" {
|
||||
t.Fatalf("got %s", got)
|
||||
}
|
||||
if got := keyWithContentType("folder/result.mp4", "video/quicktime"); got != "folder/result.mov" {
|
||||
t.Fatalf("got %s", got)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,75 @@
|
||||
// 本文件负责将剧集视频按清单写入 ZIP,仅打包原始媒体字节,不执行转码或拼接。
|
||||
package videoarchive
|
||||
|
||||
import (
|
||||
"archive/zip"
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"strings"
|
||||
"unicode"
|
||||
)
|
||||
|
||||
const maxFilenameRunes = 180
|
||||
|
||||
// Entry 描述 ZIP 中单个视频的对象存储位置和归档名称。
|
||||
type Entry struct {
|
||||
ObjectKey string
|
||||
Name string
|
||||
}
|
||||
|
||||
// Manifest 描述一个剧集视频压缩包的下载名称和文件清单。
|
||||
type Manifest struct {
|
||||
Name string
|
||||
Entries []Entry
|
||||
}
|
||||
|
||||
// OpenObject 按对象键打开原始媒体流,调用方负责提供具体存储实现。
|
||||
type OpenObject func(context.Context, string) (io.ReadCloser, error)
|
||||
|
||||
// Write 按清单顺序写入 ZIP,并确保每个对象流在复制后立即关闭。
|
||||
func Write(ctx context.Context, target io.Writer, entries []Entry, open OpenObject) error {
|
||||
writer := zip.NewWriter(target)
|
||||
for _, entry := range entries {
|
||||
body, err := open(ctx, entry.ObjectKey)
|
||||
if err != nil {
|
||||
return fmt.Errorf("打开归档视频 %q 失败: %w", entry.Name, err)
|
||||
}
|
||||
header := &zip.FileHeader{Name: entry.Name, Method: zip.Deflate}
|
||||
header.SetMode(0o644)
|
||||
part, createErr := writer.CreateHeader(header)
|
||||
if createErr == nil {
|
||||
_, createErr = io.Copy(part, body)
|
||||
}
|
||||
closeErr := body.Close()
|
||||
if createErr != nil {
|
||||
return fmt.Errorf("写入归档视频 %q 失败: %w", entry.Name, createErr)
|
||||
}
|
||||
if closeErr != nil {
|
||||
return fmt.Errorf("关闭归档视频 %q 失败: %w", entry.Name, closeErr)
|
||||
}
|
||||
}
|
||||
if err := writer.Close(); err != nil {
|
||||
return fmt.Errorf("完成 ZIP 写入失败: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// SafeFilename 清理文件系统不允许的字符并限制长度,避免归档名形成路径或无法保存。
|
||||
func SafeFilename(value, fallback string) string {
|
||||
cleaned := strings.Map(func(r rune) rune {
|
||||
if unicode.IsControl(r) || strings.ContainsRune(`\/:*?"<>|`, r) {
|
||||
return -1
|
||||
}
|
||||
return r
|
||||
}, strings.TrimSpace(value))
|
||||
cleaned = strings.Trim(cleaned, ". ")
|
||||
if cleaned == "" {
|
||||
cleaned = fallback
|
||||
}
|
||||
runes := []rune(cleaned)
|
||||
if len(runes) > maxFilenameRunes {
|
||||
cleaned = string(runes[:maxFilenameRunes])
|
||||
}
|
||||
return cleaned
|
||||
}
|
||||
@@ -0,0 +1,665 @@
|
||||
package worker
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"regexp"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"juhe-factory/api/internal/billing"
|
||||
dramapkg "juhe-factory/api/internal/drama"
|
||||
"juhe-factory/api/internal/model"
|
||||
"juhe-factory/api/internal/provider"
|
||||
"juhe-factory/api/internal/provider/apimart"
|
||||
queuepkg "juhe-factory/api/internal/queue"
|
||||
"juhe-factory/api/internal/security"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/hibiken/asynq"
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/clause"
|
||||
)
|
||||
|
||||
var (
|
||||
errDramaCancelled = errors.New("drama parse cancelled")
|
||||
errDramaInactive = errors.New("drama parse inactive")
|
||||
)
|
||||
|
||||
const maxDramaModelAttempts = 5
|
||||
|
||||
type DramaParse struct {
|
||||
DB *gorm.DB
|
||||
Queue *asynq.Client
|
||||
Encryptor *security.Encryptor
|
||||
Provider apimart.Client
|
||||
WorkerID string
|
||||
}
|
||||
|
||||
type dramaParseSnapshot struct {
|
||||
Task model.DramaParseTask
|
||||
Source model.EpisodeSource
|
||||
EpisodeNo int
|
||||
EpisodeName string
|
||||
ProjectName string
|
||||
StyleName string
|
||||
EraType string
|
||||
CustomEra string
|
||||
ModelName string
|
||||
BaseURL string
|
||||
APIKeyCiphertext string
|
||||
}
|
||||
|
||||
type dramaParseConfig struct {
|
||||
EpisodeNo int
|
||||
EpisodeName string
|
||||
ProjectName string
|
||||
StyleName string
|
||||
EraType string
|
||||
CustomEra string
|
||||
ModelName string
|
||||
BaseURL string
|
||||
APIKeyCiphertext string
|
||||
}
|
||||
|
||||
type parsedEntity struct {
|
||||
CanonicalName string `json:"canonical_name"`
|
||||
Name string `json:"name"`
|
||||
Aliases []string `json:"aliases"`
|
||||
Description string `json:"description"`
|
||||
ImagePrompt string `json:"image_prompt"`
|
||||
Attributes map[string]any `json:"attributes"`
|
||||
}
|
||||
|
||||
type parsedEntities struct {
|
||||
Characters []parsedEntity `json:"characters"`
|
||||
Roles []parsedEntity `json:"roles"`
|
||||
Scenes []parsedEntity `json:"scenes"`
|
||||
Props []parsedEntity `json:"props"`
|
||||
}
|
||||
|
||||
type parsedStoryboard struct {
|
||||
Title string `json:"title"`
|
||||
ScriptContent string `json:"script_content"`
|
||||
Content string `json:"content"`
|
||||
PromptContent string `json:"prompt_content"`
|
||||
Dialogue []map[string]any `json:"dialogue"`
|
||||
AssetNames struct {
|
||||
Characters []string `json:"characters"`
|
||||
Scenes []string `json:"scenes"`
|
||||
Props []string `json:"props"`
|
||||
} `json:"asset_names"`
|
||||
DurationSeconds int `json:"duration_seconds"`
|
||||
SourceExcerpt string `json:"source_excerpt"`
|
||||
}
|
||||
|
||||
type dramaParsedResult struct {
|
||||
Entities parsedEntities `json:"entities"`
|
||||
Storyboards []parsedStoryboard `json:"storyboards"`
|
||||
}
|
||||
|
||||
type tokenTotals struct {
|
||||
Input, Output int64
|
||||
Source string
|
||||
Segments []map[string]any
|
||||
}
|
||||
|
||||
func NewDramaParse(db *gorm.DB, queue *asynq.Client, encryptor *security.Encryptor, client *http.Client) *DramaParse {
|
||||
return &DramaParse{DB: db, Queue: queue, Encryptor: encryptor, Provider: apimart.NewClient(dramaHTTPClient(client)), WorkerID: uuid.NewString()}
|
||||
}
|
||||
|
||||
func dramaHTTPClient(client *http.Client) *http.Client {
|
||||
cloned := *client
|
||||
if transport, ok := client.Transport.(*http.Transport); ok {
|
||||
transport = transport.Clone()
|
||||
transport.TLSHandshakeTimeout = 30 * time.Second
|
||||
transport.ResponseHeaderTimeout = 10 * time.Minute
|
||||
cloned.Transport = transport
|
||||
}
|
||||
return &cloned
|
||||
}
|
||||
|
||||
func (w *DramaParse) Register(mux *asynq.ServeMux) {
|
||||
mux.HandleFunc(queuepkg.TypeParseDramaEpisode, w.parseEpisode)
|
||||
}
|
||||
|
||||
// Recover 恢复排队、等待重试或租约已过期的剧本解析任务。
|
||||
func (w *DramaParse) Recover(ctx context.Context) error {
|
||||
var tasks []model.DramaParseTask
|
||||
if err := w.DB.WithContext(ctx).Where("status='queued' OR (status IN ('running','retry_wait') AND (lease_until IS NULL OR lease_until<?))", time.Now()).Find(&tasks).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
var firstErr error
|
||||
for _, task := range tasks {
|
||||
w.DB.Model(&model.DramaParseTask{}).Where("id=?", task.ID).Updates(map[string]any{"status": "queued", "lease_owner": nil, "lease_until": nil})
|
||||
if err := queuepkg.EnqueueID(w.Queue, queuepkg.TypeParseDramaEpisode, task.ID, 0); err != nil && firstErr == nil {
|
||||
firstErr = err
|
||||
}
|
||||
}
|
||||
return firstErr
|
||||
}
|
||||
|
||||
func (w *DramaParse) parseEpisode(ctx context.Context, raw *asynq.Task) error {
|
||||
taskID, err := decodeID(raw)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
snapshot, err := w.claim(ctx, taskID)
|
||||
if err != nil {
|
||||
if errors.Is(err, errDramaCancelled) {
|
||||
return w.cancel(snapshot.Task)
|
||||
}
|
||||
if errors.Is(err, errDramaInactive) || snapshot.Task.ID == uuid.Nil {
|
||||
return nil
|
||||
}
|
||||
return w.fail(snapshot.Task, "task_claim_failed", err.Error())
|
||||
}
|
||||
apiKey, err := w.Encryptor.Decrypt(snapshot.APIKeyCiphertext)
|
||||
if err != nil {
|
||||
return w.fail(snapshot.Task, "channel_secret_error", "渠道密钥不可用")
|
||||
}
|
||||
workCtx, cancel := context.WithTimeout(ctx, 12*time.Minute)
|
||||
defer cancel()
|
||||
var cancelled atomic.Bool
|
||||
go w.watchCancellation(workCtx, cancel, taskID, &cancelled)
|
||||
result, tokens, contextSnapshot, err := w.callModel(workCtx, snapshot, apiKey)
|
||||
if err != nil {
|
||||
if cancelled.Load() {
|
||||
return w.cancel(snapshot.Task, tokens)
|
||||
}
|
||||
return w.fail(snapshot.Task, classifyDramaError(err), err.Error(), tokens)
|
||||
}
|
||||
if cancelled.Load() {
|
||||
return w.cancel(snapshot.Task, tokens)
|
||||
}
|
||||
if err = w.persist(snapshot, result, tokens, contextSnapshot); err != nil {
|
||||
if errors.Is(err, errDramaCancelled) {
|
||||
return w.cancel(snapshot.Task, tokens)
|
||||
}
|
||||
return w.fail(snapshot.Task, "result_persist_failed", err.Error())
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (w *DramaParse) claim(ctx context.Context, taskID uuid.UUID) (dramaParseSnapshot, error) {
|
||||
var snapshot dramaParseSnapshot
|
||||
err := w.DB.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
var task model.DramaParseTask
|
||||
if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Where("id=?", taskID).Take(&task).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
snapshot.Task = task
|
||||
if task.Status == "cancel_requested" {
|
||||
return errDramaCancelled
|
||||
}
|
||||
if task.Status != "queued" && task.Status != "retry_wait" {
|
||||
return errDramaInactive
|
||||
}
|
||||
now := time.Now()
|
||||
lease := now.Add(13 * time.Minute)
|
||||
if err := tx.Model(&task).Updates(map[string]any{"status": "running", "attempt_count": gorm.Expr("attempt_count+1"), "started_at": now, "lease_owner": w.WorkerID, "lease_until": lease, "error_code": nil, "error_message": nil}).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if err := tx.Model(&model.ProjectEpisode{}).Where("id=?", task.EpisodeID).Update("analysis_message", "剧本解析中").Error; err != nil {
|
||||
return err
|
||||
}
|
||||
snapshot.Task.Status = "running"
|
||||
if err := tx.Where("episode_id=? AND content_sha256=?", task.EpisodeID, task.SourceSHA256).Take(&snapshot.Source).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
var config dramaParseConfig
|
||||
result := tx.Raw(`SELECT episode.episode_no,episode.name AS episode_name,project.name AS project_name,style.name AS style_name,
|
||||
project.era_type,coalesce(project.custom_era,'') AS custom_era,model.name AS model_name,channel.base_url,channel.api_key_ciphertext
|
||||
FROM project_episodes episode JOIN creative_projects project ON project.id=episode.project_id
|
||||
JOIN project_styles style ON style.id=project.style_id
|
||||
JOIN models model ON model.id=? AND model.enabled=true AND model.deleted_at IS NULL JOIN channels channel ON channel.id=model.channel_id AND channel.id=? AND channel.enabled=true AND channel.deleted_at IS NULL
|
||||
WHERE episode.id=? AND project.id=? AND project.project_type='premium_drama'`, task.ModelID, task.ChannelID, task.EpisodeID, task.ProjectID).Scan(&config)
|
||||
if result.Error != nil {
|
||||
return result.Error
|
||||
}
|
||||
if result.RowsAffected == 0 {
|
||||
return errors.New("解析模型或渠道配置不可用")
|
||||
}
|
||||
snapshot.EpisodeNo = config.EpisodeNo
|
||||
snapshot.EpisodeName = config.EpisodeName
|
||||
snapshot.ProjectName = config.ProjectName
|
||||
snapshot.StyleName = config.StyleName
|
||||
snapshot.EraType = config.EraType
|
||||
snapshot.CustomEra = config.CustomEra
|
||||
snapshot.ModelName = config.ModelName
|
||||
snapshot.BaseURL = config.BaseURL
|
||||
snapshot.APIKeyCiphertext = config.APIKeyCiphertext
|
||||
return nil
|
||||
})
|
||||
return snapshot, err
|
||||
}
|
||||
|
||||
func (w *DramaParse) watchCancellation(ctx context.Context, cancel context.CancelFunc, taskID uuid.UUID, flag *atomic.Bool) {
|
||||
ticker := time.NewTicker(time.Second)
|
||||
defer ticker.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-ticker.C:
|
||||
var status string
|
||||
if w.DB.Model(&model.DramaParseTask{}).Where("id=?", taskID).Pluck("status", &status).Error == nil && status == "cancel_requested" {
|
||||
flag.Store(true)
|
||||
cancel()
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (w *DramaParse) callModel(ctx context.Context, snapshot dramaParseSnapshot, apiKey string) (dramaParsedResult, tokenTotals, json.RawMessage, error) {
|
||||
if err := dramapkg.ValidateEpisodeContent(snapshot.Source.RawContent); err != nil {
|
||||
return dramaParsedResult{}, tokenTotals{}, nil, err
|
||||
}
|
||||
assets, err := w.buildAssetContext(snapshot.Task.ProjectID)
|
||||
if err != nil {
|
||||
return dramaParsedResult{}, tokenTotals{}, nil, err
|
||||
}
|
||||
contextJSON, _ := json.Marshal(map[string]any{"assets": assets})
|
||||
userContent := buildDramaUserPrompt(snapshot, snapshot.Source.RawContent, assets)
|
||||
payload := map[string]any{
|
||||
"model": snapshot.ModelName,
|
||||
"messages": []any{map[string]any{"role": "system", "content": snapshot.Task.PromptSnapshot}, map[string]any{"role": "user", "content": userContent}},
|
||||
"temperature": 0.2,
|
||||
"response_format": map[string]any{"type": "json_object"},
|
||||
}
|
||||
pricing, err := billing.ParseTextPricingSnapshot(snapshot.Task.BillingSnapshot)
|
||||
if err != nil {
|
||||
return dramaParsedResult{}, tokenTotals{}, contextJSON, err
|
||||
}
|
||||
if pricing.Mode == billing.TextBillingPerToken {
|
||||
if err := w.DB.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
_, reserveErr := billing.ReserveTextCall(tx, snapshot.Task.UserID, snapshot.Task.ID.String(), "parse", pricing, billing.EstimateTextTokens(snapshot.Task.PromptSnapshot+userContent), "剧本解析")
|
||||
return reserveErr
|
||||
}); err != nil {
|
||||
return dramaParsedResult{}, tokenTotals{}, contextJSON, err
|
||||
}
|
||||
}
|
||||
chat, callErr := w.chatWithRetry(ctx, snapshot.BaseURL, apiKey, payload)
|
||||
if callErr != nil {
|
||||
return dramaParsedResult{}, tokenTotals{}, contextJSON, callErr
|
||||
}
|
||||
input, output := chat.InputTokens, chat.OutputTokens
|
||||
source := "upstream"
|
||||
if input == 0 {
|
||||
input = billing.EstimateTextTokens(snapshot.Task.PromptSnapshot + userContent)
|
||||
source = "local"
|
||||
}
|
||||
if output == 0 {
|
||||
output = billing.EstimateTextTokens(chat.Content)
|
||||
if source == "upstream" {
|
||||
source = "mixed"
|
||||
}
|
||||
}
|
||||
if pricing.Mode == billing.TextBillingPerToken {
|
||||
if err := w.DB.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
settled, settleErr := billing.SettleTextCall(tx, snapshot.Task.UserID, snapshot.Task.ID.String(), "parse", pricing, billing.TextUsage{Input: chat.InputTokens, Output: chat.OutputTokens}, "剧本解析")
|
||||
if settleErr != nil {
|
||||
return settleErr
|
||||
}
|
||||
if err := tx.Model(&model.DramaParseTask{}).Where("id=?", snapshot.Task.ID).Updates(map[string]any{
|
||||
"estimated_points": gorm.Expr("estimated_points+?::numeric", settled),
|
||||
"prepaid_points": gorm.Expr("prepaid_points+?::numeric", settled),
|
||||
"input_tokens": chat.InputTokens, "output_tokens": chat.OutputTokens, "total_tokens": chat.InputTokens + chat.OutputTokens,
|
||||
"token_count_source": source, "usage_raw": gorm.Expr("?::jsonb", string(chat.UsageRaw)),
|
||||
}).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}); err != nil {
|
||||
return dramaParsedResult{}, tokenTotals{}, contextJSON, err
|
||||
}
|
||||
}
|
||||
usage := map[string]any{"index": 1, "input_tokens": input, "output_tokens": output, "source": source}
|
||||
if len(chat.UsageRaw) > 0 {
|
||||
var upstreamUsage any
|
||||
if json.Unmarshal(chat.UsageRaw, &upstreamUsage) == nil {
|
||||
usage["upstream_usage"] = upstreamUsage
|
||||
}
|
||||
}
|
||||
totals := tokenTotals{Input: input, Output: output, Source: source, Segments: []map[string]any{usage}}
|
||||
parsed, parseErr := parseDramaJSON(chat.Content)
|
||||
if parseErr != nil {
|
||||
finishReason := strings.TrimSpace(chat.FinishReason)
|
||||
if finishReason == "" {
|
||||
finishReason = "未知"
|
||||
}
|
||||
return dramaParsedResult{}, totals, contextJSON, fmt.Errorf("文本模型返回无法解析(结束原因:%s): %w", finishReason, parseErr)
|
||||
}
|
||||
if len(parsed.Storyboards) == 0 {
|
||||
return parsed, totals, contextJSON, errors.New("文本模型未返回有效分镜")
|
||||
}
|
||||
return parsed, totals, contextJSON, nil
|
||||
}
|
||||
|
||||
func (w *DramaParse) chatWithRetry(ctx context.Context, baseURL, apiKey string, payload map[string]any) (apimart.ChatResult, error) {
|
||||
var result apimart.ChatResult
|
||||
var err error
|
||||
for attempt := 0; attempt < maxDramaModelAttempts; attempt++ {
|
||||
if attempt > 0 {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return result, ctx.Err()
|
||||
case <-time.After(time.Duration(1<<attempt) * time.Second):
|
||||
}
|
||||
}
|
||||
result, err = w.Provider.ChatWithUsage(ctx, baseURL, apiKey, payload)
|
||||
if err == nil {
|
||||
return result, nil
|
||||
}
|
||||
var httpErr *provider.HTTPError
|
||||
if errors.As(err, &httpErr) && httpErr.StatusCode < 500 && httpErr.StatusCode != http.StatusTooManyRequests {
|
||||
return result, err
|
||||
}
|
||||
}
|
||||
return result, err
|
||||
}
|
||||
|
||||
func (w *DramaParse) buildAssetContext(projectID uuid.UUID) ([]map[string]any, error) {
|
||||
assets := make([]map[string]any, 0)
|
||||
if err := w.DB.Table("project_assets").Select("asset_type,name,aliases,description,attributes").Where("project_id=? AND deleted_at IS NULL", projectID).Order("asset_type,name").Find(&assets).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return assets, nil
|
||||
}
|
||||
|
||||
func buildDramaUserPrompt(snapshot dramaParseSnapshot, source string, assets []map[string]any) string {
|
||||
assetsJSON, _ := json.Marshal(assets)
|
||||
return fmt.Sprintf("项目:%s\n视觉风格:%s\n故事时代:%s\n当前剧集:第%d集 %s\n项目资产:\n%s\n\n当前原文:\n%s\n\n解析要求:必须按上述视觉风格和故事时代解析角色、场景、道具及分镜,并在资产的description、image_prompt和分镜的prompt_content中体现与之一致的可视化特征。\n只返回合法JSON,顶层只能包含entities和storyboards。entities包含characters、scenes、props,每个资产包含name、description、image_prompt;storyboards每项包含title、script_content、prompt_content、dialogue、asset_names、duration_seconds、source_excerpt。禁止返回摘要或其他顶层字段,不要返回Markdown或解释。", snapshot.ProjectName, snapshot.StyleName, dramaEraName(snapshot.EraType, snapshot.CustomEra), snapshot.EpisodeNo, snapshot.EpisodeName, string(assetsJSON), source)
|
||||
}
|
||||
|
||||
func dramaEraName(eraType, customEra string) string {
|
||||
if eraType == "other" && strings.TrimSpace(customEra) != "" {
|
||||
return strings.TrimSpace(customEra)
|
||||
}
|
||||
if name := map[string]string{
|
||||
"modern_city": "现代都市",
|
||||
"ancient_history": "古代历史",
|
||||
"ancient_fantasy": "古代玄幻",
|
||||
"ancient_xianxia": "古代仙侠",
|
||||
"future_scifi": "未来科幻",
|
||||
}[eraType]; name != "" {
|
||||
return name
|
||||
}
|
||||
return eraType
|
||||
}
|
||||
|
||||
func parseDramaJSON(value string) (dramaParsedResult, error) {
|
||||
value = strings.TrimSpace(value)
|
||||
think := regexp.MustCompile(`(?s)<think>.*?</think>`)
|
||||
value = think.ReplaceAllString(value, "")
|
||||
value = strings.TrimPrefix(value, "```json")
|
||||
value = strings.TrimPrefix(value, "```")
|
||||
value = strings.TrimSuffix(value, "```")
|
||||
value = strings.TrimSpace(value)
|
||||
start := strings.Index(value, "{")
|
||||
end := strings.LastIndex(value, "}")
|
||||
if start < 0 || end <= start {
|
||||
return dramaParsedResult{}, errors.New("文本模型返回内容不是JSON")
|
||||
}
|
||||
var result dramaParsedResult
|
||||
if err := json.Unmarshal([]byte(value[start:end+1]), &result); err != nil {
|
||||
return result, fmt.Errorf("文本模型返回JSON无效: %w", err)
|
||||
}
|
||||
result.Entities.Characters = append(result.Entities.Characters, result.Entities.Roles...)
|
||||
valid := result.Storyboards[:0]
|
||||
for _, item := range result.Storyboards {
|
||||
item.ScriptContent = strings.TrimSpace(firstNonEmpty(item.ScriptContent, item.Content))
|
||||
item.PromptContent = strings.TrimSpace(item.PromptContent)
|
||||
if item.ScriptContent == "" && item.PromptContent == "" {
|
||||
continue
|
||||
}
|
||||
if item.DurationSeconds < 5 || item.DurationSeconds > 15 {
|
||||
item.DurationSeconds = 5
|
||||
}
|
||||
valid = append(valid, item)
|
||||
}
|
||||
result.Storyboards = valid
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func (w *DramaParse) persist(snapshot dramaParseSnapshot, result dramaParsedResult, tokens tokenTotals, contextSnapshot json.RawMessage) error {
|
||||
normalized, _ := json.Marshal(result)
|
||||
usage, _ := json.Marshal(map[string]any{"segments": tokens.Segments})
|
||||
return w.DB.Transaction(func(tx *gorm.DB) error {
|
||||
var task model.DramaParseTask
|
||||
if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Where("id=?", snapshot.Task.ID).Take(&task).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if task.Status == "cancel_requested" {
|
||||
return errDramaCancelled
|
||||
}
|
||||
if task.Status != "running" {
|
||||
return errors.New("任务状态已变化")
|
||||
}
|
||||
assetIDs, err := w.upsertEntities(tx, snapshot, result.Entities)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
now := time.Now()
|
||||
protected := make([]model.EpisodeStoryboard, 0)
|
||||
if err := tx.Where("episode_id=? AND deleted_at IS NULL AND user_edited=true", snapshot.Task.EpisodeID).Find(&protected).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
protectedBySequence := make(map[int]model.EpisodeStoryboard, len(protected))
|
||||
for _, storyboard := range protected {
|
||||
protectedBySequence[storyboard.SequenceNo] = storyboard
|
||||
}
|
||||
if err := tx.Model(&model.EpisodeStoryboard{}).Where("episode_id=? AND deleted_at IS NULL AND user_edited=false", snapshot.Task.EpisodeID).Update("deleted_at", now).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
startMS := int64(0)
|
||||
for index, item := range result.Storyboards {
|
||||
sequence := index + 1
|
||||
if existing, ok := protectedBySequence[sequence]; ok {
|
||||
duration := existing.DurationSeconds
|
||||
if duration < 5 || duration > 15 {
|
||||
duration = 5
|
||||
}
|
||||
endMS := startMS + int64(duration*1000)
|
||||
if err := tx.Model(&model.EpisodeStoryboard{}).Where("id=?", existing.ID).Updates(map[string]any{"start_ms": startMS, "end_ms": endMS, "duration_seconds": duration}).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
startMS = endMS
|
||||
continue
|
||||
}
|
||||
duration := item.DurationSeconds
|
||||
if duration < 5 || duration > 15 {
|
||||
duration = 5
|
||||
}
|
||||
refs := buildAssetRefs(item, assetIDs)
|
||||
refsJSON, _ := json.Marshal(refs)
|
||||
dialogueJSON, _ := json.Marshal(item.Dialogue)
|
||||
endMS := startMS + int64(duration*1000)
|
||||
storyboard := model.EpisodeStoryboard{ID: uuid.New(), EpisodeID: &snapshot.Task.EpisodeID, SequenceNo: sequence, StableKey: uuid.NewString(), StartMS: startMS, EndMS: endMS, DurationSeconds: duration, Title: truncateWorker(item.Title, 160), ScriptContent: item.ScriptContent, PromptContent: firstNonEmpty(item.PromptContent, item.ScriptContent), SourceExcerpt: item.SourceExcerpt, Dialogue: dialogueJSON, AssetRefs: refsJSON, Locked: false, Status: "idle", SourceParseTaskID: &task.ID}
|
||||
if err := tx.Create(&storyboard).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
startMS = endMS
|
||||
}
|
||||
total := tokens.Input + tokens.Output
|
||||
finished := time.Now()
|
||||
if err := tx.Model(&task).Updates(map[string]any{"status": "succeeded", "normalized_result": gorm.Expr("?::jsonb", string(normalized)), "context_snapshot": gorm.Expr("?::jsonb", string(contextSnapshot)), "input_tokens": tokens.Input, "output_tokens": tokens.Output, "total_tokens": total, "token_count_source": tokens.Source, "usage_raw": gorm.Expr("?::jsonb", string(usage)), "actual_points": task.PrepaidPoints, "request_char_count": len([]rune(snapshot.Source.RawContent)), "response_char_count": len([]rune(string(normalized))), "finished_at": finished, "lease_owner": nil, "lease_until": nil}).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Model(&model.ProjectEpisode{}).Where("id=?", snapshot.Task.EpisodeID).Updates(map[string]any{"status": "review", "analysis_message": "剧本解析完成"}).Error
|
||||
})
|
||||
}
|
||||
|
||||
func (w *DramaParse) upsertEntities(tx *gorm.DB, snapshot dramaParseSnapshot, entities parsedEntities) (map[string]uuid.UUID, error) {
|
||||
result := map[string]uuid.UUID{}
|
||||
groups := []struct {
|
||||
Type string
|
||||
Items []parsedEntity
|
||||
}{{"character", entities.Characters}, {"scene", entities.Scenes}, {"prop", entities.Props}}
|
||||
for _, group := range groups {
|
||||
for _, item := range group.Items {
|
||||
rawName := strings.TrimSpace(firstNonEmpty(item.CanonicalName, item.Name))
|
||||
if rawName == "" {
|
||||
continue
|
||||
}
|
||||
imagePrompt := strings.TrimSpace(firstNonEmpty(item.ImagePrompt, item.Description))
|
||||
name := truncateWorker(rawName, 200)
|
||||
aliasValues := append([]string{}, item.Aliases...)
|
||||
if rawName != name {
|
||||
aliasValues = append(aliasValues, rawName)
|
||||
}
|
||||
var existing model.ProjectAsset
|
||||
err := tx.Where("project_id=? AND asset_type=? AND lower(name)=lower(?) AND deleted_at IS NULL", snapshot.Task.ProjectID, group.Type, name).Take(&existing).Error
|
||||
aliases, _ := json.Marshal(uniqueStrings(aliasValues))
|
||||
attributes, _ := json.Marshal(item.Attributes)
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
existing = model.ProjectAsset{ID: uuid.New(), ProjectID: snapshot.Task.ProjectID, SourceEpisodeID: &snapshot.Task.EpisodeID, AssetType: group.Type, Name: name, Description: item.Description, ImagePrompt: imagePrompt, AIDescription: item.Description, Aliases: aliases, Attributes: attributes, Appearances: json.RawMessage("[]"), SourceParseTaskID: &snapshot.Task.ID}
|
||||
if err = tx.Create(&existing).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
} else if err != nil {
|
||||
return nil, err
|
||||
} else {
|
||||
updates := map[string]any{"ai_description": item.Description, "aliases": gorm.Expr("?::jsonb", string(mergeJSONArray(existing.Aliases, aliasValues))), "source_parse_task_id": snapshot.Task.ID}
|
||||
if !existing.UserEdited {
|
||||
updates["description"] = item.Description
|
||||
updates["image_prompt"] = imagePrompt
|
||||
updates["attributes"] = gorm.Expr("?::jsonb", string(attributes))
|
||||
}
|
||||
if err = tx.Model(&existing).Updates(updates).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
result[group.Type+":"+strings.ToLower(rawName)] = existing.ID
|
||||
result[group.Type+":"+strings.ToLower(name)] = existing.ID
|
||||
for _, alias := range aliasValues {
|
||||
result[group.Type+":"+strings.ToLower(strings.TrimSpace(alias))] = existing.ID
|
||||
}
|
||||
}
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func buildAssetRefs(item parsedStoryboard, ids map[string]uuid.UUID) []map[string]any {
|
||||
refs := make([]map[string]any, 0)
|
||||
groups := []struct {
|
||||
Type string
|
||||
Names []string
|
||||
}{{"character", item.AssetNames.Characters}, {"scene", item.AssetNames.Scenes}, {"prop", item.AssetNames.Props}}
|
||||
seen := map[uuid.UUID]bool{}
|
||||
for _, group := range groups {
|
||||
for _, name := range group.Names {
|
||||
if id, ok := ids[group.Type+":"+strings.ToLower(strings.TrimSpace(name))]; ok && !seen[id] {
|
||||
refs = append(refs, map[string]any{"id": id, "name": name, "type": group.Type})
|
||||
seen[id] = true
|
||||
}
|
||||
}
|
||||
}
|
||||
return refs
|
||||
}
|
||||
|
||||
func (w *DramaParse) fail(task model.DramaParseTask, code, message string, tokenUsage ...tokenTotals) error {
|
||||
now := time.Now()
|
||||
return w.DB.Transaction(func(tx *gorm.DB) error {
|
||||
refunded, err := billing.RefundTextCalls(tx, task.UserID, task.ID.String(), "剧本解析失败返还")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
updates := map[string]any{"status": "failed", "actual_points": "0.00", "error_code": code, "error_message": truncateWorker(message, 4000), "finished_at": now, "lease_owner": nil, "lease_until": nil}
|
||||
if refunded {
|
||||
updates["cost_refunded"] = true
|
||||
}
|
||||
if len(tokenUsage) > 0 {
|
||||
addTokenUpdates(updates, tokenUsage[0])
|
||||
}
|
||||
if err := tx.Model(&model.DramaParseTask{}).Where("id=? AND status<>'cancelled'", task.ID).Updates(updates).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Model(&model.ProjectEpisode{}).Where("id=?", task.EpisodeID).Updates(map[string]any{"status": "failed", "analysis_message": truncateWorker(message, 1000)}).Error
|
||||
})
|
||||
}
|
||||
func (w *DramaParse) cancel(task model.DramaParseTask, tokenUsage ...tokenTotals) error {
|
||||
now := time.Now()
|
||||
return w.DB.Transaction(func(tx *gorm.DB) error {
|
||||
amount, err := billing.CommitTextReserves(tx, task.UserID, task.ID.String(), "剧本解析取消")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
updates := map[string]any{"status": "cancelled", "actual_points": amount, "prepaid_points": amount, "finished_at": now, "lease_owner": nil, "lease_until": nil, "error_code": "user_cancelled", "error_message": "用户主动取消,费用不退"}
|
||||
if len(tokenUsage) > 0 {
|
||||
addTokenUpdates(updates, tokenUsage[0])
|
||||
}
|
||||
if err := tx.Model(&model.DramaParseTask{}).Where("id=?", task.ID).Updates(updates).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Model(&model.ProjectEpisode{}).Where("id=?", task.EpisodeID).Updates(map[string]any{"status": "draft", "analysis_message": "解析已取消"}).Error
|
||||
})
|
||||
}
|
||||
|
||||
func addTokenUpdates(updates map[string]any, tokens tokenTotals) {
|
||||
if tokens.Input == 0 && tokens.Output == 0 {
|
||||
return
|
||||
}
|
||||
usage, _ := json.Marshal(map[string]any{"segments": tokens.Segments})
|
||||
updates["input_tokens"] = tokens.Input
|
||||
updates["output_tokens"] = tokens.Output
|
||||
updates["total_tokens"] = tokens.Input + tokens.Output
|
||||
updates["token_count_source"] = tokens.Source
|
||||
updates["usage_raw"] = gorm.Expr("?::jsonb", string(usage))
|
||||
}
|
||||
|
||||
func estimateTextTokens(value string) int64 {
|
||||
return billing.EstimateTextTokens(value)
|
||||
}
|
||||
func uniqueStrings(values []string) []string {
|
||||
seen := map[string]bool{}
|
||||
out := make([]string, 0, len(values))
|
||||
for _, value := range values {
|
||||
value = strings.TrimSpace(value)
|
||||
key := strings.ToLower(value)
|
||||
if value != "" && !seen[key] {
|
||||
seen[key] = true
|
||||
out = append(out, value)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
func mergeJSONArray(raw json.RawMessage, values []string) json.RawMessage {
|
||||
var existing []string
|
||||
_ = json.Unmarshal(raw, &existing)
|
||||
encoded, _ := json.Marshal(uniqueStrings(append(existing, values...)))
|
||||
return encoded
|
||||
}
|
||||
func firstNonEmpty(values ...string) string {
|
||||
for _, value := range values {
|
||||
if strings.TrimSpace(value) != "" {
|
||||
return strings.TrimSpace(value)
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
func normalizeCompare(value string) string {
|
||||
return strings.Join(strings.Fields(strings.ToLower(value)), "")
|
||||
}
|
||||
func truncateWorker(value string, limit int) string {
|
||||
r := []rune(value)
|
||||
if len(r) > limit {
|
||||
return string(r[:limit])
|
||||
}
|
||||
return value
|
||||
}
|
||||
func classifyDramaError(err error) string {
|
||||
var httpErr *provider.HTTPError
|
||||
if errors.As(err, &httpErr) {
|
||||
return "upstream_failed"
|
||||
}
|
||||
if errors.Is(err, context.DeadlineExceeded) {
|
||||
return "upstream_timeout"
|
||||
}
|
||||
return "parse_failed"
|
||||
}
|
||||
@@ -0,0 +1,95 @@
|
||||
package worker
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// TestClassifyDramaErrorTreatsTimeoutAsUpstreamFailure 验证超时不会再进入提交未知状态。
|
||||
func TestClassifyDramaErrorTreatsTimeoutAsUpstreamFailure(t *testing.T) {
|
||||
if code := classifyDramaError(context.DeadlineExceeded); code != "upstream_timeout" {
|
||||
t.Fatalf("timeout code = %q, want upstream_timeout", code)
|
||||
}
|
||||
}
|
||||
|
||||
// TestDramaModelRetryLimit 验证剧本解析调用的自动重试次数。
|
||||
func TestDramaModelRetryLimit(t *testing.T) {
|
||||
if maxDramaModelAttempts != 5 {
|
||||
t.Fatalf("retry attempts = %d, want 5", maxDramaModelAttempts)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDramaHTTPClientAllowsLongTextResponses(t *testing.T) {
|
||||
sharedTransport := &http.Transport{TLSHandshakeTimeout: 10 * time.Second, ResponseHeaderTimeout: time.Minute}
|
||||
shared := &http.Client{Transport: sharedTransport}
|
||||
client := dramaHTTPClient(shared)
|
||||
transport, ok := client.Transport.(*http.Transport)
|
||||
if !ok {
|
||||
t.Fatal("drama client transport was not preserved")
|
||||
}
|
||||
if transport.TLSHandshakeTimeout != 30*time.Second || transport.ResponseHeaderTimeout != 10*time.Minute {
|
||||
t.Fatalf("unexpected drama timeouts: TLS=%s response=%s", transport.TLSHandshakeTimeout, transport.ResponseHeaderTimeout)
|
||||
}
|
||||
if sharedTransport.TLSHandshakeTimeout != 10*time.Second || sharedTransport.ResponseHeaderTimeout != time.Minute {
|
||||
t.Fatal("shared AI transport was modified")
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseDramaJSONNormalizesStoryboard(t *testing.T) {
|
||||
result, err := parseDramaJSON("<think>hidden</think>```json\n{\"entities\":{\"roles\":[{\"name\":\"林夏\",\"image_prompt\":\"二十五岁女性,黑色长发\"}]},\"storyboards\":[{\"content\":\"走进房间\",\"duration_seconds\":99}]}\n```")
|
||||
if err != nil {
|
||||
t.Fatalf("parseDramaJSON returned error: %v", err)
|
||||
}
|
||||
if len(result.Entities.Characters) != 1 || len(result.Storyboards) != 1 {
|
||||
t.Fatalf("unexpected normalized result: %#v", result)
|
||||
}
|
||||
if result.Storyboards[0].ScriptContent != "走进房间" || result.Storyboards[0].DurationSeconds != 5 {
|
||||
t.Fatalf("storyboard was not normalized: %#v", result.Storyboards[0])
|
||||
}
|
||||
if result.Entities.Characters[0].ImagePrompt != "二十五岁女性,黑色长发" {
|
||||
t.Fatalf("asset image prompt was not preserved: %#v", result.Entities.Characters[0])
|
||||
}
|
||||
}
|
||||
|
||||
func TestEstimateTextTokens(t *testing.T) {
|
||||
if estimateTextTokens("中文ABCD") != 3 {
|
||||
t.Fatalf("unexpected token estimate")
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildDramaUserPromptUsesOnlyCurrentContext(t *testing.T) {
|
||||
snapshot := dramaParseSnapshot{
|
||||
ProjectName: "测试项目",
|
||||
StyleName: "国风水墨",
|
||||
EraType: "ancient_xianxia",
|
||||
EpisodeNo: 4,
|
||||
EpisodeName: "再入山门",
|
||||
}
|
||||
prompt := buildDramaUserPrompt(snapshot, "当前集正文", []map[string]any{{"name": "林夏"}})
|
||||
for _, expected := range []string{
|
||||
"视觉风格:国风水墨",
|
||||
"故事时代:古代仙侠",
|
||||
"项目资产:",
|
||||
"林夏",
|
||||
"当前原文:\n当前集正文",
|
||||
"按上述视觉风格和故事时代解析角色、场景、道具及分镜",
|
||||
} {
|
||||
if !strings.Contains(prompt, expected) {
|
||||
t.Fatalf("prompt does not contain %q: %s", expected, prompt)
|
||||
}
|
||||
}
|
||||
for _, forbidden := range []string{"continuity", "解析分段", "前序"} {
|
||||
if strings.Contains(prompt, forbidden) {
|
||||
t.Fatalf("prompt contains forbidden content %q: %s", forbidden, prompt)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestDramaEraNameUsesCustomEra(t *testing.T) {
|
||||
if got := dramaEraName("other", " 民国时期 "); got != "民国时期" {
|
||||
t.Fatalf("unexpected custom era: %q", got)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,850 @@
|
||||
package worker
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"math/rand"
|
||||
"net"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"juhe-factory/api/internal/billing"
|
||||
mediakey "juhe-factory/api/internal/media"
|
||||
"juhe-factory/api/internal/model"
|
||||
"juhe-factory/api/internal/provider"
|
||||
"juhe-factory/api/internal/provider/apimart"
|
||||
queuepkg "juhe-factory/api/internal/queue"
|
||||
"juhe-factory/api/internal/security"
|
||||
"juhe-factory/api/internal/storage"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/hibiken/asynq"
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/clause"
|
||||
)
|
||||
|
||||
var completedStatuses = map[string]bool{"completed": true, "complete": true, "succeeded": true, "success": true, "done": true}
|
||||
var failedStatuses = map[string]bool{"failed": true, "failure": true, "error": true, "cancelled": true, "canceled": true}
|
||||
|
||||
const maxSubmitAttempts = 5
|
||||
|
||||
type Generation struct {
|
||||
DB *gorm.DB
|
||||
Queue *asynq.Client
|
||||
COS *storage.COS
|
||||
Encryptor *security.Encryptor
|
||||
Provider apimart.Client
|
||||
Seedance2 apimart.Seedance2
|
||||
DownloadHTTP *http.Client
|
||||
PollInterval time.Duration
|
||||
WorkerID string
|
||||
}
|
||||
|
||||
type channelSnapshot struct {
|
||||
ID uuid.UUID
|
||||
BaseURL string
|
||||
APIKeyCiphertext string
|
||||
MaxConcurrency int
|
||||
MaxUserConcurrency int
|
||||
Enabled bool
|
||||
}
|
||||
|
||||
type submitSnapshot struct {
|
||||
Task model.GenerationTask
|
||||
Channel channelSnapshot
|
||||
Model string
|
||||
}
|
||||
|
||||
func NewGeneration(db *gorm.DB, queue *asynq.Client, cos *storage.COS, encryptor *security.Encryptor, client *http.Client, pollInterval time.Duration) *Generation {
|
||||
return &Generation{DB: db, Queue: queue, COS: cos, Encryptor: encryptor, Provider: apimart.NewClient(client), Seedance2: apimart.NewSeedance2(client), DownloadHTTP: client, PollInterval: pollInterval, WorkerID: uuid.NewString()}
|
||||
}
|
||||
|
||||
func (w *Generation) Register(mux *asynq.ServeMux) {
|
||||
mux.HandleFunc(queuepkg.TypeDispatchChannel, w.dispatchChannel)
|
||||
mux.HandleFunc(queuepkg.TypeSubmitTask, w.submitTask)
|
||||
mux.HandleFunc(queuepkg.TypePollTask, w.pollTask)
|
||||
mux.HandleFunc(queuepkg.TypeDownloadTask, w.downloadTask)
|
||||
}
|
||||
|
||||
func decodeID(task *asynq.Task) (uuid.UUID, error) {
|
||||
var payload queuepkg.IDPayload
|
||||
if err := json.Unmarshal(task.Payload(), &payload); err != nil {
|
||||
return uuid.Nil, err
|
||||
}
|
||||
if payload.ID == uuid.Nil {
|
||||
return uuid.Nil, errors.New("任务标识无效")
|
||||
}
|
||||
return payload.ID, nil
|
||||
}
|
||||
|
||||
func (w *Generation) dispatchChannel(ctx context.Context, raw *asynq.Task) error {
|
||||
channelID, err := decodeID(raw)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
selected := make([]uuid.UUID, 0)
|
||||
err = w.DB.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
var channel channelSnapshot
|
||||
if err := tx.Raw(`SELECT id,base_url,api_key_ciphertext,max_concurrency,max_user_concurrency,enabled
|
||||
FROM channels WHERE id=? AND deleted_at IS NULL FOR UPDATE`, channelID).Scan(&channel).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if channel.ID == uuid.Nil || !channel.Enabled {
|
||||
return nil
|
||||
}
|
||||
var active int64
|
||||
if err := tx.Model(&model.GenerationTask{}).Where("channel_id=? AND slot_reserved=true AND slot_released_at IS NULL", channelID).Count(&active).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
available := channel.MaxConcurrency - int(active)
|
||||
if available <= 0 {
|
||||
return nil
|
||||
}
|
||||
activeByUser := map[uuid.UUID]int{}
|
||||
rows := make([]struct {
|
||||
UserID uuid.UUID
|
||||
Count int
|
||||
}, 0)
|
||||
if err := tx.Model(&model.GenerationTask{}).Select("user_id,count(*) AS count").Where("channel_id=? AND slot_reserved=true AND slot_released_at IS NULL", channelID).Group("user_id").Scan(&rows).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
for _, row := range rows {
|
||||
activeByUser[row.UserID] = row.Count
|
||||
}
|
||||
var candidates []model.GenerationTask
|
||||
now := time.Now()
|
||||
if err := tx.Raw(`WITH active_users AS (
|
||||
SELECT user_id,count(*) AS active_count FROM generation_tasks
|
||||
WHERE channel_id=? AND slot_reserved=true AND slot_released_at IS NULL GROUP BY user_id
|
||||
), ranked AS (
|
||||
SELECT pending.id,row_number() OVER (PARTITION BY pending.user_id ORDER BY pending.created_at,pending.id) AS user_rank,
|
||||
(?-coalesce(active.active_count,0)) AS allowance
|
||||
FROM generation_tasks pending LEFT JOIN active_users active ON active.user_id=pending.user_id
|
||||
WHERE pending.channel_id=? AND pending.status='pending_submission' AND
|
||||
(pending.project_id IS NOT NULL OR (pending.task_type='image_generation' AND pending.input_data->>'product_image'='true'))
|
||||
AND (pending.next_submit_at IS NULL OR pending.next_submit_at<=?)
|
||||
AND (pending.lease_until IS NULL OR pending.lease_until<?)
|
||||
)
|
||||
SELECT task.* FROM generation_tasks task JOIN ranked ON ranked.id=task.id
|
||||
WHERE ranked.allowance>0 AND ranked.user_rank<=ranked.allowance
|
||||
ORDER BY ranked.user_rank,task.created_at,task.id LIMIT ? FOR UPDATE OF task SKIP LOCKED`,
|
||||
channelID, channel.MaxUserConcurrency, channelID, now, now, available).Scan(&candidates).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if len(candidates) == 0 {
|
||||
return nil
|
||||
}
|
||||
for _, candidate := range selectFairCandidates(candidates, activeByUser, available, channel.MaxUserConcurrency) {
|
||||
leaseUntil := now.Add(2 * time.Minute)
|
||||
result := tx.Model(&model.GenerationTask{}).Where("id=? AND status='pending_submission'", candidate.ID).Updates(map[string]any{
|
||||
"status": "submitting", "slot_reserved": true, "slot_reserved_at": now,
|
||||
"slot_released_at": nil, "lease_owner": w.WorkerID, "lease_until": leaseUntil,
|
||||
"submit_attempts": gorm.Expr("submit_attempts+1"),
|
||||
})
|
||||
if result.Error != nil {
|
||||
return result.Error
|
||||
}
|
||||
if result.RowsAffected == 1 {
|
||||
selected = append(selected, candidate.ID)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
for _, taskID := range selected {
|
||||
if err := queuepkg.EnqueueID(w.Queue, queuepkg.TypeSubmitTask, taskID, 0); err != nil {
|
||||
slog.Error("提交任务入队失败", "task_id", taskID, "error", err)
|
||||
w.returnUnenqueuedTask(ctx, taskID, err.Error())
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func selectFairCandidates(candidates []model.GenerationTask, activeByUser map[uuid.UUID]int, available, maxPerUser int) []model.GenerationTask {
|
||||
if available <= 0 || maxPerUser <= 0 {
|
||||
return nil
|
||||
}
|
||||
groups := map[uuid.UUID][]model.GenerationTask{}
|
||||
order := make([]uuid.UUID, 0)
|
||||
for _, candidate := range candidates {
|
||||
if _, exists := groups[candidate.UserID]; !exists {
|
||||
order = append(order, candidate.UserID)
|
||||
}
|
||||
groups[candidate.UserID] = append(groups[candidate.UserID], candidate)
|
||||
}
|
||||
selected := make([]model.GenerationTask, 0, available)
|
||||
for available > 0 {
|
||||
progressed := false
|
||||
for _, userID := range order {
|
||||
if available == 0 {
|
||||
break
|
||||
}
|
||||
if activeByUser[userID] >= maxPerUser || len(groups[userID]) == 0 {
|
||||
continue
|
||||
}
|
||||
selected = append(selected, groups[userID][0])
|
||||
groups[userID] = groups[userID][1:]
|
||||
activeByUser[userID]++
|
||||
available--
|
||||
progressed = true
|
||||
}
|
||||
if !progressed {
|
||||
break
|
||||
}
|
||||
}
|
||||
return selected
|
||||
}
|
||||
|
||||
func (w *Generation) loadSubmitSnapshot(ctx context.Context, taskID uuid.UUID) (submitSnapshot, error) {
|
||||
var result submitSnapshot
|
||||
if err := w.DB.WithContext(ctx).Where("id=?", taskID).First(&result.Task).Error; err != nil {
|
||||
return result, err
|
||||
}
|
||||
if result.Task.ChannelID == nil || result.Task.ModelID == nil {
|
||||
return result, errors.New("任务未绑定模型渠道")
|
||||
}
|
||||
if err := w.DB.WithContext(ctx).Raw(`SELECT id,base_url,api_key_ciphertext,max_concurrency,max_user_concurrency,enabled
|
||||
FROM channels WHERE id=? AND deleted_at IS NULL`, *result.Task.ChannelID).Scan(&result.Channel).Error; err != nil {
|
||||
return result, err
|
||||
}
|
||||
if !result.Channel.Enabled {
|
||||
return result, errors.New("任务渠道已停用")
|
||||
}
|
||||
if err := w.DB.WithContext(ctx).Table("models").Select("name").Where("id=? AND enabled=true AND deleted_at IS NULL", *result.Task.ModelID).Scan(&result.Model).Error; err != nil {
|
||||
return result, err
|
||||
}
|
||||
if result.Model == "" {
|
||||
return result, errors.New("任务模型不可用")
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func (w *Generation) submitTask(ctx context.Context, raw *asynq.Task) error {
|
||||
taskID, err := decodeID(raw)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
snapshot, err := w.loadSubmitSnapshot(ctx, taskID)
|
||||
if err != nil {
|
||||
// 用户可以删除卡死的活动任务;队列收到遗留消息时直接结束,避免无意义重试。
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
return nil
|
||||
}
|
||||
return w.failAndRefund(ctx, taskID, "submit_configuration_error", err.Error())
|
||||
}
|
||||
if snapshot.Task.Status != "submitting" {
|
||||
if snapshot.Task.Status == "pending_submission" && snapshot.Task.ChannelID != nil {
|
||||
return queuepkg.EnqueueID(w.Queue, queuepkg.TypeDispatchChannel, *snapshot.Task.ChannelID, 0)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
apiKey, err := w.Encryptor.Decrypt(snapshot.Channel.APIKeyCiphertext)
|
||||
if err != nil {
|
||||
return w.failAndRefund(ctx, taskID, "channel_secret_error", "渠道密钥不可用")
|
||||
}
|
||||
payload, err := buildProviderPayload(snapshot)
|
||||
if err != nil {
|
||||
return w.failAndRefund(ctx, taskID, "invalid_payload", err.Error())
|
||||
}
|
||||
requestCtx, cancel := context.WithTimeout(ctx, 90*time.Second)
|
||||
defer cancel()
|
||||
var result provider.SubmitResult
|
||||
if snapshot.Task.TaskType == "video_generation" {
|
||||
if payloadJSON, marshalErr := json.Marshal(payload); marshalErr == nil {
|
||||
fmt.Printf("console.log video_generation payload: %s\n", payloadJSON)
|
||||
}
|
||||
result, err = w.Seedance2.Submit(requestCtx, snapshot.Channel.BaseURL, apiKey, snapshot.Task.RequestID, payload)
|
||||
} else {
|
||||
result, err = w.Provider.SubmitImage(requestCtx, snapshot.Channel.BaseURL, apiKey, snapshot.Task.RequestID, payload)
|
||||
}
|
||||
if err != nil {
|
||||
var httpErr *provider.HTTPError
|
||||
if errors.As(err, &httpErr) {
|
||||
if httpErr.StatusCode == http.StatusTooManyRequests || httpErr.StatusCode >= 500 {
|
||||
if !canRetrySubmission(snapshot.Task.SubmitAttempts) {
|
||||
return w.failAndRefund(ctx, taskID, "upstream_unavailable", err.Error())
|
||||
}
|
||||
return w.requeueRejected(ctx, snapshot.Task, err.Error(), retryDelay(snapshot.Task.SubmitAttempts, httpErr.RetryAfter))
|
||||
}
|
||||
return w.failAndRefund(ctx, taskID, "upstream_rejected", err.Error())
|
||||
}
|
||||
// 传输异常统一按提交失败重试;达到上限后结束任务并返还预扣积分。
|
||||
if !canRetrySubmission(snapshot.Task.SubmitAttempts) {
|
||||
return w.failAndRefund(ctx, taskID, "upstream_unavailable", err.Error())
|
||||
}
|
||||
return w.requeueRejected(ctx, snapshot.Task, err.Error(), retryDelay(snapshot.Task.SubmitAttempts, ""))
|
||||
}
|
||||
now := time.Now()
|
||||
if result.URL != "" {
|
||||
if err := w.markResultReady(ctx, taskID, result.URL); err != nil {
|
||||
return err
|
||||
}
|
||||
_ = queuepkg.EnqueueID(w.Queue, queuepkg.TypeDownloadTask, taskID, 0)
|
||||
w.enqueueDispatch(snapshot.Task.ChannelID)
|
||||
return nil
|
||||
}
|
||||
nextPoll := now.Add(w.PollInterval + jitter(3*time.Second))
|
||||
update := w.DB.WithContext(ctx).Model(&model.GenerationTask{}).Where("id=? AND status='submitting'", taskID).Updates(map[string]any{
|
||||
"status": "submitted", "upstream_task_id": result.TaskID, "submitted_at": now,
|
||||
"next_poll_at": nextPoll, "lease_owner": nil, "lease_until": nil, "error_code": nil, "error_message": nil,
|
||||
})
|
||||
if update.Error != nil {
|
||||
return update.Error
|
||||
}
|
||||
if update.RowsAffected == 1 {
|
||||
_ = queuepkg.EnqueueID(w.Queue, queuepkg.TypePollTask, taskID, time.Until(nextPoll))
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func buildProviderPayload(snapshot submitSnapshot) (map[string]any, error) {
|
||||
if snapshot.Task.TaskType == "video_generation" {
|
||||
return apimart.BuildSeedance2Payload(snapshot.Model, snapshot.Task.InputData)
|
||||
}
|
||||
var input map[string]any
|
||||
if err := json.Unmarshal(snapshot.Task.InputData, &input); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
payload := map[string]any{"model": snapshot.Model, "prompt": strings.TrimSpace(fmt.Sprint(input["prompt"]))}
|
||||
if payload["prompt"] == "" {
|
||||
return nil, errors.New("生成提示词不能为空")
|
||||
}
|
||||
if snapshot.Task.TaskType == "image_generation" {
|
||||
// 图片画幅由各业务模块在入队前规范化,Worker 只负责透传统一字段。
|
||||
payload["size"] = stringOr(input["size"], "16:9")
|
||||
payload["resolution"] = stringOr(input["resolution"], "1k")
|
||||
} else {
|
||||
payload["size"] = stringOr(input["resolution"], "1k")
|
||||
}
|
||||
if n, ok := input["n"]; ok {
|
||||
payload["n"] = n
|
||||
} else {
|
||||
payload["n"] = 1
|
||||
}
|
||||
if urls, ok := input["image_urls"]; ok {
|
||||
payload["image_urls"] = urls
|
||||
}
|
||||
return payload, nil
|
||||
}
|
||||
|
||||
func (w *Generation) pollTask(ctx context.Context, raw *asynq.Task) error {
|
||||
taskID, err := decodeID(raw)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
snapshot, err := w.loadSubmitSnapshot(ctx, taskID)
|
||||
if err != nil || (snapshot.Task.Status != "submitted" && snapshot.Task.Status != "processing" && snapshot.Task.Status != "cancel_requested") {
|
||||
return nil
|
||||
}
|
||||
apiKey, err := w.Encryptor.Decrypt(snapshot.Channel.APIKeyCiphertext)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
requestCtx, cancel := context.WithTimeout(ctx, 30*time.Second)
|
||||
var result provider.PollResult
|
||||
var pollErr error
|
||||
if snapshot.Task.TaskType == "video_generation" {
|
||||
result, pollErr = w.Seedance2.Poll(requestCtx, snapshot.Channel.BaseURL, apiKey, snapshot.Task.UpstreamTaskID)
|
||||
} else {
|
||||
result, pollErr = w.Provider.PollImage(requestCtx, snapshot.Channel.BaseURL, apiKey, snapshot.Task.UpstreamTaskID)
|
||||
}
|
||||
cancel()
|
||||
if pollErr != nil {
|
||||
if snapshot.Task.Status == "cancel_requested" && snapshot.Task.PollAttempts >= 10 {
|
||||
if err := w.finishCancellation(ctx, taskID, true, "上游任务状态长时间不可用,已自动结束取消"); err != nil {
|
||||
return err
|
||||
}
|
||||
w.enqueueDispatch(snapshot.Task.ChannelID)
|
||||
return nil
|
||||
}
|
||||
next := time.Now().Add(w.PollInterval + jitter(4*time.Second))
|
||||
w.DB.WithContext(ctx).Model(&model.GenerationTask{}).Where("id=? AND status IN ('submitted','processing','cancel_requested')", taskID).Updates(map[string]any{
|
||||
"poll_attempts": gorm.Expr("poll_attempts+1"), "next_poll_at": next, "error_message": pollErr.Error(),
|
||||
})
|
||||
_ = queuepkg.EnqueueID(w.Queue, queuepkg.TypePollTask, taskID, time.Until(next))
|
||||
return nil
|
||||
}
|
||||
if completedStatuses[result.Status] {
|
||||
if snapshot.Task.Status == "cancel_requested" {
|
||||
if err := w.finishCancellation(ctx, taskID, false, "上游任务已结束"); err != nil {
|
||||
return err
|
||||
}
|
||||
w.enqueueDispatch(snapshot.Task.ChannelID)
|
||||
return nil
|
||||
}
|
||||
if result.URL == "" {
|
||||
next := time.Now().Add(w.PollInterval)
|
||||
w.DB.WithContext(ctx).Model(&model.GenerationTask{}).Where("id=?", taskID).Updates(map[string]any{"next_poll_at": next, "error_message": "中转站已完成但结果地址尚未就绪"})
|
||||
_ = queuepkg.EnqueueID(w.Queue, queuepkg.TypePollTask, taskID, time.Until(next))
|
||||
return nil
|
||||
}
|
||||
if err := w.markResultReady(ctx, taskID, result.URL); err != nil {
|
||||
return err
|
||||
}
|
||||
_ = queuepkg.EnqueueID(w.Queue, queuepkg.TypeDownloadTask, taskID, 0)
|
||||
w.enqueueDispatch(snapshot.Task.ChannelID)
|
||||
return nil
|
||||
}
|
||||
if failedStatuses[result.Status] {
|
||||
if snapshot.Task.Status == "cancel_requested" {
|
||||
if err := w.finishCancellation(ctx, taskID, true, "上游任务已取消"); err != nil {
|
||||
return err
|
||||
}
|
||||
} else {
|
||||
if err := w.failAndRefund(ctx, taskID, "upstream_failed", result.Error); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
w.enqueueDispatch(snapshot.Task.ChannelID)
|
||||
return nil
|
||||
}
|
||||
next := time.Now().Add(w.PollInterval + jitter(4*time.Second))
|
||||
updates := map[string]any{"poll_attempts": gorm.Expr("poll_attempts+1"), "next_poll_at": next, "error_code": nil, "error_message": nil}
|
||||
if snapshot.Task.Status != "cancel_requested" {
|
||||
updates["status"] = "processing"
|
||||
}
|
||||
w.DB.WithContext(ctx).Model(&model.GenerationTask{}).Where("id=? AND status IN ('submitted','processing','cancel_requested')", taskID).Updates(updates)
|
||||
_ = queuepkg.EnqueueID(w.Queue, queuepkg.TypePollTask, taskID, time.Until(next))
|
||||
return nil
|
||||
}
|
||||
|
||||
func (w *Generation) markResultReady(ctx context.Context, taskID uuid.UUID, resultURL string) error {
|
||||
return w.DB.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
var task model.GenerationTask
|
||||
if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Where("id=?", taskID).First(&task).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if task.Status == "cancelled" || task.Status == "cancel_requested" {
|
||||
return nil
|
||||
}
|
||||
updates := map[string]any{"status": "result_ready", "upstream_result_url": resultURL, "next_poll_at": nil, "lease_owner": nil, "lease_until": nil}
|
||||
if task.SlotReserved && task.SlotReleasedAt == nil {
|
||||
now := time.Now()
|
||||
updates["slot_released_at"] = now
|
||||
}
|
||||
return tx.Model(&task).Updates(updates).Error
|
||||
})
|
||||
}
|
||||
|
||||
func (w *Generation) downloadTask(ctx context.Context, raw *asynq.Task) error {
|
||||
taskID, err := decodeID(raw)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
var row struct {
|
||||
model.GenerationTask
|
||||
SequenceNo int
|
||||
EpisodeNo int
|
||||
}
|
||||
if err := w.DB.WithContext(ctx).Table("generation_tasks gt").Select("gt.*,coalesce(sb.sequence_no,0) AS sequence_no,coalesce(e.episode_no,0) AS episode_no").Joins("LEFT JOIN episode_storyboards sb ON sb.id=gt.storyboard_id").Joins("LEFT JOIN project_episodes e ON e.id=gt.episode_id").Where("gt.id=?", taskID).Scan(&row).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if row.Status != "result_ready" && row.Status != "downloading" {
|
||||
return nil
|
||||
}
|
||||
if row.UpstreamResultURL == "" || w.COS == nil {
|
||||
return w.failWithoutRefund(ctx, taskID, "result_storage_error", "结果存储信息不完整")
|
||||
}
|
||||
w.DB.WithContext(ctx).Model(&model.GenerationTask{}).Where("id=? AND status='result_ready'", taskID).Update("status", "downloading")
|
||||
mediaID := uuid.New()
|
||||
outputType, fallbackType, maxBytes := "video", "video/mp4", w.COS.MaxVideoBytes()
|
||||
displayName := fmt.Sprintf("分镜%04d生成结果", row.SequenceNo)
|
||||
key := ""
|
||||
if row.TaskType == "image_generation" {
|
||||
outputType, fallbackType, maxBytes = "image", "image/png", w.COS.MaxImageBytes()
|
||||
var input map[string]any
|
||||
_ = json.Unmarshal(row.InputData, &input)
|
||||
if row.ProjectID == nil {
|
||||
key = mediakey.StandaloneTaskOutput(row.UserID, taskID, "image", fallbackType)
|
||||
displayName = "独立生成图"
|
||||
} else {
|
||||
if fmt.Sprint(input["target_type"]) == "storyboard" && row.StoryboardID != nil {
|
||||
if row.EpisodeID != nil {
|
||||
key = mediakey.TaskOutput(*row.ProjectID, *row.EpisodeID, row.SequenceNo, taskID, mediaID, "image", fallbackType)
|
||||
} else {
|
||||
key = mediakey.ProjectTaskOutput(*row.ProjectID, row.SequenceNo, taskID, mediaID, "image", fallbackType)
|
||||
}
|
||||
displayName = fmt.Sprintf("分镜%04d生成图", row.SequenceNo)
|
||||
} else {
|
||||
assetID, parseErr := uuid.Parse(fmt.Sprint(input["asset_id"]))
|
||||
if parseErr != nil {
|
||||
return w.failWithoutRefund(ctx, taskID, "result_storage_error", "资产生图任务缺少资产标识")
|
||||
}
|
||||
var asset struct{ AssetType, Name string }
|
||||
if err := w.DB.WithContext(ctx).Table("project_assets").Select("asset_type,name").Where("id=? AND project_id=? AND deleted_at IS NULL", assetID, *row.ProjectID).Take(&asset).Error; err != nil {
|
||||
return w.failWithoutRefund(ctx, taskID, "result_storage_error", "资产生图任务绑定的资产不存在")
|
||||
}
|
||||
key = mediakey.AssetOutput(*row.ProjectID, asset.AssetType, assetID, taskID, mediaID, fallbackType)
|
||||
displayName = asset.Name + "生成图"
|
||||
}
|
||||
}
|
||||
} else {
|
||||
if row.ProjectID != nil {
|
||||
if row.EpisodeID != nil {
|
||||
key = mediakey.TaskOutput(*row.ProjectID, *row.EpisodeID, row.SequenceNo, taskID, mediaID, outputType, fallbackType)
|
||||
} else {
|
||||
key = mediakey.ProjectTaskOutput(*row.ProjectID, row.SequenceNo, taskID, mediaID, outputType, fallbackType)
|
||||
}
|
||||
} else {
|
||||
key = mediakey.StandaloneTaskOutput(row.UserID, taskID, "video", fallbackType)
|
||||
}
|
||||
}
|
||||
requestCtx, cancel := context.WithTimeout(ctx, 5*time.Minute)
|
||||
url, storedKey, contentType, size, downloadErr := w.COS.PutFromURL(requestCtx, w.DownloadHTTP, row.UpstreamResultURL, key, fallbackType, maxBytes)
|
||||
cancel()
|
||||
if downloadErr != nil {
|
||||
attempts := row.DownloadAttempts + 1
|
||||
if attempts < 3 {
|
||||
w.DB.WithContext(ctx).Model(&model.GenerationTask{}).Where("id=?", taskID).Updates(map[string]any{"status": "result_ready", "download_attempts": attempts, "error_message": downloadErr.Error()})
|
||||
_ = queuepkg.EnqueueID(w.Queue, queuepkg.TypeDownloadTask, taskID, time.Duration(attempts*10)*time.Second)
|
||||
return nil
|
||||
}
|
||||
return w.failWithoutRefund(ctx, taskID, "result_download_failed", downloadErr.Error())
|
||||
}
|
||||
key = storedKey
|
||||
if (outputType == "image" && !strings.HasPrefix(contentType, "image/")) || (outputType == "video" && !strings.HasPrefix(contentType, "video/")) {
|
||||
_ = w.COS.Delete(ctx, key)
|
||||
return w.failWithoutRefund(ctx, taskID, "invalid_result_media", "中转站结果媒体类型与任务不匹配")
|
||||
}
|
||||
ownerID := row.UserID
|
||||
mediaAsset := model.MediaAsset{ID: mediaID, OwnerUserID: &ownerID, StorageProvider: "cos", ObjectKey: key, PublicURL: url, DisplayName: displayName, MimeType: contentType, SizeBytes: size}
|
||||
metadata, _ := json.Marshal(map[string]any{"source": "apimart", "stored_at": time.Now()})
|
||||
output := model.GenerationOutput{TaskID: taskID, MediaAssetID: mediaID, OutputType: outputType, SequenceNo: 1, Metadata: metadata}
|
||||
reusedOutput := false
|
||||
cancelledDuringDownload := false
|
||||
err = w.DB.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
var lockedTask model.GenerationTask
|
||||
if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Select("id", "status").Where("id=?", taskID).Take(&lockedTask).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if lockedTask.Status == "cancelled" || lockedTask.Status == "cancel_requested" {
|
||||
cancelledDuringDownload = true
|
||||
return nil
|
||||
}
|
||||
var existing model.GenerationOutput
|
||||
existingErr := tx.Where("task_id=? AND sequence_no=?", taskID, 1).Take(&existing).Error
|
||||
if existingErr == nil {
|
||||
output = existing
|
||||
reusedOutput = true
|
||||
} else if errors.Is(existingErr, gorm.ErrRecordNotFound) {
|
||||
if err := tx.Create(&mediaAsset).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if err := tx.Create(&output).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
} else {
|
||||
return existingErr
|
||||
}
|
||||
now := time.Now()
|
||||
cost := row.PrepaidPoints
|
||||
if err := tx.Model(&model.GenerationTask{}).Where("id=?", taskID).Updates(map[string]any{"status": "succeeded", "actual_points": cost, "finished_at": now, "error_code": nil, "error_message": nil}).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if row.TaskType == "video_generation" && row.StoryboardID != nil {
|
||||
premiumDrama, premiumErr := isPremiumDramaGeneration(tx, row.GenerationTask)
|
||||
if premiumErr != nil {
|
||||
return premiumErr
|
||||
}
|
||||
if premiumDrama {
|
||||
if err := tx.Exec(`UPDATE generation_outputs SET metadata=jsonb_set(coalesce(metadata,'{}'::jsonb),'{candidate}','true'::jsonb,true)
|
||||
WHERE id=(SELECT active_output_id FROM episode_storyboards WHERE id=?)`, *row.StoryboardID).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return tx.Model(&model.EpisodeStoryboard{}).Where("id=?", *row.StoryboardID).Updates(map[string]any{"active_output_id": output.ID, "status": "completed"}).Error
|
||||
}
|
||||
if row.TaskType == "image_generation" && row.ProjectID != nil {
|
||||
var input map[string]any
|
||||
_ = json.Unmarshal(row.InputData, &input)
|
||||
if fmt.Sprint(input["target_type"]) == "storyboard" && row.StoryboardID != nil {
|
||||
return tx.Model(&model.EpisodeStoryboard{}).Where("id=?", *row.StoryboardID).Updates(map[string]any{"thumbnail_asset_id": output.MediaAssetID, "status": "completed"}).Error
|
||||
}
|
||||
if assetID, parseErr := uuid.Parse(fmt.Sprint(input["asset_id"])); parseErr == nil {
|
||||
return tx.Model(&model.ProjectAsset{}).Where("id=?", assetID).Update("image_asset_id", output.MediaAssetID).Error
|
||||
}
|
||||
}
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
_ = w.COS.Delete(ctx, key)
|
||||
return err
|
||||
}
|
||||
if cancelledDuringDownload {
|
||||
_ = w.COS.Delete(ctx, key)
|
||||
return w.finishCancellation(ctx, taskID, false, "用户取消生成")
|
||||
}
|
||||
if reusedOutput {
|
||||
_ = w.COS.Delete(ctx, key)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (w *Generation) requeueRejected(ctx context.Context, task model.GenerationTask, message string, delay time.Duration) error {
|
||||
if err := w.DB.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
var current model.GenerationTask
|
||||
if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Where("id=?", task.ID).First(¤t).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
next := time.Now().Add(delay)
|
||||
updates := map[string]any{"status": "pending_submission", "next_submit_at": next, "lease_owner": nil, "lease_until": nil, "error_message": message}
|
||||
if current.SlotReserved && current.SlotReleasedAt == nil {
|
||||
now := time.Now()
|
||||
updates["slot_released_at"] = now
|
||||
}
|
||||
return tx.Model(¤t).Updates(updates).Error
|
||||
}); err != nil {
|
||||
return err
|
||||
}
|
||||
if task.ChannelID != nil {
|
||||
return queuepkg.EnqueueID(w.Queue, queuepkg.TypeDispatchChannel, *task.ChannelID, delay)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (w *Generation) returnUnenqueuedTask(ctx context.Context, taskID uuid.UUID, message string) {
|
||||
var channelID *uuid.UUID
|
||||
_ = w.DB.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
var task model.GenerationTask
|
||||
if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Where("id=?", taskID).First(&task).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if task.Status != "submitting" {
|
||||
return nil
|
||||
}
|
||||
channelID = task.ChannelID
|
||||
now := time.Now()
|
||||
return tx.Model(&task).Updates(map[string]any{
|
||||
"status": "pending_submission", "next_submit_at": now.Add(5 * time.Second),
|
||||
"slot_released_at": now, "lease_owner": nil, "lease_until": nil,
|
||||
"error_message": truncate(message, 1000),
|
||||
}).Error
|
||||
})
|
||||
w.enqueueDispatchAfter(channelID, 5*time.Second)
|
||||
}
|
||||
|
||||
func (w *Generation) finishCancellation(ctx context.Context, taskID uuid.UUID, refund bool, message string) error {
|
||||
return w.DB.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
var task model.GenerationTask
|
||||
if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Where("id=?", taskID).First(&task).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if task.Status == "cancelled" {
|
||||
return nil
|
||||
}
|
||||
now := time.Now()
|
||||
updates := map[string]any{
|
||||
"status": "cancelled", "finished_at": now, "next_poll_at": nil,
|
||||
"lease_owner": nil, "lease_until": nil, "error_code": nil, "error_message": message,
|
||||
}
|
||||
if task.SlotReserved && task.SlotReleasedAt == nil {
|
||||
updates["slot_released_at"] = now
|
||||
}
|
||||
premiumDrama, err := isPremiumDramaGeneration(tx, task)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
// Premium-drama generation keeps its prepaid charge on cancellation.
|
||||
refund = refund && !premiumDrama
|
||||
if refund {
|
||||
refunded, err := billing.RefundGenerationTask(tx, &task, "视频生成失败返还")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if refunded {
|
||||
updates["cost_refunded"] = true
|
||||
}
|
||||
updates["actual_points"] = "0.00"
|
||||
} else {
|
||||
updates["actual_points"] = task.PrepaidPoints
|
||||
}
|
||||
if err := tx.Model(&task).Updates(updates).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if task.TaskType == "video_generation" && task.StoryboardID != nil {
|
||||
return tx.Model(&model.EpisodeStoryboard{}).Where("id=?", *task.StoryboardID).Update("status", "idle").Error
|
||||
}
|
||||
return nil
|
||||
})
|
||||
}
|
||||
|
||||
func (w *Generation) failAndRefund(ctx context.Context, taskID uuid.UUID, code, message string) error {
|
||||
return w.finishFailure(ctx, taskID, code, message, true)
|
||||
}
|
||||
|
||||
func (w *Generation) failWithoutRefund(ctx context.Context, taskID uuid.UUID, code, message string) error {
|
||||
return w.finishFailure(ctx, taskID, code, message, false)
|
||||
}
|
||||
|
||||
func (w *Generation) finishFailure(ctx context.Context, taskID uuid.UUID, code, message string, refund bool) error {
|
||||
return w.DB.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
var task model.GenerationTask
|
||||
if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Where("id=?", taskID).First(&task).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if task.Status == "succeeded" || task.Status == "failed" || task.Status == "cancelled" {
|
||||
return nil
|
||||
}
|
||||
now := time.Now()
|
||||
updates := map[string]any{"status": "failed", "error_code": code, "error_message": truncate(message, 4000), "finished_at": now, "lease_owner": nil, "lease_until": nil}
|
||||
if task.SlotReserved && task.SlotReleasedAt == nil {
|
||||
updates["slot_released_at"] = now
|
||||
}
|
||||
premiumDrama, err := isPremiumDramaGeneration(tx, task)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if premiumDrama && code != "upstream_rejected" && code != "upstream_unavailable" && code != "upstream_failed" {
|
||||
refund = false
|
||||
}
|
||||
if refund {
|
||||
refunded, err := billing.RefundGenerationTask(tx, &task, "视频生成失败返还")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if refunded {
|
||||
updates["cost_refunded"] = true
|
||||
}
|
||||
updates["actual_points"] = "0.00"
|
||||
} else {
|
||||
updates["actual_points"] = task.PrepaidPoints
|
||||
}
|
||||
if err := tx.Model(&task).Updates(updates).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if task.TaskType == "video_generation" && task.StoryboardID != nil {
|
||||
return tx.Model(&model.EpisodeStoryboard{}).Where("id=?", *task.StoryboardID).Update("status", "failed").Error
|
||||
}
|
||||
return nil
|
||||
})
|
||||
}
|
||||
|
||||
func isPremiumDramaGeneration(tx *gorm.DB, task model.GenerationTask) (bool, error) {
|
||||
if task.ProjectID == nil {
|
||||
return false, nil
|
||||
}
|
||||
var projectType string
|
||||
if err := tx.Table("creative_projects").Where("id=?", *task.ProjectID).Pluck("project_type", &projectType).Error; err != nil {
|
||||
return false, err
|
||||
}
|
||||
return projectType == "premium_drama", nil
|
||||
}
|
||||
|
||||
func (w *Generation) enqueueDispatch(channelID *uuid.UUID) {
|
||||
w.enqueueDispatchAfter(channelID, 0)
|
||||
}
|
||||
|
||||
func (w *Generation) enqueueDispatchAfter(channelID *uuid.UUID, delay time.Duration) {
|
||||
if channelID != nil {
|
||||
_ = queuepkg.EnqueueID(w.Queue, queuepkg.TypeDispatchChannel, *channelID, delay)
|
||||
}
|
||||
}
|
||||
|
||||
// Recover 恢复服务中断或旧版本遗留的生成任务,并重新投递待处理工作。
|
||||
func (w *Generation) Recover(ctx context.Context) error {
|
||||
now := time.Now()
|
||||
if err := w.DB.WithContext(ctx).Model(&model.GenerationTask{}).
|
||||
Where("task_type IN ('image_generation','video_generation') AND status='submitting' AND lease_until<?", now).
|
||||
Updates(map[string]any{
|
||||
"status": "pending_submission", "next_submit_at": now, "slot_released_at": now,
|
||||
"error_code": nil, "error_message": "提交中断,正在自动重试", "lease_owner": nil, "lease_until": nil,
|
||||
}).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
channels := make([]struct {
|
||||
ChannelID uuid.UUID
|
||||
NextSubmitAt time.Time
|
||||
}, 0)
|
||||
if err := w.DB.WithContext(ctx).Model(&model.GenerationTask{}).
|
||||
Select("channel_id,min(coalesce(next_submit_at,?)) AS next_submit_at", now).
|
||||
Where("status='pending_submission' AND channel_id IS NOT NULL AND (project_id IS NOT NULL OR (task_type='image_generation' AND input_data->>'product_image'='true'))").Group("channel_id").Scan(&channels).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
for _, item := range channels {
|
||||
delay := time.Until(item.NextSubmitAt)
|
||||
if delay < 0 {
|
||||
delay = 0
|
||||
}
|
||||
_ = queuepkg.EnqueueID(w.Queue, queuepkg.TypeDispatchChannel, item.ChannelID, delay)
|
||||
}
|
||||
var polls []uuid.UUID
|
||||
_ = w.DB.WithContext(ctx).Model(&model.GenerationTask{}).Where("(project_id IS NOT NULL OR (task_type='image_generation' AND input_data->>'product_image'='true')) AND (status IN ('submitted','processing') OR (status='cancel_requested' AND upstream_task_id IS NOT NULL AND upstream_task_id<>''))").Pluck("id", &polls).Error
|
||||
for _, id := range polls {
|
||||
_ = queuepkg.EnqueueID(w.Queue, queuepkg.TypePollTask, id, jitter(5*time.Second))
|
||||
}
|
||||
var downloads []uuid.UUID
|
||||
_ = w.DB.WithContext(ctx).Model(&model.GenerationTask{}).Where("(project_id IS NOT NULL OR (task_type='image_generation' AND input_data->>'product_image'='true')) AND status IN ('result_ready','downloading')").Pluck("id", &downloads).Error
|
||||
for _, id := range downloads {
|
||||
_ = queuepkg.EnqueueID(w.Queue, queuepkg.TypeDownloadTask, id, jitter(5*time.Second))
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func stringOr(value any, fallback string) string {
|
||||
text := strings.TrimSpace(fmt.Sprint(value))
|
||||
if text == "" || text == "<nil>" {
|
||||
return fallback
|
||||
}
|
||||
return text
|
||||
}
|
||||
|
||||
func intOr(value any, fallback int) int {
|
||||
parsed, err := strconv.Atoi(fmt.Sprint(value))
|
||||
if err != nil {
|
||||
return fallback
|
||||
}
|
||||
return parsed
|
||||
}
|
||||
|
||||
func retryDelay(attempt int, retryAfter string) time.Duration {
|
||||
if seconds, err := strconv.Atoi(retryAfter); err == nil && seconds > 0 {
|
||||
return time.Duration(seconds)*time.Second + jitter(3*time.Second)
|
||||
}
|
||||
schedule := []time.Duration{15 * time.Second, 30 * time.Second, time.Minute, 2 * time.Minute, 5 * time.Minute}
|
||||
if attempt < 1 {
|
||||
attempt = 1
|
||||
}
|
||||
if attempt > len(schedule) {
|
||||
attempt = len(schedule)
|
||||
}
|
||||
return schedule[attempt-1] + jitter(5*time.Second)
|
||||
}
|
||||
|
||||
// canRetrySubmission 判断生成任务是否仍在允许的自动提交次数内。
|
||||
func canRetrySubmission(attempt int) bool {
|
||||
return attempt < maxSubmitAttempts
|
||||
}
|
||||
|
||||
func jitter(max time.Duration) time.Duration {
|
||||
if max <= 0 {
|
||||
return 0
|
||||
}
|
||||
return time.Duration(rand.Int63n(int64(max)))
|
||||
}
|
||||
|
||||
func truncate(value string, max int) string {
|
||||
if len(value) > max {
|
||||
return value[:max]
|
||||
}
|
||||
return value
|
||||
}
|
||||
|
||||
func NewHTTPClient(maxConnections int) *http.Client {
|
||||
if maxConnections < 10 {
|
||||
maxConnections = 10
|
||||
}
|
||||
transport := &http.Transport{
|
||||
Proxy: http.ProxyFromEnvironment, DialContext: (&net.Dialer{Timeout: 15 * time.Second, KeepAlive: 30 * time.Second}).DialContext,
|
||||
MaxIdleConns: maxConnections, MaxIdleConnsPerHost: maxConnections, MaxConnsPerHost: maxConnections,
|
||||
// ResponseHeaderTimeout 设为 0 表示不限制,由各调用点的 context 超时统一控制,避免非流式 AI 推理在 60 秒内未返回响应头被误判超时
|
||||
IdleConnTimeout: 90 * time.Second, TLSHandshakeTimeout: 10 * time.Second, ResponseHeaderTimeout: 0,
|
||||
}
|
||||
return &http.Client{Transport: transport}
|
||||
}
|
||||
@@ -0,0 +1,63 @@
|
||||
package worker
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"juhe-factory/api/internal/model"
|
||||
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
func TestSelectFairCandidatesHonorsTotalCapacity(t *testing.T) {
|
||||
users := make([]uuid.UUID, 50)
|
||||
for index := range users {
|
||||
users[index] = uuid.New()
|
||||
}
|
||||
candidates := make([]model.GenerationTask, 0, 501)
|
||||
for index := 0; index < 501; index++ {
|
||||
candidates = append(candidates, model.GenerationTask{ID: uuid.New(), UserID: users[index%len(users)]})
|
||||
}
|
||||
selected := selectFairCandidates(candidates, map[uuid.UUID]int{}, 500, 10)
|
||||
if len(selected) != 500 {
|
||||
t.Fatalf("selected %d tasks, want 500", len(selected))
|
||||
}
|
||||
}
|
||||
|
||||
func TestSelectFairCandidatesHonorsUserCapacity(t *testing.T) {
|
||||
userID := uuid.New()
|
||||
candidates := make([]model.GenerationTask, 11)
|
||||
for index := range candidates {
|
||||
candidates[index] = model.GenerationTask{ID: uuid.New(), UserID: userID}
|
||||
}
|
||||
selected := selectFairCandidates(candidates, map[uuid.UUID]int{}, 500, 10)
|
||||
if len(selected) != 10 {
|
||||
t.Fatalf("selected %d tasks, want 10", len(selected))
|
||||
}
|
||||
}
|
||||
|
||||
func TestSelectFairCandidatesRoundRobin(t *testing.T) {
|
||||
first, second := uuid.New(), uuid.New()
|
||||
candidates := []model.GenerationTask{
|
||||
{ID: uuid.New(), UserID: first}, {ID: uuid.New(), UserID: first},
|
||||
{ID: uuid.New(), UserID: second}, {ID: uuid.New(), UserID: second},
|
||||
}
|
||||
selected := selectFairCandidates(candidates, map[uuid.UUID]int{}, 4, 10)
|
||||
want := []uuid.UUID{first, second, first, second}
|
||||
for index := range want {
|
||||
if selected[index].UserID != want[index] {
|
||||
t.Fatalf("position %d user %s, want %s", index, selected[index].UserID, want[index])
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestCanRetrySubmission 验证提交超时会自动重试,并在第五次失败后停止。
|
||||
func TestCanRetrySubmission(t *testing.T) {
|
||||
for attempt := 1; attempt < maxSubmitAttempts; attempt++ {
|
||||
if !canRetrySubmission(attempt) {
|
||||
t.Fatalf("attempt %d should be retryable", attempt)
|
||||
}
|
||||
}
|
||||
if canRetrySubmission(maxSubmitAttempts) {
|
||||
t.Fatalf("attempt %d should reach retry limit", maxSubmitAttempts)
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,51 @@
|
||||
// 媒体反推任务范围测试,确保单分镜任务不会污染剧集或项目的总处理状态。
|
||||
package worker
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
// TestIsWholeVideoAnalysis 验证仅未绑定分镜的任务可更新反推范围总状态。
|
||||
func TestIsWholeVideoAnalysis(t *testing.T) {
|
||||
storyboardID := uuid.New()
|
||||
tests := []struct {
|
||||
name string
|
||||
snapshot analysisSnapshot
|
||||
want bool
|
||||
}{
|
||||
{name: "整段视频反推", snapshot: analysisSnapshot{}, want: true},
|
||||
{name: "单分镜重新反推", snapshot: analysisSnapshot{StoryboardID: &storyboardID}, want: false},
|
||||
}
|
||||
|
||||
for _, test := range tests {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
if got := isWholeVideoAnalysis(test.snapshot); got != test.want {
|
||||
t.Fatalf("isWholeVideoAnalysis() = %v, want %v", got, test.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestRedrawStoryboardsReady 验证只有全部分镜提示词完整时才恢复剧本反推入口。
|
||||
func TestRedrawStoryboardsReady(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
total int64
|
||||
incomplete int64
|
||||
want bool
|
||||
}{
|
||||
{name: "没有分镜", total: 0, incomplete: 0, want: false},
|
||||
{name: "仍有失败分镜", total: 7, incomplete: 1, want: false},
|
||||
{name: "全部分镜完整", total: 7, incomplete: 0, want: true},
|
||||
}
|
||||
|
||||
for _, test := range tests {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
if got := redrawStoryboardsReady(test.total, test.incomplete); got != test.want {
|
||||
t.Fatalf("redrawStoryboardsReady(%d, %d) = %v, want %v", test.total, test.incomplete, got, test.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,141 @@
|
||||
package worker
|
||||
|
||||
import "testing"
|
||||
|
||||
func TestBuildFixedSegmentsUsesFiveToFifteenSeconds(t *testing.T) {
|
||||
for _, duration := range []float64{5, 15, 19, 31, 33, 39, 45, 46, 121.4} {
|
||||
segments := buildFixedSegments(duration)
|
||||
if len(segments) == 0 {
|
||||
t.Fatalf("duration %.1f returned no segments", duration)
|
||||
}
|
||||
for _, segment := range segments {
|
||||
length := segment.End - segment.Start
|
||||
if length < 5 || length > 15.0001 {
|
||||
t.Fatalf("duration %.1f produced %.3f second segment", duration, length)
|
||||
}
|
||||
}
|
||||
if difference := segments[len(segments)-1].End - duration; difference < -0.001 || difference > 0.001 {
|
||||
t.Fatalf("duration %.1f ends at %.3f", duration, segments[len(segments)-1].End)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestNormalizeAnalysisDurationUsesBoundaryTolerance(t *testing.T) {
|
||||
tests := []struct {
|
||||
input float64
|
||||
want float64
|
||||
}{
|
||||
{input: 4.01, want: 5},
|
||||
{input: 4.999, want: 5},
|
||||
{input: 5, want: 5},
|
||||
{input: 15, want: 15},
|
||||
{input: 15.069002, want: 15},
|
||||
{input: 15.999, want: 15},
|
||||
{input: 16, want: 16},
|
||||
}
|
||||
for _, test := range tests {
|
||||
if got := normalizeAnalysisDuration(test.input); got != test.want {
|
||||
t.Fatalf("normalizeAnalysisDuration(%f) = %f, want %f", test.input, got, test.want)
|
||||
}
|
||||
}
|
||||
segments := buildFixedSegments(normalizeAnalysisDuration(15.069002))
|
||||
if len(segments) != 1 || segments[0].End != 15 {
|
||||
t.Fatalf("15.069002 seconds produced %#v, want one 15 second segment", segments)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildFixedSegmentsRedistributesShortTail(t *testing.T) {
|
||||
tests := []struct {
|
||||
duration float64
|
||||
want []float64
|
||||
}{
|
||||
{duration: 33, want: []float64{15, 9, 9}},
|
||||
{duration: 39, want: []float64{15, 15, 9}},
|
||||
}
|
||||
for _, test := range tests {
|
||||
segments := buildFixedSegments(test.duration)
|
||||
if len(segments) != len(test.want) {
|
||||
t.Fatalf("duration %.1f produced %d segments, want %d", test.duration, len(segments), len(test.want))
|
||||
}
|
||||
for index, segment := range segments {
|
||||
if got := segment.End - segment.Start; got != test.want[index] {
|
||||
t.Fatalf("duration %.1f segment %d = %.1f, want %.1f", test.duration, index+1, got, test.want[index])
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestFrameTimestampsPreserveBoundariesAndLimit(t *testing.T) {
|
||||
segment := fixedSegment{Index: 1, Start: 10, End: 23}
|
||||
frames := frameTimestamps(segment)
|
||||
if len(frames) > 16 {
|
||||
t.Fatalf("returned %d frames, want at most 16", len(frames))
|
||||
}
|
||||
if frames[0] > 10.1 || frames[len(frames)-1] < 22.9 {
|
||||
t.Fatalf("boundaries not preserved: first %.3f last %.3f", frames[0], frames[len(frames)-1])
|
||||
}
|
||||
}
|
||||
|
||||
func TestFrameRetryTimestampMovesBoundarySamplesIntoSegment(t *testing.T) {
|
||||
segment := fixedSegment{Index: 9, Start: 113.31, End: 121.62}
|
||||
if got, ok := frameRetryTimestamp(segment, 121.57); !ok || got != 121.07 {
|
||||
t.Fatalf("end retry = %.3f, %v; want 121.070, true", got, ok)
|
||||
}
|
||||
if got, ok := frameRetryTimestamp(segment, 113.36); !ok || got != 113.86 {
|
||||
t.Fatalf("start retry = %.3f, %v; want 113.860, true", got, ok)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMentionAssetsSkipsDialogue(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
prompt string
|
||||
want string
|
||||
}{
|
||||
{
|
||||
name: "Chinese quotes and all asset types",
|
||||
prompt: "Grace走进古堡大厅,把青铜钥匙放在桌上,说:“Ethan,去古堡大厅找青铜钥匙和神秘徽记。”神秘徽记微微发光。",
|
||||
want: "@Grace走进@古堡大厅,把@青铜钥匙放在桌上,说:“Ethan,去古堡大厅找青铜钥匙和神秘徽记。”@神秘徽记微微发光。",
|
||||
},
|
||||
{
|
||||
name: "English quotes and existing mention",
|
||||
prompt: `@Grace says: "Ethan, take the Key marked \"Grace\"." Ethan nods beside Castle.`,
|
||||
want: `@Grace says: "Ethan, take the Key marked \"Grace\"." @Ethan nods beside @Castle.`,
|
||||
},
|
||||
}
|
||||
names := []string{"Grace", "Ethan", "古堡大厅", "青铜钥匙", "神秘徽记", "Key", "Castle"}
|
||||
for _, test := range tests {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
if got := mentionAssets(test.prompt, names); got != test.want {
|
||||
t.Fatalf("mentionAssets() = %q, want %q", got, test.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestMentionAssetsUsesKnownAssetsAndPrefersLongerNames(t *testing.T) {
|
||||
names := []string{"지우", "지우의 아버지", "도현", "婚礼殿堂"}
|
||||
prompt := "新娘지우走进婚礼殿堂,지우의 아버지看向도현,说:“지우和도현已经到了。”"
|
||||
want := "新娘@지우走进@婚礼殿堂,@지우의 아버지看向@도현,说:“지우和도현已经到了。”"
|
||||
got := mentionAssets(prompt, names)
|
||||
if got != want {
|
||||
t.Fatalf("mentionAssets() = %q, want %q", got, want)
|
||||
}
|
||||
mentioned := mentionedAssetNames(got, names)
|
||||
wantMentioned := []string{"지우", "婚礼殿堂", "지우의 아버지", "도현"}
|
||||
if len(mentioned) != len(wantMentioned) {
|
||||
t.Fatalf("mentionedAssetNames() = %#v, want %#v", mentioned, wantMentioned)
|
||||
}
|
||||
for index := range wantMentioned {
|
||||
if mentioned[index] != wantMentioned[index] {
|
||||
t.Fatalf("mentionedAssetNames() = %#v, want %#v", mentioned, wantMentioned)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseSubtitle(t *testing.T) {
|
||||
items := parseSubtitle("1\n00:00:01.000 --> 00:00:02.500\n第一句\n\n2\n00:00:03,000 --> 00:00:04,000\n2025年,第二句")
|
||||
if len(items) != 2 || items[0].Text != "第一句" || items[1].Text != "2025年,第二句" {
|
||||
t.Fatalf("unexpected subtitles: %#v", items)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,502 @@
|
||||
package worker
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"juhe-factory/api/internal/billing"
|
||||
"juhe-factory/api/internal/model"
|
||||
"juhe-factory/api/internal/provider/apimart"
|
||||
queuepkg "juhe-factory/api/internal/queue"
|
||||
"juhe-factory/api/internal/security"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/hibiken/asynq"
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/clause"
|
||||
)
|
||||
|
||||
type ScriptAnalysis struct {
|
||||
DB *gorm.DB
|
||||
Queue *asynq.Client
|
||||
Encryptor *security.Encryptor
|
||||
Provider apimart.Client
|
||||
}
|
||||
|
||||
type scriptAnalysisSnapshot struct {
|
||||
Task model.GenerationTask
|
||||
AnalysisID uuid.UUID
|
||||
SourceContent string
|
||||
ModelName string
|
||||
BaseURL string
|
||||
APIKeyCiphertext string
|
||||
}
|
||||
|
||||
type scriptHighlight struct {
|
||||
Type string `json:"type"`
|
||||
Content string `json:"content"`
|
||||
}
|
||||
|
||||
type scriptEmotionStage struct {
|
||||
Stage string `json:"stage"`
|
||||
Emotion string `json:"emotion"`
|
||||
Description string `json:"description"`
|
||||
}
|
||||
|
||||
type scriptEpisodeOutline struct {
|
||||
EpisodeNo int `json:"episode_no"`
|
||||
Title string `json:"title"`
|
||||
Outline string `json:"outline"`
|
||||
Emotion string `json:"emotion"`
|
||||
}
|
||||
|
||||
type scriptAnalysisCharacter struct {
|
||||
Name string `json:"name"`
|
||||
Faction string `json:"faction"`
|
||||
Biography string `json:"biography"`
|
||||
Motivation string `json:"motivation"`
|
||||
Relationship string `json:"relationship"`
|
||||
}
|
||||
|
||||
type scriptEpisodeHook struct {
|
||||
EpisodeNo int `json:"episode_no"`
|
||||
HookType string `json:"hook_type"`
|
||||
Hook string `json:"hook"`
|
||||
Basis string `json:"basis"`
|
||||
}
|
||||
|
||||
type scriptPaywallHook struct {
|
||||
HookType string `json:"hook_type"`
|
||||
Content string `json:"content"`
|
||||
Appeal string `json:"appeal"`
|
||||
}
|
||||
|
||||
type scriptCustomAnalysis struct {
|
||||
Objective string `json:"objective"`
|
||||
Conclusion string `json:"conclusion"`
|
||||
Evidence []string `json:"evidence"`
|
||||
}
|
||||
|
||||
type scriptUnmatchedObjective struct {
|
||||
Objective string `json:"objective"`
|
||||
Reason string `json:"reason"`
|
||||
}
|
||||
|
||||
type scriptAnalysisResult struct {
|
||||
Tags struct {
|
||||
DramaTags []string `json:"drama_tags"`
|
||||
ApplicableSceneTags []string `json:"applicable_scene_tags"`
|
||||
} `json:"tags"`
|
||||
SellingPoint struct {
|
||||
Summary string `json:"summary"`
|
||||
Highlights []scriptHighlight `json:"highlights"`
|
||||
} `json:"selling_point"`
|
||||
StoryStructure struct {
|
||||
Synopsis string `json:"synopsis"`
|
||||
StoryOutline []string `json:"story_outline"`
|
||||
EmotionOutline []scriptEmotionStage `json:"emotion_outline"`
|
||||
EpisodeOutlines []scriptEpisodeOutline `json:"episode_outlines"`
|
||||
Characters []scriptAnalysisCharacter `json:"characters"`
|
||||
} `json:"story_structure"`
|
||||
OpeningHook struct {
|
||||
FirstScene string `json:"first_scene"`
|
||||
FirstLine string `json:"first_line"`
|
||||
AttentionMethods []string `json:"attention_methods"`
|
||||
Analysis string `json:"analysis"`
|
||||
Optimization string `json:"optimization"`
|
||||
} `json:"opening_hook"`
|
||||
EpisodeHooks []scriptEpisodeHook `json:"episode_hooks"`
|
||||
Episode10PaywallHooks []scriptPaywallHook `json:"episode_10_paywall_hooks"`
|
||||
CustomAnalysis []scriptCustomAnalysis `json:"custom_analysis"`
|
||||
UnmatchedObjectives []scriptUnmatchedObjective `json:"unmatched_objectives"`
|
||||
}
|
||||
|
||||
func NewScriptAnalysis(db *gorm.DB, queue *asynq.Client, encryptor *security.Encryptor, client *http.Client) *ScriptAnalysis {
|
||||
return &ScriptAnalysis{DB: db, Queue: queue, Encryptor: encryptor, Provider: apimart.NewClient(dramaHTTPClient(client))}
|
||||
}
|
||||
|
||||
func (w *ScriptAnalysis) Register(mux *asynq.ServeMux) {
|
||||
mux.HandleFunc(queuepkg.TypeAnalyzeScript, w.analyze)
|
||||
}
|
||||
|
||||
func (w *ScriptAnalysis) Recover(ctx context.Context) error {
|
||||
if err := w.DB.WithContext(ctx).Model(&model.GenerationTask{}).
|
||||
Where("task_type='script_analysis' AND status='processing'").Update("status", "submitted").Error; err != nil {
|
||||
return err
|
||||
}
|
||||
var ids []uuid.UUID
|
||||
if err := w.DB.WithContext(ctx).Model(&model.GenerationTask{}).
|
||||
Where("task_type='script_analysis' AND status='submitted'").Pluck("id", &ids).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
var firstErr error
|
||||
for _, id := range ids {
|
||||
if err := queuepkg.EnqueueID(w.Queue, queuepkg.TypeAnalyzeScript, id, 0); err != nil && firstErr == nil {
|
||||
firstErr = err
|
||||
}
|
||||
}
|
||||
var cancelledIDs []uuid.UUID
|
||||
if err := w.DB.WithContext(ctx).Model(&model.GenerationTask{}).
|
||||
Where("task_type='script_analysis' AND status='cancel_requested'").Pluck("id", &cancelledIDs).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
for _, id := range cancelledIDs {
|
||||
if err := w.finishCancellation(ctx, id, uuid.Nil); err != nil && firstErr == nil {
|
||||
firstErr = err
|
||||
}
|
||||
}
|
||||
return firstErr
|
||||
}
|
||||
|
||||
func (w *ScriptAnalysis) analyze(ctx context.Context, raw *asynq.Task) error {
|
||||
taskID, err := decodeID(raw)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
snapshot, err := w.claim(ctx, taskID)
|
||||
if err != nil {
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
return nil
|
||||
}
|
||||
return w.fail(ctx, taskID, uuid.Nil, err)
|
||||
}
|
||||
apiKey, err := w.Encryptor.Decrypt(snapshot.APIKeyCiphertext)
|
||||
if err != nil {
|
||||
return w.fail(ctx, taskID, snapshot.AnalysisID, errors.New("剧本分析渠道密钥不可用"))
|
||||
}
|
||||
workCtx, cancel := context.WithTimeout(ctx, 12*time.Minute)
|
||||
defer cancel()
|
||||
var cancelled atomic.Bool
|
||||
go w.watchCancellation(workCtx, cancel, taskID, &cancelled)
|
||||
result, chat, inputTokens, outputTokens, tokenSource, err := w.callModel(workCtx, snapshot, apiKey)
|
||||
if cancelled.Load() || w.cancellationRequested(ctx, taskID) {
|
||||
return w.finishCancellation(ctx, taskID, snapshot.AnalysisID)
|
||||
}
|
||||
if err != nil {
|
||||
return w.fail(ctx, taskID, snapshot.AnalysisID, err)
|
||||
}
|
||||
if err := w.persist(ctx, snapshot, result, chat, inputTokens, outputTokens, tokenSource); err != nil {
|
||||
return w.fail(ctx, taskID, snapshot.AnalysisID, err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (w *ScriptAnalysis) watchCancellation(ctx context.Context, cancel context.CancelFunc, taskID uuid.UUID, flag *atomic.Bool) {
|
||||
ticker := time.NewTicker(time.Second)
|
||||
defer ticker.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-ticker.C:
|
||||
if w.cancellationRequested(ctx, taskID) {
|
||||
flag.Store(true)
|
||||
cancel()
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (w *ScriptAnalysis) cancellationRequested(ctx context.Context, taskID uuid.UUID) bool {
|
||||
var status string
|
||||
return w.DB.WithContext(ctx).Model(&model.GenerationTask{}).Where("id=?", taskID).Pluck("status", &status).Error == nil &&
|
||||
status == "cancel_requested"
|
||||
}
|
||||
|
||||
func (w *ScriptAnalysis) claim(ctx context.Context, taskID uuid.UUID) (scriptAnalysisSnapshot, error) {
|
||||
var snapshot scriptAnalysisSnapshot
|
||||
err := w.DB.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).
|
||||
Where("id=? AND task_type='script_analysis' AND status='submitted'", taskID).Take(&snapshot.Task).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if snapshot.Task.ScriptAnalysisID == nil {
|
||||
return errors.New("剧本分析任务缺少剧本标识")
|
||||
}
|
||||
snapshot.AnalysisID = *snapshot.Task.ScriptAnalysisID
|
||||
var input struct {
|
||||
SourceContent string `json:"source_content"`
|
||||
}
|
||||
if err := json.Unmarshal(snapshot.Task.InputData, &input); err != nil || strings.TrimSpace(input.SourceContent) == "" {
|
||||
return errors.New("剧本分析原文快照无效")
|
||||
}
|
||||
snapshot.SourceContent = input.SourceContent
|
||||
var config struct {
|
||||
ModelName string
|
||||
BaseURL string
|
||||
APIKeyCiphertext string
|
||||
}
|
||||
result := tx.Raw(`SELECT model.name AS model_name,channel.base_url,channel.api_key_ciphertext
|
||||
FROM models model JOIN channels channel ON channel.id=model.channel_id
|
||||
WHERE model.id=? AND channel.id=? AND model.model_type='text'
|
||||
AND model.enabled=true AND model.deleted_at IS NULL AND channel.enabled=true AND channel.deleted_at IS NULL`,
|
||||
snapshot.Task.ModelID, snapshot.Task.ChannelID).Scan(&config)
|
||||
if result.Error != nil {
|
||||
return result.Error
|
||||
}
|
||||
if result.RowsAffected == 0 {
|
||||
return errors.New("剧本分析模型或渠道不可用")
|
||||
}
|
||||
snapshot.ModelName, snapshot.BaseURL, snapshot.APIKeyCiphertext = config.ModelName, config.BaseURL, config.APIKeyCiphertext
|
||||
if err := tx.Model(&snapshot.Task).Update("status", "processing").Error; err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Model(&model.ScriptAnalysis{}).Where("id=?", snapshot.AnalysisID).
|
||||
Updates(map[string]any{"analysis_status": "running", "analysis_message": "正在分析剧本"}).Error
|
||||
})
|
||||
return snapshot, err
|
||||
}
|
||||
|
||||
func (w *ScriptAnalysis) callModel(ctx context.Context, snapshot scriptAnalysisSnapshot, apiKey string) (scriptAnalysisResult, apimart.ChatResult, int64, int64, string, error) {
|
||||
var parsed scriptAnalysisResult
|
||||
var chat apimart.ChatResult
|
||||
pricing, err := billing.ParseTextPricingSnapshot(snapshot.Task.BillingSnapshot)
|
||||
if err != nil {
|
||||
return parsed, chat, 0, 0, "", err
|
||||
}
|
||||
systemPrompt := scriptAnalysisSystemPrompt(snapshot.Task.PromptSnapshot)
|
||||
userPrompt := "请分析以下剧本原文:\n\n" + snapshot.SourceContent
|
||||
if pricing.Mode == billing.TextBillingPerToken {
|
||||
estimated := billing.EstimateTextTokens(systemPrompt + userPrompt)
|
||||
if err := w.DB.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
_, reserveErr := billing.ReserveTextCall(tx, snapshot.Task.UserID, snapshot.Task.ID.String(), "analysis", pricing, estimated, "剧本分析")
|
||||
return reserveErr
|
||||
}); err != nil {
|
||||
return parsed, chat, 0, 0, "", err
|
||||
}
|
||||
}
|
||||
payload := map[string]any{
|
||||
"model": snapshot.ModelName,
|
||||
"messages": []map[string]any{
|
||||
{"role": "system", "content": systemPrompt},
|
||||
{"role": "user", "content": userPrompt},
|
||||
},
|
||||
"temperature": 0.2,
|
||||
"response_format": map[string]any{"type": "json_object"},
|
||||
}
|
||||
var parseErr error
|
||||
for attempt := 0; attempt < 2; attempt++ {
|
||||
if attempt > 0 {
|
||||
payload["temperature"] = 0
|
||||
payload["messages"] = []map[string]any{
|
||||
{"role": "system", "content": systemPrompt},
|
||||
{"role": "user", "content": userPrompt + "\n\n上次返回未通过结构校验。请严格补齐固定协议的全部字段,并按分析目标决定标准字段、补充分析或无法完成的目标,只返回完整合法 JSON。"},
|
||||
}
|
||||
}
|
||||
chat, err = w.Provider.ChatWithUsage(ctx, snapshot.BaseURL, apiKey, payload)
|
||||
if err != nil {
|
||||
return parsed, chat, 0, 0, "", err
|
||||
}
|
||||
parsed = scriptAnalysisResult{}
|
||||
parseErr = json.Unmarshal([]byte(stripJSONFence(chat.Content)), &parsed)
|
||||
if parseErr == nil {
|
||||
parseErr = validateScriptAnalysisResult(parsed)
|
||||
}
|
||||
if parseErr == nil {
|
||||
break
|
||||
}
|
||||
}
|
||||
if parseErr != nil {
|
||||
finishReason := strings.TrimSpace(chat.FinishReason)
|
||||
if finishReason == "" {
|
||||
finishReason = "未知"
|
||||
}
|
||||
return parsed, chat, 0, 0, "", fmt.Errorf("文本模型连续两次返回无法解析(结束原因:%s): %w", finishReason, parseErr)
|
||||
}
|
||||
inputTokens, outputTokens := chat.InputTokens, chat.OutputTokens
|
||||
tokenSource := "upstream"
|
||||
if inputTokens == 0 {
|
||||
inputTokens = billing.EstimateTextTokens(systemPrompt + userPrompt)
|
||||
tokenSource = "local"
|
||||
}
|
||||
if outputTokens == 0 {
|
||||
outputTokens = billing.EstimateTextTokens(chat.Content)
|
||||
if tokenSource == "upstream" {
|
||||
tokenSource = "mixed"
|
||||
}
|
||||
}
|
||||
if pricing.Mode == billing.TextBillingPerToken {
|
||||
if err := w.DB.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
settled, settleErr := billing.SettleTextCall(tx, snapshot.Task.UserID, snapshot.Task.ID.String(), "analysis", pricing,
|
||||
billing.TextUsage{Input: inputTokens, Output: outputTokens}, "剧本分析")
|
||||
if settleErr != nil {
|
||||
return settleErr
|
||||
}
|
||||
return tx.Model(&model.GenerationTask{}).Where("id=?", snapshot.Task.ID).Updates(map[string]any{
|
||||
"estimated_points": gorm.Expr("estimated_points+?::numeric", settled),
|
||||
"prepaid_points": gorm.Expr("prepaid_points+?::numeric", settled),
|
||||
}).Error
|
||||
}); err != nil {
|
||||
return parsed, chat, inputTokens, outputTokens, tokenSource, err
|
||||
}
|
||||
}
|
||||
return parsed, chat, inputTokens, outputTokens, tokenSource, nil
|
||||
}
|
||||
|
||||
func validateScriptAnalysisResult(result scriptAnalysisResult) error {
|
||||
if strings.TrimSpace(result.SellingPoint.Summary) == "" && strings.TrimSpace(result.StoryStructure.Synopsis) == "" &&
|
||||
len(result.CustomAnalysis) == 0 && len(result.UnmatchedObjectives) == 0 {
|
||||
return errors.New("缺少标准分析结果、通用分析结果或无法匹配目标说明")
|
||||
}
|
||||
if len(result.EpisodeHooks) != 0 && len(result.EpisodeHooks) != 10 {
|
||||
return errors.New("前十集钩子数量必须为10")
|
||||
}
|
||||
for index, hook := range result.EpisodeHooks {
|
||||
if hook.EpisodeNo != index+1 {
|
||||
return errors.New("前十集钩子集数或顺序无效")
|
||||
}
|
||||
}
|
||||
if len(result.Episode10PaywallHooks) != 0 &&
|
||||
(len(result.Episode10PaywallHooks) < 3 || len(result.Episode10PaywallHooks) > 5) {
|
||||
return errors.New("第十集付费卡点候选必须为3至5个")
|
||||
}
|
||||
for _, character := range result.StoryStructure.Characters {
|
||||
if strings.TrimSpace(character.Name) == "" {
|
||||
return errors.New("人物名称不能为空")
|
||||
}
|
||||
if character.Faction != "正派" && character.Faction != "反派" && character.Faction != "中立派" {
|
||||
return errors.New("人物阵营无效")
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (w *ScriptAnalysis) persist(ctx context.Context, snapshot scriptAnalysisSnapshot, result scriptAnalysisResult, chat apimart.ChatResult, inputTokens, outputTokens int64, tokenSource string) error {
|
||||
encoded, err := json.Marshal(result)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
pretty, _ := json.MarshalIndent(result, "", " ")
|
||||
return w.DB.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
var task model.GenerationTask
|
||||
if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Where("id=?", snapshot.Task.ID).Take(&task).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if task.Status == "cancel_requested" {
|
||||
return w.finishCancellationLocked(tx, &task, snapshot.AnalysisID)
|
||||
}
|
||||
if task.Status != "processing" {
|
||||
return errors.New("剧本分析任务状态已变化")
|
||||
}
|
||||
if err := tx.Where("script_analysis_id=?", snapshot.AnalysisID).Delete(&model.ScriptAnalysisCharacter{}).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
seen := make(map[string]bool)
|
||||
for index, character := range result.StoryStructure.Characters {
|
||||
name := strings.TrimSpace(character.Name)
|
||||
key := strings.ToLower(name)
|
||||
if seen[key] {
|
||||
continue
|
||||
}
|
||||
seen[key] = true
|
||||
item := model.ScriptAnalysisCharacter{
|
||||
ScriptAnalysisID: snapshot.AnalysisID, Name: name, Faction: character.Faction,
|
||||
Biography: strings.TrimSpace(character.Biography), Motivation: strings.TrimSpace(character.Motivation),
|
||||
Relationship: strings.TrimSpace(character.Relationship), SortOrder: index,
|
||||
}
|
||||
if err := tx.Create(&item).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
taskUpdates := map[string]any{
|
||||
"status": "succeeded", "finished_at": time.Now(), "actual_points": task.PrepaidPoints,
|
||||
"input_tokens": inputTokens, "output_tokens": outputTokens, "total_tokens": inputTokens + outputTokens,
|
||||
"token_count_source": tokenSource, "usage_raw": chat.UsageRaw, "result_data": gorm.Expr("?::jsonb", string(encoded)),
|
||||
"error_code": nil, "error_message": nil,
|
||||
}
|
||||
if err := tx.Model(&task).Updates(taskUpdates).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Model(&model.ScriptAnalysis{}).Where("id=?", snapshot.AnalysisID).Updates(map[string]any{
|
||||
"analysis_result": gorm.Expr("?::jsonb", string(encoded)), "result_content": string(pretty),
|
||||
"analysis_status": "succeeded", "analysis_message": "剧本分析完成",
|
||||
}).Error
|
||||
})
|
||||
}
|
||||
|
||||
func (w *ScriptAnalysis) fail(ctx context.Context, taskID, analysisID uuid.UUID, taskErr error) error {
|
||||
return w.DB.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
var task model.GenerationTask
|
||||
if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Where("id=?", taskID).Take(&task).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if task.Status == "succeeded" || task.Status == "failed" || task.Status == "cancelled" {
|
||||
return nil
|
||||
}
|
||||
if task.Status == "cancel_requested" {
|
||||
return w.finishCancellationLocked(tx, &task, analysisID)
|
||||
}
|
||||
refunded, err := billing.RefundTextGenerationTask(tx, &task, "剧本分析失败返还")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
updates := map[string]any{
|
||||
"status": "failed", "actual_points": "0.00", "finished_at": time.Now(),
|
||||
"error_code": "script_analysis_failed", "error_message": truncate(taskErr.Error(), 4000),
|
||||
}
|
||||
if refunded {
|
||||
updates["cost_refunded"] = true
|
||||
}
|
||||
if err := tx.Model(&task).Updates(updates).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if analysisID == uuid.Nil && task.ScriptAnalysisID != nil {
|
||||
analysisID = *task.ScriptAnalysisID
|
||||
}
|
||||
if analysisID == uuid.Nil {
|
||||
return nil
|
||||
}
|
||||
return tx.Model(&model.ScriptAnalysis{}).Where("id=?", analysisID).Updates(map[string]any{
|
||||
"analysis_status": "failed", "analysis_message": truncate(taskErr.Error(), 1000),
|
||||
}).Error
|
||||
})
|
||||
}
|
||||
|
||||
func (w *ScriptAnalysis) finishCancellation(ctx context.Context, taskID, analysisID uuid.UUID) error {
|
||||
return w.DB.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
var task model.GenerationTask
|
||||
if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Where("id=?", taskID).Take(&task).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
return w.finishCancellationLocked(tx, &task, analysisID)
|
||||
})
|
||||
}
|
||||
|
||||
func (w *ScriptAnalysis) finishCancellationLocked(tx *gorm.DB, task *model.GenerationTask, analysisID uuid.UUID) error {
|
||||
if task.Status == "cancelled" {
|
||||
return nil
|
||||
}
|
||||
if task.Status == "succeeded" || task.Status == "failed" {
|
||||
return nil
|
||||
}
|
||||
actual := task.PrepaidPoints
|
||||
committed, err := billing.CommitTextReserves(tx, task.UserID, task.ID.String(), "剧本分析取消")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if committed != "" && committed != "0" && committed != "0.00" {
|
||||
actual = committed
|
||||
}
|
||||
if err := tx.Model(task).Updates(map[string]any{
|
||||
"status": "cancelled", "prepaid_points": actual, "actual_points": actual, "finished_at": time.Now(),
|
||||
"error_code": "user_cancelled", "error_message": "用户取消,已扣积分不退",
|
||||
}).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if analysisID == uuid.Nil && task.ScriptAnalysisID != nil {
|
||||
analysisID = *task.ScriptAnalysisID
|
||||
}
|
||||
if analysisID == uuid.Nil {
|
||||
return nil
|
||||
}
|
||||
return tx.Model(&model.ScriptAnalysis{}).Where("id=?", analysisID).Updates(map[string]any{
|
||||
"analysis_status": "idle", "analysis_message": "剧本分析已取消,已扣积分不退",
|
||||
}).Error
|
||||
}
|
||||
@@ -0,0 +1,61 @@
|
||||
package worker
|
||||
|
||||
import "strings"
|
||||
|
||||
const scriptAnalysisOutputRules = `将分析目标能够对应到标准字段的结果写入标准字段;无法对应到标准字段的目标写入 custom_analysis;因原文信息不足或目标不适用于剧本而无法完成的目标写入 unmatched_objectives。不得为了填充标准字段而曲解分析目标。
|
||||
|
||||
必须只返回一个合法 JSON 对象,不得输出 Markdown、代码围栏、解释或额外文字。所有字段均须存在并严格符合以下结构:
|
||||
{
|
||||
"tags": {
|
||||
"drama_tags": ["标签"],
|
||||
"applicable_scene_tags": ["适用场景标签"]
|
||||
},
|
||||
"selling_point": {
|
||||
"summary": "一句话总结",
|
||||
"highlights": [
|
||||
{"type": "爽点|痛点|卖点|看点", "content": "具体内容"}
|
||||
]
|
||||
},
|
||||
"story_structure": {
|
||||
"synopsis": "短剧简介",
|
||||
"story_outline": ["按故事阶段组织的大纲"],
|
||||
"emotion_outline": [
|
||||
{"stage": "情绪阶段", "emotion": "主导情绪", "description": "情绪变化"}
|
||||
],
|
||||
"episode_outlines": [
|
||||
{"episode_no": 1, "title": "集标题", "outline": "本集大纲", "emotion": "本集情绪"}
|
||||
],
|
||||
"characters": [
|
||||
{"name": "人物名", "faction": "正派|反派|中立派", "biography": "人物小传", "motivation": "核心动机", "relationship": "关键人物关系"}
|
||||
]
|
||||
},
|
||||
"opening_hook": {
|
||||
"first_scene": "开篇第一个画面或场景",
|
||||
"first_line": "第一句台词",
|
||||
"attention_methods": ["悬念前置|激烈冲突|视觉奇观|其他明确手法"],
|
||||
"analysis": "为什么能或不能抓住观众",
|
||||
"optimization": "必要时给出优化建议"
|
||||
},
|
||||
"episode_hooks": [
|
||||
{"episode_no": 1, "hook_type": "钩子类型", "hook": "本集结尾或关键卡点钩子", "basis": "对应的原文依据或信息不足说明"}
|
||||
],
|
||||
"episode_10_paywall_hooks": [
|
||||
{"hook_type": "钩子类型", "content": "付费卡点候选", "appeal": "促使继续付费观看的原因"}
|
||||
],
|
||||
"custom_analysis": [
|
||||
{"objective": "无法对应标准字段的分析目标", "conclusion": "分析结论", "evidence": ["剧本原文依据"]}
|
||||
],
|
||||
"unmatched_objectives": [
|
||||
{"objective": "无法完成的分析目标", "reason": "无法完成的原因"}
|
||||
]
|
||||
}
|
||||
|
||||
仅当分析目标要求逐集钩子时,episode_hooks 必须按 episode_no 从1到10排列且恰好返回10项;仅当分析目标要求第10集付费卡点时,episode_10_paywall_hooks 返回3至5项。其他情况下对应数组返回空数组。`
|
||||
|
||||
func scriptAnalysisSystemPrompt(analysisGoal string) string {
|
||||
goal := strings.TrimSpace(analysisGoal)
|
||||
if goal == "" {
|
||||
return scriptAnalysisOutputRules
|
||||
}
|
||||
return goal + "\n\n" + scriptAnalysisOutputRules
|
||||
}
|
||||
@@ -0,0 +1,19 @@
|
||||
package worker
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestScriptAnalysisSystemPromptKeepsGoalAndFixedSchema(t *testing.T) {
|
||||
prompt := scriptAnalysisSystemPrompt("你是一名分析师。\n\n分析目标:\n分析拍摄预算")
|
||||
if !strings.Contains(prompt, "分析目标:\n分析拍摄预算") {
|
||||
t.Fatal("prompt does not contain the analysis goal")
|
||||
}
|
||||
if strings.Count(prompt, "你是一名分析师") != 1 {
|
||||
t.Fatal("prompt unexpectedly duplicates the database prompt")
|
||||
}
|
||||
if !strings.Contains(prompt, `"custom_analysis"`) || !strings.Contains(prompt, `"unmatched_objectives"`) {
|
||||
t.Fatal("prompt does not contain the fixed fallback schema")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,38 @@
|
||||
package worker
|
||||
|
||||
import "testing"
|
||||
|
||||
func TestValidateScriptAnalysisResult(t *testing.T) {
|
||||
var result scriptAnalysisResult
|
||||
result.SellingPoint.Summary = "核心卖点"
|
||||
result.StoryStructure.Synopsis = "剧情梗概"
|
||||
for episode := 1; episode <= 10; episode++ {
|
||||
result.EpisodeHooks = append(result.EpisodeHooks, scriptEpisodeHook{EpisodeNo: episode})
|
||||
}
|
||||
for index := 0; index < 3; index++ {
|
||||
result.Episode10PaywallHooks = append(result.Episode10PaywallHooks, scriptPaywallHook{Content: "候选"})
|
||||
}
|
||||
result.StoryStructure.Characters = []scriptAnalysisCharacter{{Name: "甲", Faction: "正派"}, {Name: "乙", Faction: "反派"}, {Name: "丙", Faction: "中立派"}}
|
||||
if err := validateScriptAnalysisResult(result); err != nil {
|
||||
t.Fatalf("expected valid analysis result: %v", err)
|
||||
}
|
||||
|
||||
result.EpisodeHooks[9].EpisodeNo = 9
|
||||
if err := validateScriptAnalysisResult(result); err == nil {
|
||||
t.Fatal("expected invalid episode hook ordering")
|
||||
}
|
||||
}
|
||||
|
||||
func TestValidateScriptAnalysisResultRejectsInvalidFaction(t *testing.T) {
|
||||
var result scriptAnalysisResult
|
||||
result.SellingPoint.Summary = "核心卖点"
|
||||
result.StoryStructure.Synopsis = "剧情梗概"
|
||||
for episode := 1; episode <= 10; episode++ {
|
||||
result.EpisodeHooks = append(result.EpisodeHooks, scriptEpisodeHook{EpisodeNo: episode})
|
||||
}
|
||||
result.Episode10PaywallHooks = make([]scriptPaywallHook, 3)
|
||||
result.StoryStructure.Characters = []scriptAnalysisCharacter{{Name: "甲", Faction: "未知"}}
|
||||
if err := validateScriptAnalysisResult(result); err == nil {
|
||||
t.Fatal("expected invalid character faction")
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user