初始化

This commit is contained in:
Ran
2026-08-25 17:59:42 +08:00
commit 4b7380dd9b
408 changed files with 327400 additions and 0 deletions
+59
View File
@@ -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
}
+250
View File
@@ -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
}
+30
View File
@@ -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)
}
}
}
+321
View File
@@ -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
}
+70
View File
@@ -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)
}
}
+25
View File
@@ -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
}
+241
View File
@@ -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
}
+76
View File
@@ -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{},
}
}
+59
View File
@@ -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()
}
+133
View File
@@ -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
+30
View File
@@ -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) }
+20
View File
@@ -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
}
+18
View File
@@ -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")
}
}
+111
View File
@@ -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")})
}
+544
View File
@@ -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()
}
}
+210
View File
@@ -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
}
+48
View File
@@ -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,
})
}
+236
View File
@@ -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)
}
+29
View File
@@ -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)
}
})
}
}
+257
View File
@@ -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})
}
+166
View File
@@ -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)
}
+226
View File
@@ -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
}
+156
View File
@@ -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"
}
+17
View File
@@ -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)
}
}
+46
View File
@@ -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" }
+299
View File
@@ -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"`
}
+19
View File
@@ -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)
}
}
+36
View File
@@ -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"`
}
+39
View File
@@ -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" }
+306
View File
@@ -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)
}
})
}
}
+328
View File
@@ -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)
}
}
+64
View File
@@ -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)
}
}
+141
View File
@@ -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)
}
}
+64
View File
@@ -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
}
+103
View File
@@ -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)
}
}
+225
View File
@@ -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")
}
}
+29
View File
@@ -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)
}
+15
View File
@@ -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,
})
}
+17
View File
@@ -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}})
}
+48
View File
@@ -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
}
+57
View File
@@ -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
}
+55
View File
@@ -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, &parallelism); 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
}
+63
View File
@@ -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[:])
}
+47
View File
@@ -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
}
+271
View File
@@ -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()
}
}
+42
View File
@@ -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)
}
})
}
}
+575
View File
@@ -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")
}
})
}
}
+179
View File
@@ -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
}
+78
View File
@@ -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)
}
}
+113
View File
@@ -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
}
+152
View File
@@ -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)
}
}
+735
View File
@@ -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"
}
+544
View File
@@ -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 })
}
+95
View File
@@ -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()
}
+232
View File
@@ -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")
}
}
+443
View File
@@ -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
}
+298
View File
@@ -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)
}
}
+272
View File
@@ -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)
}
})
}
}
+213
View File
@@ -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
}
+14
View File
@@ -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)
}
}
+75
View File
@@ -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
}
+665
View File
@@ -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"
}
+95
View File
@@ -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)
}
}
+850
View File
@@ -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(&current).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(&current).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}
}
+63
View File
@@ -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
+51
View File
@@ -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)
}
})
}
}
+141
View File
@@ -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)
}
}
+502
View File
@@ -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")
}
}