初始化
This commit is contained in:
Ran
@@ -0,0 +1,665 @@
|
||||
package worker
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"regexp"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"juhe-factory/api/internal/billing"
|
||||
dramapkg "juhe-factory/api/internal/drama"
|
||||
"juhe-factory/api/internal/model"
|
||||
"juhe-factory/api/internal/provider"
|
||||
"juhe-factory/api/internal/provider/apimart"
|
||||
queuepkg "juhe-factory/api/internal/queue"
|
||||
"juhe-factory/api/internal/security"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/hibiken/asynq"
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/clause"
|
||||
)
|
||||
|
||||
var (
|
||||
errDramaCancelled = errors.New("drama parse cancelled")
|
||||
errDramaInactive = errors.New("drama parse inactive")
|
||||
)
|
||||
|
||||
const maxDramaModelAttempts = 5
|
||||
|
||||
type DramaParse struct {
|
||||
DB *gorm.DB
|
||||
Queue *asynq.Client
|
||||
Encryptor *security.Encryptor
|
||||
Provider apimart.Client
|
||||
WorkerID string
|
||||
}
|
||||
|
||||
type dramaParseSnapshot struct {
|
||||
Task model.DramaParseTask
|
||||
Source model.EpisodeSource
|
||||
EpisodeNo int
|
||||
EpisodeName string
|
||||
ProjectName string
|
||||
StyleName string
|
||||
EraType string
|
||||
CustomEra string
|
||||
ModelName string
|
||||
BaseURL string
|
||||
APIKeyCiphertext string
|
||||
}
|
||||
|
||||
type dramaParseConfig struct {
|
||||
EpisodeNo int
|
||||
EpisodeName string
|
||||
ProjectName string
|
||||
StyleName string
|
||||
EraType string
|
||||
CustomEra string
|
||||
ModelName string
|
||||
BaseURL string
|
||||
APIKeyCiphertext string
|
||||
}
|
||||
|
||||
type parsedEntity struct {
|
||||
CanonicalName string `json:"canonical_name"`
|
||||
Name string `json:"name"`
|
||||
Aliases []string `json:"aliases"`
|
||||
Description string `json:"description"`
|
||||
ImagePrompt string `json:"image_prompt"`
|
||||
Attributes map[string]any `json:"attributes"`
|
||||
}
|
||||
|
||||
type parsedEntities struct {
|
||||
Characters []parsedEntity `json:"characters"`
|
||||
Roles []parsedEntity `json:"roles"`
|
||||
Scenes []parsedEntity `json:"scenes"`
|
||||
Props []parsedEntity `json:"props"`
|
||||
}
|
||||
|
||||
type parsedStoryboard struct {
|
||||
Title string `json:"title"`
|
||||
ScriptContent string `json:"script_content"`
|
||||
Content string `json:"content"`
|
||||
PromptContent string `json:"prompt_content"`
|
||||
Dialogue []map[string]any `json:"dialogue"`
|
||||
AssetNames struct {
|
||||
Characters []string `json:"characters"`
|
||||
Scenes []string `json:"scenes"`
|
||||
Props []string `json:"props"`
|
||||
} `json:"asset_names"`
|
||||
DurationSeconds int `json:"duration_seconds"`
|
||||
SourceExcerpt string `json:"source_excerpt"`
|
||||
}
|
||||
|
||||
type dramaParsedResult struct {
|
||||
Entities parsedEntities `json:"entities"`
|
||||
Storyboards []parsedStoryboard `json:"storyboards"`
|
||||
}
|
||||
|
||||
type tokenTotals struct {
|
||||
Input, Output int64
|
||||
Source string
|
||||
Segments []map[string]any
|
||||
}
|
||||
|
||||
func NewDramaParse(db *gorm.DB, queue *asynq.Client, encryptor *security.Encryptor, client *http.Client) *DramaParse {
|
||||
return &DramaParse{DB: db, Queue: queue, Encryptor: encryptor, Provider: apimart.NewClient(dramaHTTPClient(client)), WorkerID: uuid.NewString()}
|
||||
}
|
||||
|
||||
func dramaHTTPClient(client *http.Client) *http.Client {
|
||||
cloned := *client
|
||||
if transport, ok := client.Transport.(*http.Transport); ok {
|
||||
transport = transport.Clone()
|
||||
transport.TLSHandshakeTimeout = 30 * time.Second
|
||||
transport.ResponseHeaderTimeout = 10 * time.Minute
|
||||
cloned.Transport = transport
|
||||
}
|
||||
return &cloned
|
||||
}
|
||||
|
||||
func (w *DramaParse) Register(mux *asynq.ServeMux) {
|
||||
mux.HandleFunc(queuepkg.TypeParseDramaEpisode, w.parseEpisode)
|
||||
}
|
||||
|
||||
// Recover 恢复排队、等待重试或租约已过期的剧本解析任务。
|
||||
func (w *DramaParse) Recover(ctx context.Context) error {
|
||||
var tasks []model.DramaParseTask
|
||||
if err := w.DB.WithContext(ctx).Where("status='queued' OR (status IN ('running','retry_wait') AND (lease_until IS NULL OR lease_until<?))", time.Now()).Find(&tasks).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
var firstErr error
|
||||
for _, task := range tasks {
|
||||
w.DB.Model(&model.DramaParseTask{}).Where("id=?", task.ID).Updates(map[string]any{"status": "queued", "lease_owner": nil, "lease_until": nil})
|
||||
if err := queuepkg.EnqueueID(w.Queue, queuepkg.TypeParseDramaEpisode, task.ID, 0); err != nil && firstErr == nil {
|
||||
firstErr = err
|
||||
}
|
||||
}
|
||||
return firstErr
|
||||
}
|
||||
|
||||
func (w *DramaParse) parseEpisode(ctx context.Context, raw *asynq.Task) error {
|
||||
taskID, err := decodeID(raw)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
snapshot, err := w.claim(ctx, taskID)
|
||||
if err != nil {
|
||||
if errors.Is(err, errDramaCancelled) {
|
||||
return w.cancel(snapshot.Task)
|
||||
}
|
||||
if errors.Is(err, errDramaInactive) || snapshot.Task.ID == uuid.Nil {
|
||||
return nil
|
||||
}
|
||||
return w.fail(snapshot.Task, "task_claim_failed", err.Error())
|
||||
}
|
||||
apiKey, err := w.Encryptor.Decrypt(snapshot.APIKeyCiphertext)
|
||||
if err != nil {
|
||||
return w.fail(snapshot.Task, "channel_secret_error", "渠道密钥不可用")
|
||||
}
|
||||
workCtx, cancel := context.WithTimeout(ctx, 12*time.Minute)
|
||||
defer cancel()
|
||||
var cancelled atomic.Bool
|
||||
go w.watchCancellation(workCtx, cancel, taskID, &cancelled)
|
||||
result, tokens, contextSnapshot, err := w.callModel(workCtx, snapshot, apiKey)
|
||||
if err != nil {
|
||||
if cancelled.Load() {
|
||||
return w.cancel(snapshot.Task, tokens)
|
||||
}
|
||||
return w.fail(snapshot.Task, classifyDramaError(err), err.Error(), tokens)
|
||||
}
|
||||
if cancelled.Load() {
|
||||
return w.cancel(snapshot.Task, tokens)
|
||||
}
|
||||
if err = w.persist(snapshot, result, tokens, contextSnapshot); err != nil {
|
||||
if errors.Is(err, errDramaCancelled) {
|
||||
return w.cancel(snapshot.Task, tokens)
|
||||
}
|
||||
return w.fail(snapshot.Task, "result_persist_failed", err.Error())
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (w *DramaParse) claim(ctx context.Context, taskID uuid.UUID) (dramaParseSnapshot, error) {
|
||||
var snapshot dramaParseSnapshot
|
||||
err := w.DB.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
var task model.DramaParseTask
|
||||
if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Where("id=?", taskID).Take(&task).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
snapshot.Task = task
|
||||
if task.Status == "cancel_requested" {
|
||||
return errDramaCancelled
|
||||
}
|
||||
if task.Status != "queued" && task.Status != "retry_wait" {
|
||||
return errDramaInactive
|
||||
}
|
||||
now := time.Now()
|
||||
lease := now.Add(13 * time.Minute)
|
||||
if err := tx.Model(&task).Updates(map[string]any{"status": "running", "attempt_count": gorm.Expr("attempt_count+1"), "started_at": now, "lease_owner": w.WorkerID, "lease_until": lease, "error_code": nil, "error_message": nil}).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if err := tx.Model(&model.ProjectEpisode{}).Where("id=?", task.EpisodeID).Update("analysis_message", "剧本解析中").Error; err != nil {
|
||||
return err
|
||||
}
|
||||
snapshot.Task.Status = "running"
|
||||
if err := tx.Where("episode_id=? AND content_sha256=?", task.EpisodeID, task.SourceSHA256).Take(&snapshot.Source).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
var config dramaParseConfig
|
||||
result := tx.Raw(`SELECT episode.episode_no,episode.name AS episode_name,project.name AS project_name,style.name AS style_name,
|
||||
project.era_type,coalesce(project.custom_era,'') AS custom_era,model.name AS model_name,channel.base_url,channel.api_key_ciphertext
|
||||
FROM project_episodes episode JOIN creative_projects project ON project.id=episode.project_id
|
||||
JOIN project_styles style ON style.id=project.style_id
|
||||
JOIN models model ON model.id=? AND model.enabled=true AND model.deleted_at IS NULL JOIN channels channel ON channel.id=model.channel_id AND channel.id=? AND channel.enabled=true AND channel.deleted_at IS NULL
|
||||
WHERE episode.id=? AND project.id=? AND project.project_type='premium_drama'`, task.ModelID, task.ChannelID, task.EpisodeID, task.ProjectID).Scan(&config)
|
||||
if result.Error != nil {
|
||||
return result.Error
|
||||
}
|
||||
if result.RowsAffected == 0 {
|
||||
return errors.New("解析模型或渠道配置不可用")
|
||||
}
|
||||
snapshot.EpisodeNo = config.EpisodeNo
|
||||
snapshot.EpisodeName = config.EpisodeName
|
||||
snapshot.ProjectName = config.ProjectName
|
||||
snapshot.StyleName = config.StyleName
|
||||
snapshot.EraType = config.EraType
|
||||
snapshot.CustomEra = config.CustomEra
|
||||
snapshot.ModelName = config.ModelName
|
||||
snapshot.BaseURL = config.BaseURL
|
||||
snapshot.APIKeyCiphertext = config.APIKeyCiphertext
|
||||
return nil
|
||||
})
|
||||
return snapshot, err
|
||||
}
|
||||
|
||||
func (w *DramaParse) watchCancellation(ctx context.Context, cancel context.CancelFunc, taskID uuid.UUID, flag *atomic.Bool) {
|
||||
ticker := time.NewTicker(time.Second)
|
||||
defer ticker.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-ticker.C:
|
||||
var status string
|
||||
if w.DB.Model(&model.DramaParseTask{}).Where("id=?", taskID).Pluck("status", &status).Error == nil && status == "cancel_requested" {
|
||||
flag.Store(true)
|
||||
cancel()
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (w *DramaParse) callModel(ctx context.Context, snapshot dramaParseSnapshot, apiKey string) (dramaParsedResult, tokenTotals, json.RawMessage, error) {
|
||||
if err := dramapkg.ValidateEpisodeContent(snapshot.Source.RawContent); err != nil {
|
||||
return dramaParsedResult{}, tokenTotals{}, nil, err
|
||||
}
|
||||
assets, err := w.buildAssetContext(snapshot.Task.ProjectID)
|
||||
if err != nil {
|
||||
return dramaParsedResult{}, tokenTotals{}, nil, err
|
||||
}
|
||||
contextJSON, _ := json.Marshal(map[string]any{"assets": assets})
|
||||
userContent := buildDramaUserPrompt(snapshot, snapshot.Source.RawContent, assets)
|
||||
payload := map[string]any{
|
||||
"model": snapshot.ModelName,
|
||||
"messages": []any{map[string]any{"role": "system", "content": snapshot.Task.PromptSnapshot}, map[string]any{"role": "user", "content": userContent}},
|
||||
"temperature": 0.2,
|
||||
"response_format": map[string]any{"type": "json_object"},
|
||||
}
|
||||
pricing, err := billing.ParseTextPricingSnapshot(snapshot.Task.BillingSnapshot)
|
||||
if err != nil {
|
||||
return dramaParsedResult{}, tokenTotals{}, contextJSON, err
|
||||
}
|
||||
if pricing.Mode == billing.TextBillingPerToken {
|
||||
if err := w.DB.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
_, reserveErr := billing.ReserveTextCall(tx, snapshot.Task.UserID, snapshot.Task.ID.String(), "parse", pricing, billing.EstimateTextTokens(snapshot.Task.PromptSnapshot+userContent), "剧本解析")
|
||||
return reserveErr
|
||||
}); err != nil {
|
||||
return dramaParsedResult{}, tokenTotals{}, contextJSON, err
|
||||
}
|
||||
}
|
||||
chat, callErr := w.chatWithRetry(ctx, snapshot.BaseURL, apiKey, payload)
|
||||
if callErr != nil {
|
||||
return dramaParsedResult{}, tokenTotals{}, contextJSON, callErr
|
||||
}
|
||||
input, output := chat.InputTokens, chat.OutputTokens
|
||||
source := "upstream"
|
||||
if input == 0 {
|
||||
input = billing.EstimateTextTokens(snapshot.Task.PromptSnapshot + userContent)
|
||||
source = "local"
|
||||
}
|
||||
if output == 0 {
|
||||
output = billing.EstimateTextTokens(chat.Content)
|
||||
if source == "upstream" {
|
||||
source = "mixed"
|
||||
}
|
||||
}
|
||||
if pricing.Mode == billing.TextBillingPerToken {
|
||||
if err := w.DB.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
settled, settleErr := billing.SettleTextCall(tx, snapshot.Task.UserID, snapshot.Task.ID.String(), "parse", pricing, billing.TextUsage{Input: chat.InputTokens, Output: chat.OutputTokens}, "剧本解析")
|
||||
if settleErr != nil {
|
||||
return settleErr
|
||||
}
|
||||
if err := tx.Model(&model.DramaParseTask{}).Where("id=?", snapshot.Task.ID).Updates(map[string]any{
|
||||
"estimated_points": gorm.Expr("estimated_points+?::numeric", settled),
|
||||
"prepaid_points": gorm.Expr("prepaid_points+?::numeric", settled),
|
||||
"input_tokens": chat.InputTokens, "output_tokens": chat.OutputTokens, "total_tokens": chat.InputTokens + chat.OutputTokens,
|
||||
"token_count_source": source, "usage_raw": gorm.Expr("?::jsonb", string(chat.UsageRaw)),
|
||||
}).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}); err != nil {
|
||||
return dramaParsedResult{}, tokenTotals{}, contextJSON, err
|
||||
}
|
||||
}
|
||||
usage := map[string]any{"index": 1, "input_tokens": input, "output_tokens": output, "source": source}
|
||||
if len(chat.UsageRaw) > 0 {
|
||||
var upstreamUsage any
|
||||
if json.Unmarshal(chat.UsageRaw, &upstreamUsage) == nil {
|
||||
usage["upstream_usage"] = upstreamUsage
|
||||
}
|
||||
}
|
||||
totals := tokenTotals{Input: input, Output: output, Source: source, Segments: []map[string]any{usage}}
|
||||
parsed, parseErr := parseDramaJSON(chat.Content)
|
||||
if parseErr != nil {
|
||||
finishReason := strings.TrimSpace(chat.FinishReason)
|
||||
if finishReason == "" {
|
||||
finishReason = "未知"
|
||||
}
|
||||
return dramaParsedResult{}, totals, contextJSON, fmt.Errorf("文本模型返回无法解析（结束原因：%s）: %w", finishReason, parseErr)
|
||||
}
|
||||
if len(parsed.Storyboards) == 0 {
|
||||
return parsed, totals, contextJSON, errors.New("文本模型未返回有效分镜")
|
||||
}
|
||||
return parsed, totals, contextJSON, nil
|
||||
}
|
||||
|
||||
func (w *DramaParse) chatWithRetry(ctx context.Context, baseURL, apiKey string, payload map[string]any) (apimart.ChatResult, error) {
|
||||
var result apimart.ChatResult
|
||||
var err error
|
||||
for attempt := 0; attempt < maxDramaModelAttempts; attempt++ {
|
||||
if attempt > 0 {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return result, ctx.Err()
|
||||
case <-time.After(time.Duration(1<<attempt) * time.Second):
|
||||
}
|
||||
}
|
||||
result, err = w.Provider.ChatWithUsage(ctx, baseURL, apiKey, payload)
|
||||
if err == nil {
|
||||
return result, nil
|
||||
}
|
||||
var httpErr *provider.HTTPError
|
||||
if errors.As(err, &httpErr) && httpErr.StatusCode < 500 && httpErr.StatusCode != http.StatusTooManyRequests {
|
||||
return result, err
|
||||
}
|
||||
}
|
||||
return result, err
|
||||
}
|
||||
|
||||
func (w *DramaParse) buildAssetContext(projectID uuid.UUID) ([]map[string]any, error) {
|
||||
assets := make([]map[string]any, 0)
|
||||
if err := w.DB.Table("project_assets").Select("asset_type,name,aliases,description,attributes").Where("project_id=? AND deleted_at IS NULL", projectID).Order("asset_type,name").Find(&assets).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return assets, nil
|
||||
}
|
||||
|
||||
func buildDramaUserPrompt(snapshot dramaParseSnapshot, source string, assets []map[string]any) string {
|
||||
assetsJSON, _ := json.Marshal(assets)
|
||||
return fmt.Sprintf("项目：%s\n视觉风格：%s\n故事时代：%s\n当前剧集：第%d集 %s\n项目资产：\n%s\n\n当前原文：\n%s\n\n解析要求：必须按上述视觉风格和故事时代解析角色、场景、道具及分镜，并在资产的description、image_prompt和分镜的prompt_content中体现与之一致的可视化特征。\n只返回合法JSON，顶层只能包含entities和storyboards。entities包含characters、scenes、props，每个资产包含name、description、image_prompt；storyboards每项包含title、script_content、prompt_content、dialogue、asset_names、duration_seconds、source_excerpt。禁止返回摘要或其他顶层字段，不要返回Markdown或解释。", snapshot.ProjectName, snapshot.StyleName, dramaEraName(snapshot.EraType, snapshot.CustomEra), snapshot.EpisodeNo, snapshot.EpisodeName, string(assetsJSON), source)
|
||||
}
|
||||
|
||||
func dramaEraName(eraType, customEra string) string {
|
||||
if eraType == "other" && strings.TrimSpace(customEra) != "" {
|
||||
return strings.TrimSpace(customEra)
|
||||
}
|
||||
if name := map[string]string{
|
||||
"modern_city": "现代都市",
|
||||
"ancient_history": "古代历史",
|
||||
"ancient_fantasy": "古代玄幻",
|
||||
"ancient_xianxia": "古代仙侠",
|
||||
"future_scifi": "未来科幻",
|
||||
}[eraType]; name != "" {
|
||||
return name
|
||||
}
|
||||
return eraType
|
||||
}
|
||||
|
||||
func parseDramaJSON(value string) (dramaParsedResult, error) {
|
||||
value = strings.TrimSpace(value)
|
||||
think := regexp.MustCompile(`(?s)<think>.*?</think>`)
|
||||
value = think.ReplaceAllString(value, "")
|
||||
value = strings.TrimPrefix(value, "```json")
|
||||
value = strings.TrimPrefix(value, "```")
|
||||
value = strings.TrimSuffix(value, "```")
|
||||
value = strings.TrimSpace(value)
|
||||
start := strings.Index(value, "{")
|
||||
end := strings.LastIndex(value, "}")
|
||||
if start < 0 || end <= start {
|
||||
return dramaParsedResult{}, errors.New("文本模型返回内容不是JSON")
|
||||
}
|
||||
var result dramaParsedResult
|
||||
if err := json.Unmarshal([]byte(value[start:end+1]), &result); err != nil {
|
||||
return result, fmt.Errorf("文本模型返回JSON无效: %w", err)
|
||||
}
|
||||
result.Entities.Characters = append(result.Entities.Characters, result.Entities.Roles...)
|
||||
valid := result.Storyboards[:0]
|
||||
for _, item := range result.Storyboards {
|
||||
item.ScriptContent = strings.TrimSpace(firstNonEmpty(item.ScriptContent, item.Content))
|
||||
item.PromptContent = strings.TrimSpace(item.PromptContent)
|
||||
if item.ScriptContent == "" && item.PromptContent == "" {
|
||||
continue
|
||||
}
|
||||
if item.DurationSeconds < 5 || item.DurationSeconds > 15 {
|
||||
item.DurationSeconds = 5
|
||||
}
|
||||
valid = append(valid, item)
|
||||
}
|
||||
result.Storyboards = valid
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func (w *DramaParse) persist(snapshot dramaParseSnapshot, result dramaParsedResult, tokens tokenTotals, contextSnapshot json.RawMessage) error {
|
||||
normalized, _ := json.Marshal(result)
|
||||
usage, _ := json.Marshal(map[string]any{"segments": tokens.Segments})
|
||||
return w.DB.Transaction(func(tx *gorm.DB) error {
|
||||
var task model.DramaParseTask
|
||||
if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Where("id=?", snapshot.Task.ID).Take(&task).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if task.Status == "cancel_requested" {
|
||||
return errDramaCancelled
|
||||
}
|
||||
if task.Status != "running" {
|
||||
return errors.New("任务状态已变化")
|
||||
}
|
||||
assetIDs, err := w.upsertEntities(tx, snapshot, result.Entities)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
now := time.Now()
|
||||
protected := make([]model.EpisodeStoryboard, 0)
|
||||
if err := tx.Where("episode_id=? AND deleted_at IS NULL AND user_edited=true", snapshot.Task.EpisodeID).Find(&protected).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
protectedBySequence := make(map[int]model.EpisodeStoryboard, len(protected))
|
||||
for _, storyboard := range protected {
|
||||
protectedBySequence[storyboard.SequenceNo] = storyboard
|
||||
}
|
||||
if err := tx.Model(&model.EpisodeStoryboard{}).Where("episode_id=? AND deleted_at IS NULL AND user_edited=false", snapshot.Task.EpisodeID).Update("deleted_at", now).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
startMS := int64(0)
|
||||
for index, item := range result.Storyboards {
|
||||
sequence := index + 1
|
||||
if existing, ok := protectedBySequence[sequence]; ok {
|
||||
duration := existing.DurationSeconds
|
||||
if duration < 5 || duration > 15 {
|
||||
duration = 5
|
||||
}
|
||||
endMS := startMS + int64(duration*1000)
|
||||
if err := tx.Model(&model.EpisodeStoryboard{}).Where("id=?", existing.ID).Updates(map[string]any{"start_ms": startMS, "end_ms": endMS, "duration_seconds": duration}).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
startMS = endMS
|
||||
continue
|
||||
}
|
||||
duration := item.DurationSeconds
|
||||
if duration < 5 || duration > 15 {
|
||||
duration = 5
|
||||
}
|
||||
refs := buildAssetRefs(item, assetIDs)
|
||||
refsJSON, _ := json.Marshal(refs)
|
||||
dialogueJSON, _ := json.Marshal(item.Dialogue)
|
||||
endMS := startMS + int64(duration*1000)
|
||||
storyboard := model.EpisodeStoryboard{ID: uuid.New(), EpisodeID: &snapshot.Task.EpisodeID, SequenceNo: sequence, StableKey: uuid.NewString(), StartMS: startMS, EndMS: endMS, DurationSeconds: duration, Title: truncateWorker(item.Title, 160), ScriptContent: item.ScriptContent, PromptContent: firstNonEmpty(item.PromptContent, item.ScriptContent), SourceExcerpt: item.SourceExcerpt, Dialogue: dialogueJSON, AssetRefs: refsJSON, Locked: false, Status: "idle", SourceParseTaskID: &task.ID}
|
||||
if err := tx.Create(&storyboard).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
startMS = endMS
|
||||
}
|
||||
total := tokens.Input + tokens.Output
|
||||
finished := time.Now()
|
||||
if err := tx.Model(&task).Updates(map[string]any{"status": "succeeded", "normalized_result": gorm.Expr("?::jsonb", string(normalized)), "context_snapshot": gorm.Expr("?::jsonb", string(contextSnapshot)), "input_tokens": tokens.Input, "output_tokens": tokens.Output, "total_tokens": total, "token_count_source": tokens.Source, "usage_raw": gorm.Expr("?::jsonb", string(usage)), "actual_points": task.PrepaidPoints, "request_char_count": len([]rune(snapshot.Source.RawContent)), "response_char_count": len([]rune(string(normalized))), "finished_at": finished, "lease_owner": nil, "lease_until": nil}).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Model(&model.ProjectEpisode{}).Where("id=?", snapshot.Task.EpisodeID).Updates(map[string]any{"status": "review", "analysis_message": "剧本解析完成"}).Error
|
||||
})
|
||||
}
|
||||
|
||||
func (w *DramaParse) upsertEntities(tx *gorm.DB, snapshot dramaParseSnapshot, entities parsedEntities) (map[string]uuid.UUID, error) {
|
||||
result := map[string]uuid.UUID{}
|
||||
groups := []struct {
|
||||
Type string
|
||||
Items []parsedEntity
|
||||
}{{"character", entities.Characters}, {"scene", entities.Scenes}, {"prop", entities.Props}}
|
||||
for _, group := range groups {
|
||||
for _, item := range group.Items {
|
||||
rawName := strings.TrimSpace(firstNonEmpty(item.CanonicalName, item.Name))
|
||||
if rawName == "" {
|
||||
continue
|
||||
}
|
||||
imagePrompt := strings.TrimSpace(firstNonEmpty(item.ImagePrompt, item.Description))
|
||||
name := truncateWorker(rawName, 200)
|
||||
aliasValues := append([]string{}, item.Aliases...)
|
||||
if rawName != name {
|
||||
aliasValues = append(aliasValues, rawName)
|
||||
}
|
||||
var existing model.ProjectAsset
|
||||
err := tx.Where("project_id=? AND asset_type=? AND lower(name)=lower(?) AND deleted_at IS NULL", snapshot.Task.ProjectID, group.Type, name).Take(&existing).Error
|
||||
aliases, _ := json.Marshal(uniqueStrings(aliasValues))
|
||||
attributes, _ := json.Marshal(item.Attributes)
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
existing = model.ProjectAsset{ID: uuid.New(), ProjectID: snapshot.Task.ProjectID, SourceEpisodeID: &snapshot.Task.EpisodeID, AssetType: group.Type, Name: name, Description: item.Description, ImagePrompt: imagePrompt, AIDescription: item.Description, Aliases: aliases, Attributes: attributes, Appearances: json.RawMessage("[]"), SourceParseTaskID: &snapshot.Task.ID}
|
||||
if err = tx.Create(&existing).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
} else if err != nil {
|
||||
return nil, err
|
||||
} else {
|
||||
updates := map[string]any{"ai_description": item.Description, "aliases": gorm.Expr("?::jsonb", string(mergeJSONArray(existing.Aliases, aliasValues))), "source_parse_task_id": snapshot.Task.ID}
|
||||
if !existing.UserEdited {
|
||||
updates["description"] = item.Description
|
||||
updates["image_prompt"] = imagePrompt
|
||||
updates["attributes"] = gorm.Expr("?::jsonb", string(attributes))
|
||||
}
|
||||
if err = tx.Model(&existing).Updates(updates).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
result[group.Type+":"+strings.ToLower(rawName)] = existing.ID
|
||||
result[group.Type+":"+strings.ToLower(name)] = existing.ID
|
||||
for _, alias := range aliasValues {
|
||||
result[group.Type+":"+strings.ToLower(strings.TrimSpace(alias))] = existing.ID
|
||||
}
|
||||
}
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func buildAssetRefs(item parsedStoryboard, ids map[string]uuid.UUID) []map[string]any {
|
||||
refs := make([]map[string]any, 0)
|
||||
groups := []struct {
|
||||
Type string
|
||||
Names []string
|
||||
}{{"character", item.AssetNames.Characters}, {"scene", item.AssetNames.Scenes}, {"prop", item.AssetNames.Props}}
|
||||
seen := map[uuid.UUID]bool{}
|
||||
for _, group := range groups {
|
||||
for _, name := range group.Names {
|
||||
if id, ok := ids[group.Type+":"+strings.ToLower(strings.TrimSpace(name))]; ok && !seen[id] {
|
||||
refs = append(refs, map[string]any{"id": id, "name": name, "type": group.Type})
|
||||
seen[id] = true
|
||||
}
|
||||
}
|
||||
}
|
||||
return refs
|
||||
}
|
||||
|
||||
func (w *DramaParse) fail(task model.DramaParseTask, code, message string, tokenUsage ...tokenTotals) error {
|
||||
now := time.Now()
|
||||
return w.DB.Transaction(func(tx *gorm.DB) error {
|
||||
refunded, err := billing.RefundTextCalls(tx, task.UserID, task.ID.String(), "剧本解析失败返还")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
updates := map[string]any{"status": "failed", "actual_points": "0.00", "error_code": code, "error_message": truncateWorker(message, 4000), "finished_at": now, "lease_owner": nil, "lease_until": nil}
|
||||
if refunded {
|
||||
updates["cost_refunded"] = true
|
||||
}
|
||||
if len(tokenUsage) > 0 {
|
||||
addTokenUpdates(updates, tokenUsage[0])
|
||||
}
|
||||
if err := tx.Model(&model.DramaParseTask{}).Where("id=? AND status<>'cancelled'", task.ID).Updates(updates).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Model(&model.ProjectEpisode{}).Where("id=?", task.EpisodeID).Updates(map[string]any{"status": "failed", "analysis_message": truncateWorker(message, 1000)}).Error
|
||||
})
|
||||
}
|
||||
func (w *DramaParse) cancel(task model.DramaParseTask, tokenUsage ...tokenTotals) error {
|
||||
now := time.Now()
|
||||
return w.DB.Transaction(func(tx *gorm.DB) error {
|
||||
amount, err := billing.CommitTextReserves(tx, task.UserID, task.ID.String(), "剧本解析取消")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
updates := map[string]any{"status": "cancelled", "actual_points": amount, "prepaid_points": amount, "finished_at": now, "lease_owner": nil, "lease_until": nil, "error_code": "user_cancelled", "error_message": "用户主动取消，费用不退"}
|
||||
if len(tokenUsage) > 0 {
|
||||
addTokenUpdates(updates, tokenUsage[0])
|
||||
}
|
||||
if err := tx.Model(&model.DramaParseTask{}).Where("id=?", task.ID).Updates(updates).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Model(&model.ProjectEpisode{}).Where("id=?", task.EpisodeID).Updates(map[string]any{"status": "draft", "analysis_message": "解析已取消"}).Error
|
||||
})
|
||||
}
|
||||
|
||||
func addTokenUpdates(updates map[string]any, tokens tokenTotals) {
|
||||
if tokens.Input == 0 && tokens.Output == 0 {
|
||||
return
|
||||
}
|
||||
usage, _ := json.Marshal(map[string]any{"segments": tokens.Segments})
|
||||
updates["input_tokens"] = tokens.Input
|
||||
updates["output_tokens"] = tokens.Output
|
||||
updates["total_tokens"] = tokens.Input + tokens.Output
|
||||
updates["token_count_source"] = tokens.Source
|
||||
updates["usage_raw"] = gorm.Expr("?::jsonb", string(usage))
|
||||
}
|
||||
|
||||
func estimateTextTokens(value string) int64 {
|
||||
return billing.EstimateTextTokens(value)
|
||||
}
|
||||
func uniqueStrings(values []string) []string {
|
||||
seen := map[string]bool{}
|
||||
out := make([]string, 0, len(values))
|
||||
for _, value := range values {
|
||||
value = strings.TrimSpace(value)
|
||||
key := strings.ToLower(value)
|
||||
if value != "" && !seen[key] {
|
||||
seen[key] = true
|
||||
out = append(out, value)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
func mergeJSONArray(raw json.RawMessage, values []string) json.RawMessage {
|
||||
var existing []string
|
||||
_ = json.Unmarshal(raw, &existing)
|
||||
encoded, _ := json.Marshal(uniqueStrings(append(existing, values...)))
|
||||
return encoded
|
||||
}
|
||||
func firstNonEmpty(values ...string) string {
|
||||
for _, value := range values {
|
||||
if strings.TrimSpace(value) != "" {
|
||||
return strings.TrimSpace(value)
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
func normalizeCompare(value string) string {
|
||||
return strings.Join(strings.Fields(strings.ToLower(value)), "")
|
||||
}
|
||||
func truncateWorker(value string, limit int) string {
|
||||
r := []rune(value)
|
||||
if len(r) > limit {
|
||||
return string(r[:limit])
|
||||
}
|
||||
return value
|
||||
}
|
||||
func classifyDramaError(err error) string {
|
||||
var httpErr *provider.HTTPError
|
||||
if errors.As(err, &httpErr) {
|
||||
return "upstream_failed"
|
||||
}
|
||||
if errors.Is(err, context.DeadlineExceeded) {
|
||||
return "upstream_timeout"
|
||||
}
|
||||
return "parse_failed"
|
||||
}
|
||||
@@ -0,0 +1,95 @@
|
||||
package worker
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// TestClassifyDramaErrorTreatsTimeoutAsUpstreamFailure 验证超时不会再进入提交未知状态。
|
||||
func TestClassifyDramaErrorTreatsTimeoutAsUpstreamFailure(t *testing.T) {
|
||||
if code := classifyDramaError(context.DeadlineExceeded); code != "upstream_timeout" {
|
||||
t.Fatalf("timeout code = %q, want upstream_timeout", code)
|
||||
}
|
||||
}
|
||||
|
||||
// TestDramaModelRetryLimit 验证剧本解析调用的自动重试次数。
|
||||
func TestDramaModelRetryLimit(t *testing.T) {
|
||||
if maxDramaModelAttempts != 5 {
|
||||
t.Fatalf("retry attempts = %d, want 5", maxDramaModelAttempts)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDramaHTTPClientAllowsLongTextResponses(t *testing.T) {
|
||||
sharedTransport := &http.Transport{TLSHandshakeTimeout: 10 * time.Second, ResponseHeaderTimeout: time.Minute}
|
||||
shared := &http.Client{Transport: sharedTransport}
|
||||
client := dramaHTTPClient(shared)
|
||||
transport, ok := client.Transport.(*http.Transport)
|
||||
if !ok {
|
||||
t.Fatal("drama client transport was not preserved")
|
||||
}
|
||||
if transport.TLSHandshakeTimeout != 30*time.Second || transport.ResponseHeaderTimeout != 10*time.Minute {
|
||||
t.Fatalf("unexpected drama timeouts: TLS=%s response=%s", transport.TLSHandshakeTimeout, transport.ResponseHeaderTimeout)
|
||||
}
|
||||
if sharedTransport.TLSHandshakeTimeout != 10*time.Second || sharedTransport.ResponseHeaderTimeout != time.Minute {
|
||||
t.Fatal("shared AI transport was modified")
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseDramaJSONNormalizesStoryboard(t *testing.T) {
|
||||
result, err := parseDramaJSON("<think>hidden</think>```json\n{\"entities\":{\"roles\":[{\"name\":\"林夏\",\"image_prompt\":\"二十五岁女性，黑色长发\"}]},\"storyboards\":[{\"content\":\"走进房间\",\"duration_seconds\":99}]}\n```")
|
||||
if err != nil {
|
||||
t.Fatalf("parseDramaJSON returned error: %v", err)
|
||||
}
|
||||
if len(result.Entities.Characters) != 1 || len(result.Storyboards) != 1 {
|
||||
t.Fatalf("unexpected normalized result: %#v", result)
|
||||
}
|
||||
if result.Storyboards[0].ScriptContent != "走进房间" || result.Storyboards[0].DurationSeconds != 5 {
|
||||
t.Fatalf("storyboard was not normalized: %#v", result.Storyboards[0])
|
||||
}
|
||||
if result.Entities.Characters[0].ImagePrompt != "二十五岁女性，黑色长发" {
|
||||
t.Fatalf("asset image prompt was not preserved: %#v", result.Entities.Characters[0])
|
||||
}
|
||||
}
|
||||
|
||||
func TestEstimateTextTokens(t *testing.T) {
|
||||
if estimateTextTokens("中文ABCD") != 3 {
|
||||
t.Fatalf("unexpected token estimate")
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildDramaUserPromptUsesOnlyCurrentContext(t *testing.T) {
|
||||
snapshot := dramaParseSnapshot{
|
||||
ProjectName: "测试项目",
|
||||
StyleName: "国风水墨",
|
||||
EraType: "ancient_xianxia",
|
||||
EpisodeNo: 4,
|
||||
EpisodeName: "再入山门",
|
||||
}
|
||||
prompt := buildDramaUserPrompt(snapshot, "当前集正文", []map[string]any{{"name": "林夏"}})
|
||||
for _, expected := range []string{
|
||||
"视觉风格：国风水墨",
|
||||
"故事时代：古代仙侠",
|
||||
"项目资产：",
|
||||
"林夏",
|
||||
"当前原文：\n当前集正文",
|
||||
"按上述视觉风格和故事时代解析角色、场景、道具及分镜",
|
||||
} {
|
||||
if !strings.Contains(prompt, expected) {
|
||||
t.Fatalf("prompt does not contain %q: %s", expected, prompt)
|
||||
}
|
||||
}
|
||||
for _, forbidden := range []string{"continuity", "解析分段", "前序"} {
|
||||
if strings.Contains(prompt, forbidden) {
|
||||
t.Fatalf("prompt contains forbidden content %q: %s", forbidden, prompt)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestDramaEraNameUsesCustomEra(t *testing.T) {
|
||||
if got := dramaEraName("other", " 民国时期 "); got != "民国时期" {
|
||||
t.Fatalf("unexpected custom era: %q", got)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,850 @@
|
||||
package worker
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"math/rand"
|
||||
"net"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"juhe-factory/api/internal/billing"
|
||||
mediakey "juhe-factory/api/internal/media"
|
||||
"juhe-factory/api/internal/model"
|
||||
"juhe-factory/api/internal/provider"
|
||||
"juhe-factory/api/internal/provider/apimart"
|
||||
queuepkg "juhe-factory/api/internal/queue"
|
||||
"juhe-factory/api/internal/security"
|
||||
"juhe-factory/api/internal/storage"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/hibiken/asynq"
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/clause"
|
||||
)
|
||||
|
||||
var completedStatuses = map[string]bool{"completed": true, "complete": true, "succeeded": true, "success": true, "done": true}
|
||||
var failedStatuses = map[string]bool{"failed": true, "failure": true, "error": true, "cancelled": true, "canceled": true}
|
||||
|
||||
const maxSubmitAttempts = 5
|
||||
|
||||
type Generation struct {
|
||||
DB *gorm.DB
|
||||
Queue *asynq.Client
|
||||
COS *storage.COS
|
||||
Encryptor *security.Encryptor
|
||||
Provider apimart.Client
|
||||
Seedance2 apimart.Seedance2
|
||||
DownloadHTTP *http.Client
|
||||
PollInterval time.Duration
|
||||
WorkerID string
|
||||
}
|
||||
|
||||
type channelSnapshot struct {
|
||||
ID uuid.UUID
|
||||
BaseURL string
|
||||
APIKeyCiphertext string
|
||||
MaxConcurrency int
|
||||
MaxUserConcurrency int
|
||||
Enabled bool
|
||||
}
|
||||
|
||||
type submitSnapshot struct {
|
||||
Task model.GenerationTask
|
||||
Channel channelSnapshot
|
||||
Model string
|
||||
}
|
||||
|
||||
func NewGeneration(db *gorm.DB, queue *asynq.Client, cos *storage.COS, encryptor *security.Encryptor, client *http.Client, pollInterval time.Duration) *Generation {
|
||||
return &Generation{DB: db, Queue: queue, COS: cos, Encryptor: encryptor, Provider: apimart.NewClient(client), Seedance2: apimart.NewSeedance2(client), DownloadHTTP: client, PollInterval: pollInterval, WorkerID: uuid.NewString()}
|
||||
}
|
||||
|
||||
func (w *Generation) Register(mux *asynq.ServeMux) {
|
||||
mux.HandleFunc(queuepkg.TypeDispatchChannel, w.dispatchChannel)
|
||||
mux.HandleFunc(queuepkg.TypeSubmitTask, w.submitTask)
|
||||
mux.HandleFunc(queuepkg.TypePollTask, w.pollTask)
|
||||
mux.HandleFunc(queuepkg.TypeDownloadTask, w.downloadTask)
|
||||
}
|
||||
|
||||
func decodeID(task *asynq.Task) (uuid.UUID, error) {
|
||||
var payload queuepkg.IDPayload
|
||||
if err := json.Unmarshal(task.Payload(), &payload); err != nil {
|
||||
return uuid.Nil, err
|
||||
}
|
||||
if payload.ID == uuid.Nil {
|
||||
return uuid.Nil, errors.New("任务标识无效")
|
||||
}
|
||||
return payload.ID, nil
|
||||
}
|
||||
|
||||
func (w *Generation) dispatchChannel(ctx context.Context, raw *asynq.Task) error {
|
||||
channelID, err := decodeID(raw)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
selected := make([]uuid.UUID, 0)
|
||||
err = w.DB.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
var channel channelSnapshot
|
||||
if err := tx.Raw(`SELECT id,base_url,api_key_ciphertext,max_concurrency,max_user_concurrency,enabled
|
||||
FROM channels WHERE id=? AND deleted_at IS NULL FOR UPDATE`, channelID).Scan(&channel).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if channel.ID == uuid.Nil || !channel.Enabled {
|
||||
return nil
|
||||
}
|
||||
var active int64
|
||||
if err := tx.Model(&model.GenerationTask{}).Where("channel_id=? AND slot_reserved=true AND slot_released_at IS NULL", channelID).Count(&active).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
available := channel.MaxConcurrency - int(active)
|
||||
if available <= 0 {
|
||||
return nil
|
||||
}
|
||||
activeByUser := map[uuid.UUID]int{}
|
||||
rows := make([]struct {
|
||||
UserID uuid.UUID
|
||||
Count int
|
||||
}, 0)
|
||||
if err := tx.Model(&model.GenerationTask{}).Select("user_id,count(*) AS count").Where("channel_id=? AND slot_reserved=true AND slot_released_at IS NULL", channelID).Group("user_id").Scan(&rows).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
for _, row := range rows {
|
||||
activeByUser[row.UserID] = row.Count
|
||||
}
|
||||
var candidates []model.GenerationTask
|
||||
now := time.Now()
|
||||
if err := tx.Raw(`WITH active_users AS (
|
||||
SELECT user_id,count(*) AS active_count FROM generation_tasks
|
||||
WHERE channel_id=? AND slot_reserved=true AND slot_released_at IS NULL GROUP BY user_id
|
||||
), ranked AS (
|
||||
SELECT pending.id,row_number() OVER (PARTITION BY pending.user_id ORDER BY pending.created_at,pending.id) AS user_rank,
|
||||
(?-coalesce(active.active_count,0)) AS allowance
|
||||
FROM generation_tasks pending LEFT JOIN active_users active ON active.user_id=pending.user_id
|
||||
WHERE pending.channel_id=? AND pending.status='pending_submission' AND
|
||||
(pending.project_id IS NOT NULL OR (pending.task_type='image_generation' AND pending.input_data->>'product_image'='true'))
|
||||
AND (pending.next_submit_at IS NULL OR pending.next_submit_at<=?)
|
||||
AND (pending.lease_until IS NULL OR pending.lease_until<?)
|
||||
)
|
||||
SELECT task.* FROM generation_tasks task JOIN ranked ON ranked.id=task.id
|
||||
WHERE ranked.allowance>0 AND ranked.user_rank<=ranked.allowance
|
||||
ORDER BY ranked.user_rank,task.created_at,task.id LIMIT ? FOR UPDATE OF task SKIP LOCKED`,
|
||||
channelID, channel.MaxUserConcurrency, channelID, now, now, available).Scan(&candidates).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if len(candidates) == 0 {
|
||||
return nil
|
||||
}
|
||||
for _, candidate := range selectFairCandidates(candidates, activeByUser, available, channel.MaxUserConcurrency) {
|
||||
leaseUntil := now.Add(2 * time.Minute)
|
||||
result := tx.Model(&model.GenerationTask{}).Where("id=? AND status='pending_submission'", candidate.ID).Updates(map[string]any{
|
||||
"status": "submitting", "slot_reserved": true, "slot_reserved_at": now,
|
||||
"slot_released_at": nil, "lease_owner": w.WorkerID, "lease_until": leaseUntil,
|
||||
"submit_attempts": gorm.Expr("submit_attempts+1"),
|
||||
})
|
||||
if result.Error != nil {
|
||||
return result.Error
|
||||
}
|
||||
if result.RowsAffected == 1 {
|
||||
selected = append(selected, candidate.ID)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
for _, taskID := range selected {
|
||||
if err := queuepkg.EnqueueID(w.Queue, queuepkg.TypeSubmitTask, taskID, 0); err != nil {
|
||||
slog.Error("提交任务入队失败", "task_id", taskID, "error", err)
|
||||
w.returnUnenqueuedTask(ctx, taskID, err.Error())
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func selectFairCandidates(candidates []model.GenerationTask, activeByUser map[uuid.UUID]int, available, maxPerUser int) []model.GenerationTask {
|
||||
if available <= 0 || maxPerUser <= 0 {
|
||||
return nil
|
||||
}
|
||||
groups := map[uuid.UUID][]model.GenerationTask{}
|
||||
order := make([]uuid.UUID, 0)
|
||||
for _, candidate := range candidates {
|
||||
if _, exists := groups[candidate.UserID]; !exists {
|
||||
order = append(order, candidate.UserID)
|
||||
}
|
||||
groups[candidate.UserID] = append(groups[candidate.UserID], candidate)
|
||||
}
|
||||
selected := make([]model.GenerationTask, 0, available)
|
||||
for available > 0 {
|
||||
progressed := false
|
||||
for _, userID := range order {
|
||||
if available == 0 {
|
||||
break
|
||||
}
|
||||
if activeByUser[userID] >= maxPerUser || len(groups[userID]) == 0 {
|
||||
continue
|
||||
}
|
||||
selected = append(selected, groups[userID][0])
|
||||
groups[userID] = groups[userID][1:]
|
||||
activeByUser[userID]++
|
||||
available--
|
||||
progressed = true
|
||||
}
|
||||
if !progressed {
|
||||
break
|
||||
}
|
||||
}
|
||||
return selected
|
||||
}
|
||||
|
||||
func (w *Generation) loadSubmitSnapshot(ctx context.Context, taskID uuid.UUID) (submitSnapshot, error) {
|
||||
var result submitSnapshot
|
||||
if err := w.DB.WithContext(ctx).Where("id=?", taskID).First(&result.Task).Error; err != nil {
|
||||
return result, err
|
||||
}
|
||||
if result.Task.ChannelID == nil || result.Task.ModelID == nil {
|
||||
return result, errors.New("任务未绑定模型渠道")
|
||||
}
|
||||
if err := w.DB.WithContext(ctx).Raw(`SELECT id,base_url,api_key_ciphertext,max_concurrency,max_user_concurrency,enabled
|
||||
FROM channels WHERE id=? AND deleted_at IS NULL`, *result.Task.ChannelID).Scan(&result.Channel).Error; err != nil {
|
||||
return result, err
|
||||
}
|
||||
if !result.Channel.Enabled {
|
||||
return result, errors.New("任务渠道已停用")
|
||||
}
|
||||
if err := w.DB.WithContext(ctx).Table("models").Select("name").Where("id=? AND enabled=true AND deleted_at IS NULL", *result.Task.ModelID).Scan(&result.Model).Error; err != nil {
|
||||
return result, err
|
||||
}
|
||||
if result.Model == "" {
|
||||
return result, errors.New("任务模型不可用")
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func (w *Generation) submitTask(ctx context.Context, raw *asynq.Task) error {
|
||||
taskID, err := decodeID(raw)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
snapshot, err := w.loadSubmitSnapshot(ctx, taskID)
|
||||
if err != nil {
|
||||
// 用户可以删除卡死的活动任务；队列收到遗留消息时直接结束，避免无意义重试。
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
return nil
|
||||
}
|
||||
return w.failAndRefund(ctx, taskID, "submit_configuration_error", err.Error())
|
||||
}
|
||||
if snapshot.Task.Status != "submitting" {
|
||||
if snapshot.Task.Status == "pending_submission" && snapshot.Task.ChannelID != nil {
|
||||
return queuepkg.EnqueueID(w.Queue, queuepkg.TypeDispatchChannel, *snapshot.Task.ChannelID, 0)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
apiKey, err := w.Encryptor.Decrypt(snapshot.Channel.APIKeyCiphertext)
|
||||
if err != nil {
|
||||
return w.failAndRefund(ctx, taskID, "channel_secret_error", "渠道密钥不可用")
|
||||
}
|
||||
payload, err := buildProviderPayload(snapshot)
|
||||
if err != nil {
|
||||
return w.failAndRefund(ctx, taskID, "invalid_payload", err.Error())
|
||||
}
|
||||
requestCtx, cancel := context.WithTimeout(ctx, 90*time.Second)
|
||||
defer cancel()
|
||||
var result provider.SubmitResult
|
||||
if snapshot.Task.TaskType == "video_generation" {
|
||||
if payloadJSON, marshalErr := json.Marshal(payload); marshalErr == nil {
|
||||
fmt.Printf("console.log video_generation payload: %s\n", payloadJSON)
|
||||
}
|
||||
result, err = w.Seedance2.Submit(requestCtx, snapshot.Channel.BaseURL, apiKey, snapshot.Task.RequestID, payload)
|
||||
} else {
|
||||
result, err = w.Provider.SubmitImage(requestCtx, snapshot.Channel.BaseURL, apiKey, snapshot.Task.RequestID, payload)
|
||||
}
|
||||
if err != nil {
|
||||
var httpErr *provider.HTTPError
|
||||
if errors.As(err, &httpErr) {
|
||||
if httpErr.StatusCode == http.StatusTooManyRequests || httpErr.StatusCode >= 500 {
|
||||
if !canRetrySubmission(snapshot.Task.SubmitAttempts) {
|
||||
return w.failAndRefund(ctx, taskID, "upstream_unavailable", err.Error())
|
||||
}
|
||||
return w.requeueRejected(ctx, snapshot.Task, err.Error(), retryDelay(snapshot.Task.SubmitAttempts, httpErr.RetryAfter))
|
||||
}
|
||||
return w.failAndRefund(ctx, taskID, "upstream_rejected", err.Error())
|
||||
}
|
||||
// 传输异常统一按提交失败重试；达到上限后结束任务并返还预扣积分。
|
||||
if !canRetrySubmission(snapshot.Task.SubmitAttempts) {
|
||||
return w.failAndRefund(ctx, taskID, "upstream_unavailable", err.Error())
|
||||
}
|
||||
return w.requeueRejected(ctx, snapshot.Task, err.Error(), retryDelay(snapshot.Task.SubmitAttempts, ""))
|
||||
}
|
||||
now := time.Now()
|
||||
if result.URL != "" {
|
||||
if err := w.markResultReady(ctx, taskID, result.URL); err != nil {
|
||||
return err
|
||||
}
|
||||
_ = queuepkg.EnqueueID(w.Queue, queuepkg.TypeDownloadTask, taskID, 0)
|
||||
w.enqueueDispatch(snapshot.Task.ChannelID)
|
||||
return nil
|
||||
}
|
||||
nextPoll := now.Add(w.PollInterval + jitter(3*time.Second))
|
||||
update := w.DB.WithContext(ctx).Model(&model.GenerationTask{}).Where("id=? AND status='submitting'", taskID).Updates(map[string]any{
|
||||
"status": "submitted", "upstream_task_id": result.TaskID, "submitted_at": now,
|
||||
"next_poll_at": nextPoll, "lease_owner": nil, "lease_until": nil, "error_code": nil, "error_message": nil,
|
||||
})
|
||||
if update.Error != nil {
|
||||
return update.Error
|
||||
}
|
||||
if update.RowsAffected == 1 {
|
||||
_ = queuepkg.EnqueueID(w.Queue, queuepkg.TypePollTask, taskID, time.Until(nextPoll))
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func buildProviderPayload(snapshot submitSnapshot) (map[string]any, error) {
|
||||
if snapshot.Task.TaskType == "video_generation" {
|
||||
return apimart.BuildSeedance2Payload(snapshot.Model, snapshot.Task.InputData)
|
||||
}
|
||||
var input map[string]any
|
||||
if err := json.Unmarshal(snapshot.Task.InputData, &input); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
payload := map[string]any{"model": snapshot.Model, "prompt": strings.TrimSpace(fmt.Sprint(input["prompt"]))}
|
||||
if payload["prompt"] == "" {
|
||||
return nil, errors.New("生成提示词不能为空")
|
||||
}
|
||||
if snapshot.Task.TaskType == "image_generation" {
|
||||
// 图片画幅由各业务模块在入队前规范化，Worker 只负责透传统一字段。
|
||||
payload["size"] = stringOr(input["size"], "16:9")
|
||||
payload["resolution"] = stringOr(input["resolution"], "1k")
|
||||
} else {
|
||||
payload["size"] = stringOr(input["resolution"], "1k")
|
||||
}
|
||||
if n, ok := input["n"]; ok {
|
||||
payload["n"] = n
|
||||
} else {
|
||||
payload["n"] = 1
|
||||
}
|
||||
if urls, ok := input["image_urls"]; ok {
|
||||
payload["image_urls"] = urls
|
||||
}
|
||||
return payload, nil
|
||||
}
|
||||
|
||||
func (w *Generation) pollTask(ctx context.Context, raw *asynq.Task) error {
|
||||
taskID, err := decodeID(raw)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
snapshot, err := w.loadSubmitSnapshot(ctx, taskID)
|
||||
if err != nil || (snapshot.Task.Status != "submitted" && snapshot.Task.Status != "processing" && snapshot.Task.Status != "cancel_requested") {
|
||||
return nil
|
||||
}
|
||||
apiKey, err := w.Encryptor.Decrypt(snapshot.Channel.APIKeyCiphertext)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
requestCtx, cancel := context.WithTimeout(ctx, 30*time.Second)
|
||||
var result provider.PollResult
|
||||
var pollErr error
|
||||
if snapshot.Task.TaskType == "video_generation" {
|
||||
result, pollErr = w.Seedance2.Poll(requestCtx, snapshot.Channel.BaseURL, apiKey, snapshot.Task.UpstreamTaskID)
|
||||
} else {
|
||||
result, pollErr = w.Provider.PollImage(requestCtx, snapshot.Channel.BaseURL, apiKey, snapshot.Task.UpstreamTaskID)
|
||||
}
|
||||
cancel()
|
||||
if pollErr != nil {
|
||||
if snapshot.Task.Status == "cancel_requested" && snapshot.Task.PollAttempts >= 10 {
|
||||
if err := w.finishCancellation(ctx, taskID, true, "上游任务状态长时间不可用，已自动结束取消"); err != nil {
|
||||
return err
|
||||
}
|
||||
w.enqueueDispatch(snapshot.Task.ChannelID)
|
||||
return nil
|
||||
}
|
||||
next := time.Now().Add(w.PollInterval + jitter(4*time.Second))
|
||||
w.DB.WithContext(ctx).Model(&model.GenerationTask{}).Where("id=? AND status IN ('submitted','processing','cancel_requested')", taskID).Updates(map[string]any{
|
||||
"poll_attempts": gorm.Expr("poll_attempts+1"), "next_poll_at": next, "error_message": pollErr.Error(),
|
||||
})
|
||||
_ = queuepkg.EnqueueID(w.Queue, queuepkg.TypePollTask, taskID, time.Until(next))
|
||||
return nil
|
||||
}
|
||||
if completedStatuses[result.Status] {
|
||||
if snapshot.Task.Status == "cancel_requested" {
|
||||
if err := w.finishCancellation(ctx, taskID, false, "上游任务已结束"); err != nil {
|
||||
return err
|
||||
}
|
||||
w.enqueueDispatch(snapshot.Task.ChannelID)
|
||||
return nil
|
||||
}
|
||||
if result.URL == "" {
|
||||
next := time.Now().Add(w.PollInterval)
|
||||
w.DB.WithContext(ctx).Model(&model.GenerationTask{}).Where("id=?", taskID).Updates(map[string]any{"next_poll_at": next, "error_message": "中转站已完成但结果地址尚未就绪"})
|
||||
_ = queuepkg.EnqueueID(w.Queue, queuepkg.TypePollTask, taskID, time.Until(next))
|
||||
return nil
|
||||
}
|
||||
if err := w.markResultReady(ctx, taskID, result.URL); err != nil {
|
||||
return err
|
||||
}
|
||||
_ = queuepkg.EnqueueID(w.Queue, queuepkg.TypeDownloadTask, taskID, 0)
|
||||
w.enqueueDispatch(snapshot.Task.ChannelID)
|
||||
return nil
|
||||
}
|
||||
if failedStatuses[result.Status] {
|
||||
if snapshot.Task.Status == "cancel_requested" {
|
||||
if err := w.finishCancellation(ctx, taskID, true, "上游任务已取消"); err != nil {
|
||||
return err
|
||||
}
|
||||
} else {
|
||||
if err := w.failAndRefund(ctx, taskID, "upstream_failed", result.Error); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
w.enqueueDispatch(snapshot.Task.ChannelID)
|
||||
return nil
|
||||
}
|
||||
next := time.Now().Add(w.PollInterval + jitter(4*time.Second))
|
||||
updates := map[string]any{"poll_attempts": gorm.Expr("poll_attempts+1"), "next_poll_at": next, "error_code": nil, "error_message": nil}
|
||||
if snapshot.Task.Status != "cancel_requested" {
|
||||
updates["status"] = "processing"
|
||||
}
|
||||
w.DB.WithContext(ctx).Model(&model.GenerationTask{}).Where("id=? AND status IN ('submitted','processing','cancel_requested')", taskID).Updates(updates)
|
||||
_ = queuepkg.EnqueueID(w.Queue, queuepkg.TypePollTask, taskID, time.Until(next))
|
||||
return nil
|
||||
}
|
||||
|
||||
func (w *Generation) markResultReady(ctx context.Context, taskID uuid.UUID, resultURL string) error {
|
||||
return w.DB.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
var task model.GenerationTask
|
||||
if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Where("id=?", taskID).First(&task).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if task.Status == "cancelled" || task.Status == "cancel_requested" {
|
||||
return nil
|
||||
}
|
||||
updates := map[string]any{"status": "result_ready", "upstream_result_url": resultURL, "next_poll_at": nil, "lease_owner": nil, "lease_until": nil}
|
||||
if task.SlotReserved && task.SlotReleasedAt == nil {
|
||||
now := time.Now()
|
||||
updates["slot_released_at"] = now
|
||||
}
|
||||
return tx.Model(&task).Updates(updates).Error
|
||||
})
|
||||
}
|
||||
|
||||
func (w *Generation) downloadTask(ctx context.Context, raw *asynq.Task) error {
|
||||
taskID, err := decodeID(raw)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
var row struct {
|
||||
model.GenerationTask
|
||||
SequenceNo int
|
||||
EpisodeNo int
|
||||
}
|
||||
if err := w.DB.WithContext(ctx).Table("generation_tasks gt").Select("gt.*,coalesce(sb.sequence_no,0) AS sequence_no,coalesce(e.episode_no,0) AS episode_no").Joins("LEFT JOIN episode_storyboards sb ON sb.id=gt.storyboard_id").Joins("LEFT JOIN project_episodes e ON e.id=gt.episode_id").Where("gt.id=?", taskID).Scan(&row).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if row.Status != "result_ready" && row.Status != "downloading" {
|
||||
return nil
|
||||
}
|
||||
if row.UpstreamResultURL == "" || w.COS == nil {
|
||||
return w.failWithoutRefund(ctx, taskID, "result_storage_error", "结果存储信息不完整")
|
||||
}
|
||||
w.DB.WithContext(ctx).Model(&model.GenerationTask{}).Where("id=? AND status='result_ready'", taskID).Update("status", "downloading")
|
||||
mediaID := uuid.New()
|
||||
outputType, fallbackType, maxBytes := "video", "video/mp4", w.COS.MaxVideoBytes()
|
||||
displayName := fmt.Sprintf("分镜%04d生成结果", row.SequenceNo)
|
||||
key := ""
|
||||
if row.TaskType == "image_generation" {
|
||||
outputType, fallbackType, maxBytes = "image", "image/png", w.COS.MaxImageBytes()
|
||||
var input map[string]any
|
||||
_ = json.Unmarshal(row.InputData, &input)
|
||||
if row.ProjectID == nil {
|
||||
key = mediakey.StandaloneTaskOutput(row.UserID, taskID, "image", fallbackType)
|
||||
displayName = "独立生成图"
|
||||
} else {
|
||||
if fmt.Sprint(input["target_type"]) == "storyboard" && row.StoryboardID != nil {
|
||||
if row.EpisodeID != nil {
|
||||
key = mediakey.TaskOutput(*row.ProjectID, *row.EpisodeID, row.SequenceNo, taskID, mediaID, "image", fallbackType)
|
||||
} else {
|
||||
key = mediakey.ProjectTaskOutput(*row.ProjectID, row.SequenceNo, taskID, mediaID, "image", fallbackType)
|
||||
}
|
||||
displayName = fmt.Sprintf("分镜%04d生成图", row.SequenceNo)
|
||||
} else {
|
||||
assetID, parseErr := uuid.Parse(fmt.Sprint(input["asset_id"]))
|
||||
if parseErr != nil {
|
||||
return w.failWithoutRefund(ctx, taskID, "result_storage_error", "资产生图任务缺少资产标识")
|
||||
}
|
||||
var asset struct{ AssetType, Name string }
|
||||
if err := w.DB.WithContext(ctx).Table("project_assets").Select("asset_type,name").Where("id=? AND project_id=? AND deleted_at IS NULL", assetID, *row.ProjectID).Take(&asset).Error; err != nil {
|
||||
return w.failWithoutRefund(ctx, taskID, "result_storage_error", "资产生图任务绑定的资产不存在")
|
||||
}
|
||||
key = mediakey.AssetOutput(*row.ProjectID, asset.AssetType, assetID, taskID, mediaID, fallbackType)
|
||||
displayName = asset.Name + "生成图"
|
||||
}
|
||||
}
|
||||
} else {
|
||||
if row.ProjectID != nil {
|
||||
if row.EpisodeID != nil {
|
||||
key = mediakey.TaskOutput(*row.ProjectID, *row.EpisodeID, row.SequenceNo, taskID, mediaID, outputType, fallbackType)
|
||||
} else {
|
||||
key = mediakey.ProjectTaskOutput(*row.ProjectID, row.SequenceNo, taskID, mediaID, outputType, fallbackType)
|
||||
}
|
||||
} else {
|
||||
key = mediakey.StandaloneTaskOutput(row.UserID, taskID, "video", fallbackType)
|
||||
}
|
||||
}
|
||||
requestCtx, cancel := context.WithTimeout(ctx, 5*time.Minute)
|
||||
url, storedKey, contentType, size, downloadErr := w.COS.PutFromURL(requestCtx, w.DownloadHTTP, row.UpstreamResultURL, key, fallbackType, maxBytes)
|
||||
cancel()
|
||||
if downloadErr != nil {
|
||||
attempts := row.DownloadAttempts + 1
|
||||
if attempts < 3 {
|
||||
w.DB.WithContext(ctx).Model(&model.GenerationTask{}).Where("id=?", taskID).Updates(map[string]any{"status": "result_ready", "download_attempts": attempts, "error_message": downloadErr.Error()})
|
||||
_ = queuepkg.EnqueueID(w.Queue, queuepkg.TypeDownloadTask, taskID, time.Duration(attempts*10)*time.Second)
|
||||
return nil
|
||||
}
|
||||
return w.failWithoutRefund(ctx, taskID, "result_download_failed", downloadErr.Error())
|
||||
}
|
||||
key = storedKey
|
||||
if (outputType == "image" && !strings.HasPrefix(contentType, "image/")) || (outputType == "video" && !strings.HasPrefix(contentType, "video/")) {
|
||||
_ = w.COS.Delete(ctx, key)
|
||||
return w.failWithoutRefund(ctx, taskID, "invalid_result_media", "中转站结果媒体类型与任务不匹配")
|
||||
}
|
||||
ownerID := row.UserID
|
||||
mediaAsset := model.MediaAsset{ID: mediaID, OwnerUserID: &ownerID, StorageProvider: "cos", ObjectKey: key, PublicURL: url, DisplayName: displayName, MimeType: contentType, SizeBytes: size}
|
||||
metadata, _ := json.Marshal(map[string]any{"source": "apimart", "stored_at": time.Now()})
|
||||
output := model.GenerationOutput{TaskID: taskID, MediaAssetID: mediaID, OutputType: outputType, SequenceNo: 1, Metadata: metadata}
|
||||
reusedOutput := false
|
||||
cancelledDuringDownload := false
|
||||
err = w.DB.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
var lockedTask model.GenerationTask
|
||||
if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Select("id", "status").Where("id=?", taskID).Take(&lockedTask).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if lockedTask.Status == "cancelled" || lockedTask.Status == "cancel_requested" {
|
||||
cancelledDuringDownload = true
|
||||
return nil
|
||||
}
|
||||
var existing model.GenerationOutput
|
||||
existingErr := tx.Where("task_id=? AND sequence_no=?", taskID, 1).Take(&existing).Error
|
||||
if existingErr == nil {
|
||||
output = existing
|
||||
reusedOutput = true
|
||||
} else if errors.Is(existingErr, gorm.ErrRecordNotFound) {
|
||||
if err := tx.Create(&mediaAsset).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if err := tx.Create(&output).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
} else {
|
||||
return existingErr
|
||||
}
|
||||
now := time.Now()
|
||||
cost := row.PrepaidPoints
|
||||
if err := tx.Model(&model.GenerationTask{}).Where("id=?", taskID).Updates(map[string]any{"status": "succeeded", "actual_points": cost, "finished_at": now, "error_code": nil, "error_message": nil}).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if row.TaskType == "video_generation" && row.StoryboardID != nil {
|
||||
premiumDrama, premiumErr := isPremiumDramaGeneration(tx, row.GenerationTask)
|
||||
if premiumErr != nil {
|
||||
return premiumErr
|
||||
}
|
||||
if premiumDrama {
|
||||
if err := tx.Exec(`UPDATE generation_outputs SET metadata=jsonb_set(coalesce(metadata,'{}'::jsonb),'{candidate}','true'::jsonb,true)
|
||||
WHERE id=(SELECT active_output_id FROM episode_storyboards WHERE id=?)`, *row.StoryboardID).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return tx.Model(&model.EpisodeStoryboard{}).Where("id=?", *row.StoryboardID).Updates(map[string]any{"active_output_id": output.ID, "status": "completed"}).Error
|
||||
}
|
||||
if row.TaskType == "image_generation" && row.ProjectID != nil {
|
||||
var input map[string]any
|
||||
_ = json.Unmarshal(row.InputData, &input)
|
||||
if fmt.Sprint(input["target_type"]) == "storyboard" && row.StoryboardID != nil {
|
||||
return tx.Model(&model.EpisodeStoryboard{}).Where("id=?", *row.StoryboardID).Updates(map[string]any{"thumbnail_asset_id": output.MediaAssetID, "status": "completed"}).Error
|
||||
}
|
||||
if assetID, parseErr := uuid.Parse(fmt.Sprint(input["asset_id"])); parseErr == nil {
|
||||
return tx.Model(&model.ProjectAsset{}).Where("id=?", assetID).Update("image_asset_id", output.MediaAssetID).Error
|
||||
}
|
||||
}
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
_ = w.COS.Delete(ctx, key)
|
||||
return err
|
||||
}
|
||||
if cancelledDuringDownload {
|
||||
_ = w.COS.Delete(ctx, key)
|
||||
return w.finishCancellation(ctx, taskID, false, "用户取消生成")
|
||||
}
|
||||
if reusedOutput {
|
||||
_ = w.COS.Delete(ctx, key)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (w *Generation) requeueRejected(ctx context.Context, task model.GenerationTask, message string, delay time.Duration) error {
|
||||
if err := w.DB.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
var current model.GenerationTask
|
||||
if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Where("id=?", task.ID).First(¤t).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
next := time.Now().Add(delay)
|
||||
updates := map[string]any{"status": "pending_submission", "next_submit_at": next, "lease_owner": nil, "lease_until": nil, "error_message": message}
|
||||
if current.SlotReserved && current.SlotReleasedAt == nil {
|
||||
now := time.Now()
|
||||
updates["slot_released_at"] = now
|
||||
}
|
||||
return tx.Model(¤t).Updates(updates).Error
|
||||
}); err != nil {
|
||||
return err
|
||||
}
|
||||
if task.ChannelID != nil {
|
||||
return queuepkg.EnqueueID(w.Queue, queuepkg.TypeDispatchChannel, *task.ChannelID, delay)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (w *Generation) returnUnenqueuedTask(ctx context.Context, taskID uuid.UUID, message string) {
|
||||
var channelID *uuid.UUID
|
||||
_ = w.DB.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
var task model.GenerationTask
|
||||
if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Where("id=?", taskID).First(&task).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if task.Status != "submitting" {
|
||||
return nil
|
||||
}
|
||||
channelID = task.ChannelID
|
||||
now := time.Now()
|
||||
return tx.Model(&task).Updates(map[string]any{
|
||||
"status": "pending_submission", "next_submit_at": now.Add(5 * time.Second),
|
||||
"slot_released_at": now, "lease_owner": nil, "lease_until": nil,
|
||||
"error_message": truncate(message, 1000),
|
||||
}).Error
|
||||
})
|
||||
w.enqueueDispatchAfter(channelID, 5*time.Second)
|
||||
}
|
||||
|
||||
func (w *Generation) finishCancellation(ctx context.Context, taskID uuid.UUID, refund bool, message string) error {
|
||||
return w.DB.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
var task model.GenerationTask
|
||||
if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Where("id=?", taskID).First(&task).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if task.Status == "cancelled" {
|
||||
return nil
|
||||
}
|
||||
now := time.Now()
|
||||
updates := map[string]any{
|
||||
"status": "cancelled", "finished_at": now, "next_poll_at": nil,
|
||||
"lease_owner": nil, "lease_until": nil, "error_code": nil, "error_message": message,
|
||||
}
|
||||
if task.SlotReserved && task.SlotReleasedAt == nil {
|
||||
updates["slot_released_at"] = now
|
||||
}
|
||||
premiumDrama, err := isPremiumDramaGeneration(tx, task)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
// Premium-drama generation keeps its prepaid charge on cancellation.
|
||||
refund = refund && !premiumDrama
|
||||
if refund {
|
||||
refunded, err := billing.RefundGenerationTask(tx, &task, "视频生成失败返还")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if refunded {
|
||||
updates["cost_refunded"] = true
|
||||
}
|
||||
updates["actual_points"] = "0.00"
|
||||
} else {
|
||||
updates["actual_points"] = task.PrepaidPoints
|
||||
}
|
||||
if err := tx.Model(&task).Updates(updates).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if task.TaskType == "video_generation" && task.StoryboardID != nil {
|
||||
return tx.Model(&model.EpisodeStoryboard{}).Where("id=?", *task.StoryboardID).Update("status", "idle").Error
|
||||
}
|
||||
return nil
|
||||
})
|
||||
}
|
||||
|
||||
func (w *Generation) failAndRefund(ctx context.Context, taskID uuid.UUID, code, message string) error {
|
||||
return w.finishFailure(ctx, taskID, code, message, true)
|
||||
}
|
||||
|
||||
func (w *Generation) failWithoutRefund(ctx context.Context, taskID uuid.UUID, code, message string) error {
|
||||
return w.finishFailure(ctx, taskID, code, message, false)
|
||||
}
|
||||
|
||||
func (w *Generation) finishFailure(ctx context.Context, taskID uuid.UUID, code, message string, refund bool) error {
|
||||
return w.DB.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
var task model.GenerationTask
|
||||
if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Where("id=?", taskID).First(&task).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if task.Status == "succeeded" || task.Status == "failed" || task.Status == "cancelled" {
|
||||
return nil
|
||||
}
|
||||
now := time.Now()
|
||||
updates := map[string]any{"status": "failed", "error_code": code, "error_message": truncate(message, 4000), "finished_at": now, "lease_owner": nil, "lease_until": nil}
|
||||
if task.SlotReserved && task.SlotReleasedAt == nil {
|
||||
updates["slot_released_at"] = now
|
||||
}
|
||||
premiumDrama, err := isPremiumDramaGeneration(tx, task)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if premiumDrama && code != "upstream_rejected" && code != "upstream_unavailable" && code != "upstream_failed" {
|
||||
refund = false
|
||||
}
|
||||
if refund {
|
||||
refunded, err := billing.RefundGenerationTask(tx, &task, "视频生成失败返还")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if refunded {
|
||||
updates["cost_refunded"] = true
|
||||
}
|
||||
updates["actual_points"] = "0.00"
|
||||
} else {
|
||||
updates["actual_points"] = task.PrepaidPoints
|
||||
}
|
||||
if err := tx.Model(&task).Updates(updates).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if task.TaskType == "video_generation" && task.StoryboardID != nil {
|
||||
return tx.Model(&model.EpisodeStoryboard{}).Where("id=?", *task.StoryboardID).Update("status", "failed").Error
|
||||
}
|
||||
return nil
|
||||
})
|
||||
}
|
||||
|
||||
func isPremiumDramaGeneration(tx *gorm.DB, task model.GenerationTask) (bool, error) {
|
||||
if task.ProjectID == nil {
|
||||
return false, nil
|
||||
}
|
||||
var projectType string
|
||||
if err := tx.Table("creative_projects").Where("id=?", *task.ProjectID).Pluck("project_type", &projectType).Error; err != nil {
|
||||
return false, err
|
||||
}
|
||||
return projectType == "premium_drama", nil
|
||||
}
|
||||
|
||||
func (w *Generation) enqueueDispatch(channelID *uuid.UUID) {
|
||||
w.enqueueDispatchAfter(channelID, 0)
|
||||
}
|
||||
|
||||
func (w *Generation) enqueueDispatchAfter(channelID *uuid.UUID, delay time.Duration) {
|
||||
if channelID != nil {
|
||||
_ = queuepkg.EnqueueID(w.Queue, queuepkg.TypeDispatchChannel, *channelID, delay)
|
||||
}
|
||||
}
|
||||
|
||||
// Recover 恢复服务中断或旧版本遗留的生成任务，并重新投递待处理工作。
|
||||
func (w *Generation) Recover(ctx context.Context) error {
|
||||
now := time.Now()
|
||||
if err := w.DB.WithContext(ctx).Model(&model.GenerationTask{}).
|
||||
Where("task_type IN ('image_generation','video_generation') AND status='submitting' AND lease_until<?", now).
|
||||
Updates(map[string]any{
|
||||
"status": "pending_submission", "next_submit_at": now, "slot_released_at": now,
|
||||
"error_code": nil, "error_message": "提交中断，正在自动重试", "lease_owner": nil, "lease_until": nil,
|
||||
}).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
channels := make([]struct {
|
||||
ChannelID uuid.UUID
|
||||
NextSubmitAt time.Time
|
||||
}, 0)
|
||||
if err := w.DB.WithContext(ctx).Model(&model.GenerationTask{}).
|
||||
Select("channel_id,min(coalesce(next_submit_at,?)) AS next_submit_at", now).
|
||||
Where("status='pending_submission' AND channel_id IS NOT NULL AND (project_id IS NOT NULL OR (task_type='image_generation' AND input_data->>'product_image'='true'))").Group("channel_id").Scan(&channels).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
for _, item := range channels {
|
||||
delay := time.Until(item.NextSubmitAt)
|
||||
if delay < 0 {
|
||||
delay = 0
|
||||
}
|
||||
_ = queuepkg.EnqueueID(w.Queue, queuepkg.TypeDispatchChannel, item.ChannelID, delay)
|
||||
}
|
||||
var polls []uuid.UUID
|
||||
_ = w.DB.WithContext(ctx).Model(&model.GenerationTask{}).Where("(project_id IS NOT NULL OR (task_type='image_generation' AND input_data->>'product_image'='true')) AND (status IN ('submitted','processing') OR (status='cancel_requested' AND upstream_task_id IS NOT NULL AND upstream_task_id<>''))").Pluck("id", &polls).Error
|
||||
for _, id := range polls {
|
||||
_ = queuepkg.EnqueueID(w.Queue, queuepkg.TypePollTask, id, jitter(5*time.Second))
|
||||
}
|
||||
var downloads []uuid.UUID
|
||||
_ = w.DB.WithContext(ctx).Model(&model.GenerationTask{}).Where("(project_id IS NOT NULL OR (task_type='image_generation' AND input_data->>'product_image'='true')) AND status IN ('result_ready','downloading')").Pluck("id", &downloads).Error
|
||||
for _, id := range downloads {
|
||||
_ = queuepkg.EnqueueID(w.Queue, queuepkg.TypeDownloadTask, id, jitter(5*time.Second))
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func stringOr(value any, fallback string) string {
|
||||
text := strings.TrimSpace(fmt.Sprint(value))
|
||||
if text == "" || text == "<nil>" {
|
||||
return fallback
|
||||
}
|
||||
return text
|
||||
}
|
||||
|
||||
func intOr(value any, fallback int) int {
|
||||
parsed, err := strconv.Atoi(fmt.Sprint(value))
|
||||
if err != nil {
|
||||
return fallback
|
||||
}
|
||||
return parsed
|
||||
}
|
||||
|
||||
func retryDelay(attempt int, retryAfter string) time.Duration {
|
||||
if seconds, err := strconv.Atoi(retryAfter); err == nil && seconds > 0 {
|
||||
return time.Duration(seconds)*time.Second + jitter(3*time.Second)
|
||||
}
|
||||
schedule := []time.Duration{15 * time.Second, 30 * time.Second, time.Minute, 2 * time.Minute, 5 * time.Minute}
|
||||
if attempt < 1 {
|
||||
attempt = 1
|
||||
}
|
||||
if attempt > len(schedule) {
|
||||
attempt = len(schedule)
|
||||
}
|
||||
return schedule[attempt-1] + jitter(5*time.Second)
|
||||
}
|
||||
|
||||
// canRetrySubmission 判断生成任务是否仍在允许的自动提交次数内。
|
||||
func canRetrySubmission(attempt int) bool {
|
||||
return attempt < maxSubmitAttempts
|
||||
}
|
||||
|
||||
func jitter(max time.Duration) time.Duration {
|
||||
if max <= 0 {
|
||||
return 0
|
||||
}
|
||||
return time.Duration(rand.Int63n(int64(max)))
|
||||
}
|
||||
|
||||
func truncate(value string, max int) string {
|
||||
if len(value) > max {
|
||||
return value[:max]
|
||||
}
|
||||
return value
|
||||
}
|
||||
|
||||
func NewHTTPClient(maxConnections int) *http.Client {
|
||||
if maxConnections < 10 {
|
||||
maxConnections = 10
|
||||
}
|
||||
transport := &http.Transport{
|
||||
Proxy: http.ProxyFromEnvironment, DialContext: (&net.Dialer{Timeout: 15 * time.Second, KeepAlive: 30 * time.Second}).DialContext,
|
||||
MaxIdleConns: maxConnections, MaxIdleConnsPerHost: maxConnections, MaxConnsPerHost: maxConnections,
|
||||
// ResponseHeaderTimeout 设为 0 表示不限制，由各调用点的 context 超时统一控制，避免非流式 AI 推理在 60 秒内未返回响应头被误判超时
|
||||
IdleConnTimeout: 90 * time.Second, TLSHandshakeTimeout: 10 * time.Second, ResponseHeaderTimeout: 0,
|
||||
}
|
||||
return &http.Client{Transport: transport}
|
||||
}
|
||||
@@ -0,0 +1,63 @@
|
||||
package worker
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"juhe-factory/api/internal/model"
|
||||
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
func TestSelectFairCandidatesHonorsTotalCapacity(t *testing.T) {
|
||||
users := make([]uuid.UUID, 50)
|
||||
for index := range users {
|
||||
users[index] = uuid.New()
|
||||
}
|
||||
candidates := make([]model.GenerationTask, 0, 501)
|
||||
for index := 0; index < 501; index++ {
|
||||
candidates = append(candidates, model.GenerationTask{ID: uuid.New(), UserID: users[index%len(users)]})
|
||||
}
|
||||
selected := selectFairCandidates(candidates, map[uuid.UUID]int{}, 500, 10)
|
||||
if len(selected) != 500 {
|
||||
t.Fatalf("selected %d tasks, want 500", len(selected))
|
||||
}
|
||||
}
|
||||
|
||||
func TestSelectFairCandidatesHonorsUserCapacity(t *testing.T) {
|
||||
userID := uuid.New()
|
||||
candidates := make([]model.GenerationTask, 11)
|
||||
for index := range candidates {
|
||||
candidates[index] = model.GenerationTask{ID: uuid.New(), UserID: userID}
|
||||
}
|
||||
selected := selectFairCandidates(candidates, map[uuid.UUID]int{}, 500, 10)
|
||||
if len(selected) != 10 {
|
||||
t.Fatalf("selected %d tasks, want 10", len(selected))
|
||||
}
|
||||
}
|
||||
|
||||
func TestSelectFairCandidatesRoundRobin(t *testing.T) {
|
||||
first, second := uuid.New(), uuid.New()
|
||||
candidates := []model.GenerationTask{
|
||||
{ID: uuid.New(), UserID: first}, {ID: uuid.New(), UserID: first},
|
||||
{ID: uuid.New(), UserID: second}, {ID: uuid.New(), UserID: second},
|
||||
}
|
||||
selected := selectFairCandidates(candidates, map[uuid.UUID]int{}, 4, 10)
|
||||
want := []uuid.UUID{first, second, first, second}
|
||||
for index := range want {
|
||||
if selected[index].UserID != want[index] {
|
||||
t.Fatalf("position %d user %s, want %s", index, selected[index].UserID, want[index])
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestCanRetrySubmission 验证提交超时会自动重试，并在第五次失败后停止。
|
||||
func TestCanRetrySubmission(t *testing.T) {
|
||||
for attempt := 1; attempt < maxSubmitAttempts; attempt++ {
|
||||
if !canRetrySubmission(attempt) {
|
||||
t.Fatalf("attempt %d should be retryable", attempt)
|
||||
}
|
||||
}
|
||||
if canRetrySubmission(maxSubmitAttempts) {
|
||||
t.Fatalf("attempt %d should reach retry limit", maxSubmitAttempts)
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,51 @@
|
||||
// 媒体反推任务范围测试，确保单分镜任务不会污染剧集或项目的总处理状态。
|
||||
package worker
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
// TestIsWholeVideoAnalysis 验证仅未绑定分镜的任务可更新反推范围总状态。
|
||||
func TestIsWholeVideoAnalysis(t *testing.T) {
|
||||
storyboardID := uuid.New()
|
||||
tests := []struct {
|
||||
name string
|
||||
snapshot analysisSnapshot
|
||||
want bool
|
||||
}{
|
||||
{name: "整段视频反推", snapshot: analysisSnapshot{}, want: true},
|
||||
{name: "单分镜重新反推", snapshot: analysisSnapshot{StoryboardID: &storyboardID}, want: false},
|
||||
}
|
||||
|
||||
for _, test := range tests {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
if got := isWholeVideoAnalysis(test.snapshot); got != test.want {
|
||||
t.Fatalf("isWholeVideoAnalysis() = %v, want %v", got, test.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestRedrawStoryboardsReady 验证只有全部分镜提示词完整时才恢复剧本反推入口。
|
||||
func TestRedrawStoryboardsReady(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
total int64
|
||||
incomplete int64
|
||||
want bool
|
||||
}{
|
||||
{name: "没有分镜", total: 0, incomplete: 0, want: false},
|
||||
{name: "仍有失败分镜", total: 7, incomplete: 1, want: false},
|
||||
{name: "全部分镜完整", total: 7, incomplete: 0, want: true},
|
||||
}
|
||||
|
||||
for _, test := range tests {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
if got := redrawStoryboardsReady(test.total, test.incomplete); got != test.want {
|
||||
t.Fatalf("redrawStoryboardsReady(%d, %d) = %v, want %v", test.total, test.incomplete, got, test.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,141 @@
|
||||
package worker
|
||||
|
||||
import "testing"
|
||||
|
||||
func TestBuildFixedSegmentsUsesFiveToFifteenSeconds(t *testing.T) {
|
||||
for _, duration := range []float64{5, 15, 19, 31, 33, 39, 45, 46, 121.4} {
|
||||
segments := buildFixedSegments(duration)
|
||||
if len(segments) == 0 {
|
||||
t.Fatalf("duration %.1f returned no segments", duration)
|
||||
}
|
||||
for _, segment := range segments {
|
||||
length := segment.End - segment.Start
|
||||
if length < 5 || length > 15.0001 {
|
||||
t.Fatalf("duration %.1f produced %.3f second segment", duration, length)
|
||||
}
|
||||
}
|
||||
if difference := segments[len(segments)-1].End - duration; difference < -0.001 || difference > 0.001 {
|
||||
t.Fatalf("duration %.1f ends at %.3f", duration, segments[len(segments)-1].End)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestNormalizeAnalysisDurationUsesBoundaryTolerance(t *testing.T) {
|
||||
tests := []struct {
|
||||
input float64
|
||||
want float64
|
||||
}{
|
||||
{input: 4.01, want: 5},
|
||||
{input: 4.999, want: 5},
|
||||
{input: 5, want: 5},
|
||||
{input: 15, want: 15},
|
||||
{input: 15.069002, want: 15},
|
||||
{input: 15.999, want: 15},
|
||||
{input: 16, want: 16},
|
||||
}
|
||||
for _, test := range tests {
|
||||
if got := normalizeAnalysisDuration(test.input); got != test.want {
|
||||
t.Fatalf("normalizeAnalysisDuration(%f) = %f, want %f", test.input, got, test.want)
|
||||
}
|
||||
}
|
||||
segments := buildFixedSegments(normalizeAnalysisDuration(15.069002))
|
||||
if len(segments) != 1 || segments[0].End != 15 {
|
||||
t.Fatalf("15.069002 seconds produced %#v, want one 15 second segment", segments)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildFixedSegmentsRedistributesShortTail(t *testing.T) {
|
||||
tests := []struct {
|
||||
duration float64
|
||||
want []float64
|
||||
}{
|
||||
{duration: 33, want: []float64{15, 9, 9}},
|
||||
{duration: 39, want: []float64{15, 15, 9}},
|
||||
}
|
||||
for _, test := range tests {
|
||||
segments := buildFixedSegments(test.duration)
|
||||
if len(segments) != len(test.want) {
|
||||
t.Fatalf("duration %.1f produced %d segments, want %d", test.duration, len(segments), len(test.want))
|
||||
}
|
||||
for index, segment := range segments {
|
||||
if got := segment.End - segment.Start; got != test.want[index] {
|
||||
t.Fatalf("duration %.1f segment %d = %.1f, want %.1f", test.duration, index+1, got, test.want[index])
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestFrameTimestampsPreserveBoundariesAndLimit(t *testing.T) {
|
||||
segment := fixedSegment{Index: 1, Start: 10, End: 23}
|
||||
frames := frameTimestamps(segment)
|
||||
if len(frames) > 16 {
|
||||
t.Fatalf("returned %d frames, want at most 16", len(frames))
|
||||
}
|
||||
if frames[0] > 10.1 || frames[len(frames)-1] < 22.9 {
|
||||
t.Fatalf("boundaries not preserved: first %.3f last %.3f", frames[0], frames[len(frames)-1])
|
||||
}
|
||||
}
|
||||
|
||||
func TestFrameRetryTimestampMovesBoundarySamplesIntoSegment(t *testing.T) {
|
||||
segment := fixedSegment{Index: 9, Start: 113.31, End: 121.62}
|
||||
if got, ok := frameRetryTimestamp(segment, 121.57); !ok || got != 121.07 {
|
||||
t.Fatalf("end retry = %.3f, %v; want 121.070, true", got, ok)
|
||||
}
|
||||
if got, ok := frameRetryTimestamp(segment, 113.36); !ok || got != 113.86 {
|
||||
t.Fatalf("start retry = %.3f, %v; want 113.860, true", got, ok)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMentionAssetsSkipsDialogue(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
prompt string
|
||||
want string
|
||||
}{
|
||||
{
|
||||
name: "Chinese quotes and all asset types",
|
||||
prompt: "Grace走进古堡大厅，把青铜钥匙放在桌上，说：“Ethan，去古堡大厅找青铜钥匙和神秘徽记。”神秘徽记微微发光。",
|
||||
want: "@Grace走进@古堡大厅，把@青铜钥匙放在桌上，说：“Ethan，去古堡大厅找青铜钥匙和神秘徽记。”@神秘徽记微微发光。",
|
||||
},
|
||||
{
|
||||
name: "English quotes and existing mention",
|
||||
prompt: `@Grace says: "Ethan, take the Key marked \"Grace\"." Ethan nods beside Castle.`,
|
||||
want: `@Grace says: "Ethan, take the Key marked \"Grace\"." @Ethan nods beside @Castle.`,
|
||||
},
|
||||
}
|
||||
names := []string{"Grace", "Ethan", "古堡大厅", "青铜钥匙", "神秘徽记", "Key", "Castle"}
|
||||
for _, test := range tests {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
if got := mentionAssets(test.prompt, names); got != test.want {
|
||||
t.Fatalf("mentionAssets() = %q, want %q", got, test.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestMentionAssetsUsesKnownAssetsAndPrefersLongerNames(t *testing.T) {
|
||||
names := []string{"지우", "지우의 아버지", "도현", "婚礼殿堂"}
|
||||
prompt := "新娘지우走进婚礼殿堂，지우의 아버지看向도현，说：“지우和도현已经到了。”"
|
||||
want := "新娘@지우走进@婚礼殿堂，@지우의 아버지看向@도현，说：“지우和도현已经到了。”"
|
||||
got := mentionAssets(prompt, names)
|
||||
if got != want {
|
||||
t.Fatalf("mentionAssets() = %q, want %q", got, want)
|
||||
}
|
||||
mentioned := mentionedAssetNames(got, names)
|
||||
wantMentioned := []string{"지우", "婚礼殿堂", "지우의 아버지", "도현"}
|
||||
if len(mentioned) != len(wantMentioned) {
|
||||
t.Fatalf("mentionedAssetNames() = %#v, want %#v", mentioned, wantMentioned)
|
||||
}
|
||||
for index := range wantMentioned {
|
||||
if mentioned[index] != wantMentioned[index] {
|
||||
t.Fatalf("mentionedAssetNames() = %#v, want %#v", mentioned, wantMentioned)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseSubtitle(t *testing.T) {
|
||||
items := parseSubtitle("1\n00:00:01.000 --> 00:00:02.500\n第一句\n\n2\n00:00:03,000 --> 00:00:04,000\n2025年，第二句")
|
||||
if len(items) != 2 || items[0].Text != "第一句" || items[1].Text != "2025年，第二句" {
|
||||
t.Fatalf("unexpected subtitles: %#v", items)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,502 @@
|
||||
package worker
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"juhe-factory/api/internal/billing"
|
||||
"juhe-factory/api/internal/model"
|
||||
"juhe-factory/api/internal/provider/apimart"
|
||||
queuepkg "juhe-factory/api/internal/queue"
|
||||
"juhe-factory/api/internal/security"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/hibiken/asynq"
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/clause"
|
||||
)
|
||||
|
||||
type ScriptAnalysis struct {
|
||||
DB *gorm.DB
|
||||
Queue *asynq.Client
|
||||
Encryptor *security.Encryptor
|
||||
Provider apimart.Client
|
||||
}
|
||||
|
||||
type scriptAnalysisSnapshot struct {
|
||||
Task model.GenerationTask
|
||||
AnalysisID uuid.UUID
|
||||
SourceContent string
|
||||
ModelName string
|
||||
BaseURL string
|
||||
APIKeyCiphertext string
|
||||
}
|
||||
|
||||
type scriptHighlight struct {
|
||||
Type string `json:"type"`
|
||||
Content string `json:"content"`
|
||||
}
|
||||
|
||||
type scriptEmotionStage struct {
|
||||
Stage string `json:"stage"`
|
||||
Emotion string `json:"emotion"`
|
||||
Description string `json:"description"`
|
||||
}
|
||||
|
||||
type scriptEpisodeOutline struct {
|
||||
EpisodeNo int `json:"episode_no"`
|
||||
Title string `json:"title"`
|
||||
Outline string `json:"outline"`
|
||||
Emotion string `json:"emotion"`
|
||||
}
|
||||
|
||||
type scriptAnalysisCharacter struct {
|
||||
Name string `json:"name"`
|
||||
Faction string `json:"faction"`
|
||||
Biography string `json:"biography"`
|
||||
Motivation string `json:"motivation"`
|
||||
Relationship string `json:"relationship"`
|
||||
}
|
||||
|
||||
type scriptEpisodeHook struct {
|
||||
EpisodeNo int `json:"episode_no"`
|
||||
HookType string `json:"hook_type"`
|
||||
Hook string `json:"hook"`
|
||||
Basis string `json:"basis"`
|
||||
}
|
||||
|
||||
type scriptPaywallHook struct {
|
||||
HookType string `json:"hook_type"`
|
||||
Content string `json:"content"`
|
||||
Appeal string `json:"appeal"`
|
||||
}
|
||||
|
||||
type scriptCustomAnalysis struct {
|
||||
Objective string `json:"objective"`
|
||||
Conclusion string `json:"conclusion"`
|
||||
Evidence []string `json:"evidence"`
|
||||
}
|
||||
|
||||
type scriptUnmatchedObjective struct {
|
||||
Objective string `json:"objective"`
|
||||
Reason string `json:"reason"`
|
||||
}
|
||||
|
||||
type scriptAnalysisResult struct {
|
||||
Tags struct {
|
||||
DramaTags []string `json:"drama_tags"`
|
||||
ApplicableSceneTags []string `json:"applicable_scene_tags"`
|
||||
} `json:"tags"`
|
||||
SellingPoint struct {
|
||||
Summary string `json:"summary"`
|
||||
Highlights []scriptHighlight `json:"highlights"`
|
||||
} `json:"selling_point"`
|
||||
StoryStructure struct {
|
||||
Synopsis string `json:"synopsis"`
|
||||
StoryOutline []string `json:"story_outline"`
|
||||
EmotionOutline []scriptEmotionStage `json:"emotion_outline"`
|
||||
EpisodeOutlines []scriptEpisodeOutline `json:"episode_outlines"`
|
||||
Characters []scriptAnalysisCharacter `json:"characters"`
|
||||
} `json:"story_structure"`
|
||||
OpeningHook struct {
|
||||
FirstScene string `json:"first_scene"`
|
||||
FirstLine string `json:"first_line"`
|
||||
AttentionMethods []string `json:"attention_methods"`
|
||||
Analysis string `json:"analysis"`
|
||||
Optimization string `json:"optimization"`
|
||||
} `json:"opening_hook"`
|
||||
EpisodeHooks []scriptEpisodeHook `json:"episode_hooks"`
|
||||
Episode10PaywallHooks []scriptPaywallHook `json:"episode_10_paywall_hooks"`
|
||||
CustomAnalysis []scriptCustomAnalysis `json:"custom_analysis"`
|
||||
UnmatchedObjectives []scriptUnmatchedObjective `json:"unmatched_objectives"`
|
||||
}
|
||||
|
||||
func NewScriptAnalysis(db *gorm.DB, queue *asynq.Client, encryptor *security.Encryptor, client *http.Client) *ScriptAnalysis {
|
||||
return &ScriptAnalysis{DB: db, Queue: queue, Encryptor: encryptor, Provider: apimart.NewClient(dramaHTTPClient(client))}
|
||||
}
|
||||
|
||||
func (w *ScriptAnalysis) Register(mux *asynq.ServeMux) {
|
||||
mux.HandleFunc(queuepkg.TypeAnalyzeScript, w.analyze)
|
||||
}
|
||||
|
||||
func (w *ScriptAnalysis) Recover(ctx context.Context) error {
|
||||
if err := w.DB.WithContext(ctx).Model(&model.GenerationTask{}).
|
||||
Where("task_type='script_analysis' AND status='processing'").Update("status", "submitted").Error; err != nil {
|
||||
return err
|
||||
}
|
||||
var ids []uuid.UUID
|
||||
if err := w.DB.WithContext(ctx).Model(&model.GenerationTask{}).
|
||||
Where("task_type='script_analysis' AND status='submitted'").Pluck("id", &ids).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
var firstErr error
|
||||
for _, id := range ids {
|
||||
if err := queuepkg.EnqueueID(w.Queue, queuepkg.TypeAnalyzeScript, id, 0); err != nil && firstErr == nil {
|
||||
firstErr = err
|
||||
}
|
||||
}
|
||||
var cancelledIDs []uuid.UUID
|
||||
if err := w.DB.WithContext(ctx).Model(&model.GenerationTask{}).
|
||||
Where("task_type='script_analysis' AND status='cancel_requested'").Pluck("id", &cancelledIDs).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
for _, id := range cancelledIDs {
|
||||
if err := w.finishCancellation(ctx, id, uuid.Nil); err != nil && firstErr == nil {
|
||||
firstErr = err
|
||||
}
|
||||
}
|
||||
return firstErr
|
||||
}
|
||||
|
||||
func (w *ScriptAnalysis) analyze(ctx context.Context, raw *asynq.Task) error {
|
||||
taskID, err := decodeID(raw)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
snapshot, err := w.claim(ctx, taskID)
|
||||
if err != nil {
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
return nil
|
||||
}
|
||||
return w.fail(ctx, taskID, uuid.Nil, err)
|
||||
}
|
||||
apiKey, err := w.Encryptor.Decrypt(snapshot.APIKeyCiphertext)
|
||||
if err != nil {
|
||||
return w.fail(ctx, taskID, snapshot.AnalysisID, errors.New("剧本分析渠道密钥不可用"))
|
||||
}
|
||||
workCtx, cancel := context.WithTimeout(ctx, 12*time.Minute)
|
||||
defer cancel()
|
||||
var cancelled atomic.Bool
|
||||
go w.watchCancellation(workCtx, cancel, taskID, &cancelled)
|
||||
result, chat, inputTokens, outputTokens, tokenSource, err := w.callModel(workCtx, snapshot, apiKey)
|
||||
if cancelled.Load() || w.cancellationRequested(ctx, taskID) {
|
||||
return w.finishCancellation(ctx, taskID, snapshot.AnalysisID)
|
||||
}
|
||||
if err != nil {
|
||||
return w.fail(ctx, taskID, snapshot.AnalysisID, err)
|
||||
}
|
||||
if err := w.persist(ctx, snapshot, result, chat, inputTokens, outputTokens, tokenSource); err != nil {
|
||||
return w.fail(ctx, taskID, snapshot.AnalysisID, err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (w *ScriptAnalysis) watchCancellation(ctx context.Context, cancel context.CancelFunc, taskID uuid.UUID, flag *atomic.Bool) {
|
||||
ticker := time.NewTicker(time.Second)
|
||||
defer ticker.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-ticker.C:
|
||||
if w.cancellationRequested(ctx, taskID) {
|
||||
flag.Store(true)
|
||||
cancel()
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (w *ScriptAnalysis) cancellationRequested(ctx context.Context, taskID uuid.UUID) bool {
|
||||
var status string
|
||||
return w.DB.WithContext(ctx).Model(&model.GenerationTask{}).Where("id=?", taskID).Pluck("status", &status).Error == nil &&
|
||||
status == "cancel_requested"
|
||||
}
|
||||
|
||||
func (w *ScriptAnalysis) claim(ctx context.Context, taskID uuid.UUID) (scriptAnalysisSnapshot, error) {
|
||||
var snapshot scriptAnalysisSnapshot
|
||||
err := w.DB.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).
|
||||
Where("id=? AND task_type='script_analysis' AND status='submitted'", taskID).Take(&snapshot.Task).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if snapshot.Task.ScriptAnalysisID == nil {
|
||||
return errors.New("剧本分析任务缺少剧本标识")
|
||||
}
|
||||
snapshot.AnalysisID = *snapshot.Task.ScriptAnalysisID
|
||||
var input struct {
|
||||
SourceContent string `json:"source_content"`
|
||||
}
|
||||
if err := json.Unmarshal(snapshot.Task.InputData, &input); err != nil || strings.TrimSpace(input.SourceContent) == "" {
|
||||
return errors.New("剧本分析原文快照无效")
|
||||
}
|
||||
snapshot.SourceContent = input.SourceContent
|
||||
var config struct {
|
||||
ModelName string
|
||||
BaseURL string
|
||||
APIKeyCiphertext string
|
||||
}
|
||||
result := tx.Raw(`SELECT model.name AS model_name,channel.base_url,channel.api_key_ciphertext
|
||||
FROM models model JOIN channels channel ON channel.id=model.channel_id
|
||||
WHERE model.id=? AND channel.id=? AND model.model_type='text'
|
||||
AND model.enabled=true AND model.deleted_at IS NULL AND channel.enabled=true AND channel.deleted_at IS NULL`,
|
||||
snapshot.Task.ModelID, snapshot.Task.ChannelID).Scan(&config)
|
||||
if result.Error != nil {
|
||||
return result.Error
|
||||
}
|
||||
if result.RowsAffected == 0 {
|
||||
return errors.New("剧本分析模型或渠道不可用")
|
||||
}
|
||||
snapshot.ModelName, snapshot.BaseURL, snapshot.APIKeyCiphertext = config.ModelName, config.BaseURL, config.APIKeyCiphertext
|
||||
if err := tx.Model(&snapshot.Task).Update("status", "processing").Error; err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Model(&model.ScriptAnalysis{}).Where("id=?", snapshot.AnalysisID).
|
||||
Updates(map[string]any{"analysis_status": "running", "analysis_message": "正在分析剧本"}).Error
|
||||
})
|
||||
return snapshot, err
|
||||
}
|
||||
|
||||
func (w *ScriptAnalysis) callModel(ctx context.Context, snapshot scriptAnalysisSnapshot, apiKey string) (scriptAnalysisResult, apimart.ChatResult, int64, int64, string, error) {
|
||||
var parsed scriptAnalysisResult
|
||||
var chat apimart.ChatResult
|
||||
pricing, err := billing.ParseTextPricingSnapshot(snapshot.Task.BillingSnapshot)
|
||||
if err != nil {
|
||||
return parsed, chat, 0, 0, "", err
|
||||
}
|
||||
systemPrompt := scriptAnalysisSystemPrompt(snapshot.Task.PromptSnapshot)
|
||||
userPrompt := "请分析以下剧本原文：\n\n" + snapshot.SourceContent
|
||||
if pricing.Mode == billing.TextBillingPerToken {
|
||||
estimated := billing.EstimateTextTokens(systemPrompt + userPrompt)
|
||||
if err := w.DB.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
_, reserveErr := billing.ReserveTextCall(tx, snapshot.Task.UserID, snapshot.Task.ID.String(), "analysis", pricing, estimated, "剧本分析")
|
||||
return reserveErr
|
||||
}); err != nil {
|
||||
return parsed, chat, 0, 0, "", err
|
||||
}
|
||||
}
|
||||
payload := map[string]any{
|
||||
"model": snapshot.ModelName,
|
||||
"messages": []map[string]any{
|
||||
{"role": "system", "content": systemPrompt},
|
||||
{"role": "user", "content": userPrompt},
|
||||
},
|
||||
"temperature": 0.2,
|
||||
"response_format": map[string]any{"type": "json_object"},
|
||||
}
|
||||
var parseErr error
|
||||
for attempt := 0; attempt < 2; attempt++ {
|
||||
if attempt > 0 {
|
||||
payload["temperature"] = 0
|
||||
payload["messages"] = []map[string]any{
|
||||
{"role": "system", "content": systemPrompt},
|
||||
{"role": "user", "content": userPrompt + "\n\n上次返回未通过结构校验。请严格补齐固定协议的全部字段，并按分析目标决定标准字段、补充分析或无法完成的目标，只返回完整合法 JSON。"},
|
||||
}
|
||||
}
|
||||
chat, err = w.Provider.ChatWithUsage(ctx, snapshot.BaseURL, apiKey, payload)
|
||||
if err != nil {
|
||||
return parsed, chat, 0, 0, "", err
|
||||
}
|
||||
parsed = scriptAnalysisResult{}
|
||||
parseErr = json.Unmarshal([]byte(stripJSONFence(chat.Content)), &parsed)
|
||||
if parseErr == nil {
|
||||
parseErr = validateScriptAnalysisResult(parsed)
|
||||
}
|
||||
if parseErr == nil {
|
||||
break
|
||||
}
|
||||
}
|
||||
if parseErr != nil {
|
||||
finishReason := strings.TrimSpace(chat.FinishReason)
|
||||
if finishReason == "" {
|
||||
finishReason = "未知"
|
||||
}
|
||||
return parsed, chat, 0, 0, "", fmt.Errorf("文本模型连续两次返回无法解析（结束原因：%s）: %w", finishReason, parseErr)
|
||||
}
|
||||
inputTokens, outputTokens := chat.InputTokens, chat.OutputTokens
|
||||
tokenSource := "upstream"
|
||||
if inputTokens == 0 {
|
||||
inputTokens = billing.EstimateTextTokens(systemPrompt + userPrompt)
|
||||
tokenSource = "local"
|
||||
}
|
||||
if outputTokens == 0 {
|
||||
outputTokens = billing.EstimateTextTokens(chat.Content)
|
||||
if tokenSource == "upstream" {
|
||||
tokenSource = "mixed"
|
||||
}
|
||||
}
|
||||
if pricing.Mode == billing.TextBillingPerToken {
|
||||
if err := w.DB.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
settled, settleErr := billing.SettleTextCall(tx, snapshot.Task.UserID, snapshot.Task.ID.String(), "analysis", pricing,
|
||||
billing.TextUsage{Input: inputTokens, Output: outputTokens}, "剧本分析")
|
||||
if settleErr != nil {
|
||||
return settleErr
|
||||
}
|
||||
return tx.Model(&model.GenerationTask{}).Where("id=?", snapshot.Task.ID).Updates(map[string]any{
|
||||
"estimated_points": gorm.Expr("estimated_points+?::numeric", settled),
|
||||
"prepaid_points": gorm.Expr("prepaid_points+?::numeric", settled),
|
||||
}).Error
|
||||
}); err != nil {
|
||||
return parsed, chat, inputTokens, outputTokens, tokenSource, err
|
||||
}
|
||||
}
|
||||
return parsed, chat, inputTokens, outputTokens, tokenSource, nil
|
||||
}
|
||||
|
||||
func validateScriptAnalysisResult(result scriptAnalysisResult) error {
|
||||
if strings.TrimSpace(result.SellingPoint.Summary) == "" && strings.TrimSpace(result.StoryStructure.Synopsis) == "" &&
|
||||
len(result.CustomAnalysis) == 0 && len(result.UnmatchedObjectives) == 0 {
|
||||
return errors.New("缺少标准分析结果、通用分析结果或无法匹配目标说明")
|
||||
}
|
||||
if len(result.EpisodeHooks) != 0 && len(result.EpisodeHooks) != 10 {
|
||||
return errors.New("前十集钩子数量必须为10")
|
||||
}
|
||||
for index, hook := range result.EpisodeHooks {
|
||||
if hook.EpisodeNo != index+1 {
|
||||
return errors.New("前十集钩子集数或顺序无效")
|
||||
}
|
||||
}
|
||||
if len(result.Episode10PaywallHooks) != 0 &&
|
||||
(len(result.Episode10PaywallHooks) < 3 || len(result.Episode10PaywallHooks) > 5) {
|
||||
return errors.New("第十集付费卡点候选必须为3至5个")
|
||||
}
|
||||
for _, character := range result.StoryStructure.Characters {
|
||||
if strings.TrimSpace(character.Name) == "" {
|
||||
return errors.New("人物名称不能为空")
|
||||
}
|
||||
if character.Faction != "正派" && character.Faction != "反派" && character.Faction != "中立派" {
|
||||
return errors.New("人物阵营无效")
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (w *ScriptAnalysis) persist(ctx context.Context, snapshot scriptAnalysisSnapshot, result scriptAnalysisResult, chat apimart.ChatResult, inputTokens, outputTokens int64, tokenSource string) error {
|
||||
encoded, err := json.Marshal(result)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
pretty, _ := json.MarshalIndent(result, "", " ")
|
||||
return w.DB.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
var task model.GenerationTask
|
||||
if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Where("id=?", snapshot.Task.ID).Take(&task).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if task.Status == "cancel_requested" {
|
||||
return w.finishCancellationLocked(tx, &task, snapshot.AnalysisID)
|
||||
}
|
||||
if task.Status != "processing" {
|
||||
return errors.New("剧本分析任务状态已变化")
|
||||
}
|
||||
if err := tx.Where("script_analysis_id=?", snapshot.AnalysisID).Delete(&model.ScriptAnalysisCharacter{}).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
seen := make(map[string]bool)
|
||||
for index, character := range result.StoryStructure.Characters {
|
||||
name := strings.TrimSpace(character.Name)
|
||||
key := strings.ToLower(name)
|
||||
if seen[key] {
|
||||
continue
|
||||
}
|
||||
seen[key] = true
|
||||
item := model.ScriptAnalysisCharacter{
|
||||
ScriptAnalysisID: snapshot.AnalysisID, Name: name, Faction: character.Faction,
|
||||
Biography: strings.TrimSpace(character.Biography), Motivation: strings.TrimSpace(character.Motivation),
|
||||
Relationship: strings.TrimSpace(character.Relationship), SortOrder: index,
|
||||
}
|
||||
if err := tx.Create(&item).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
taskUpdates := map[string]any{
|
||||
"status": "succeeded", "finished_at": time.Now(), "actual_points": task.PrepaidPoints,
|
||||
"input_tokens": inputTokens, "output_tokens": outputTokens, "total_tokens": inputTokens + outputTokens,
|
||||
"token_count_source": tokenSource, "usage_raw": chat.UsageRaw, "result_data": gorm.Expr("?::jsonb", string(encoded)),
|
||||
"error_code": nil, "error_message": nil,
|
||||
}
|
||||
if err := tx.Model(&task).Updates(taskUpdates).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Model(&model.ScriptAnalysis{}).Where("id=?", snapshot.AnalysisID).Updates(map[string]any{
|
||||
"analysis_result": gorm.Expr("?::jsonb", string(encoded)), "result_content": string(pretty),
|
||||
"analysis_status": "succeeded", "analysis_message": "剧本分析完成",
|
||||
}).Error
|
||||
})
|
||||
}
|
||||
|
||||
func (w *ScriptAnalysis) fail(ctx context.Context, taskID, analysisID uuid.UUID, taskErr error) error {
|
||||
return w.DB.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
var task model.GenerationTask
|
||||
if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Where("id=?", taskID).Take(&task).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if task.Status == "succeeded" || task.Status == "failed" || task.Status == "cancelled" {
|
||||
return nil
|
||||
}
|
||||
if task.Status == "cancel_requested" {
|
||||
return w.finishCancellationLocked(tx, &task, analysisID)
|
||||
}
|
||||
refunded, err := billing.RefundTextGenerationTask(tx, &task, "剧本分析失败返还")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
updates := map[string]any{
|
||||
"status": "failed", "actual_points": "0.00", "finished_at": time.Now(),
|
||||
"error_code": "script_analysis_failed", "error_message": truncate(taskErr.Error(), 4000),
|
||||
}
|
||||
if refunded {
|
||||
updates["cost_refunded"] = true
|
||||
}
|
||||
if err := tx.Model(&task).Updates(updates).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if analysisID == uuid.Nil && task.ScriptAnalysisID != nil {
|
||||
analysisID = *task.ScriptAnalysisID
|
||||
}
|
||||
if analysisID == uuid.Nil {
|
||||
return nil
|
||||
}
|
||||
return tx.Model(&model.ScriptAnalysis{}).Where("id=?", analysisID).Updates(map[string]any{
|
||||
"analysis_status": "failed", "analysis_message": truncate(taskErr.Error(), 1000),
|
||||
}).Error
|
||||
})
|
||||
}
|
||||
|
||||
func (w *ScriptAnalysis) finishCancellation(ctx context.Context, taskID, analysisID uuid.UUID) error {
|
||||
return w.DB.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
var task model.GenerationTask
|
||||
if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Where("id=?", taskID).Take(&task).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
return w.finishCancellationLocked(tx, &task, analysisID)
|
||||
})
|
||||
}
|
||||
|
||||
func (w *ScriptAnalysis) finishCancellationLocked(tx *gorm.DB, task *model.GenerationTask, analysisID uuid.UUID) error {
|
||||
if task.Status == "cancelled" {
|
||||
return nil
|
||||
}
|
||||
if task.Status == "succeeded" || task.Status == "failed" {
|
||||
return nil
|
||||
}
|
||||
actual := task.PrepaidPoints
|
||||
committed, err := billing.CommitTextReserves(tx, task.UserID, task.ID.String(), "剧本分析取消")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if committed != "" && committed != "0" && committed != "0.00" {
|
||||
actual = committed
|
||||
}
|
||||
if err := tx.Model(task).Updates(map[string]any{
|
||||
"status": "cancelled", "prepaid_points": actual, "actual_points": actual, "finished_at": time.Now(),
|
||||
"error_code": "user_cancelled", "error_message": "用户取消，已扣积分不退",
|
||||
}).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if analysisID == uuid.Nil && task.ScriptAnalysisID != nil {
|
||||
analysisID = *task.ScriptAnalysisID
|
||||
}
|
||||
if analysisID == uuid.Nil {
|
||||
return nil
|
||||
}
|
||||
return tx.Model(&model.ScriptAnalysis{}).Where("id=?", analysisID).Updates(map[string]any{
|
||||
"analysis_status": "idle", "analysis_message": "剧本分析已取消，已扣积分不退",
|
||||
}).Error
|
||||
}
|
||||
@@ -0,0 +1,61 @@
|
||||
package worker
|
||||
|
||||
import "strings"
|
||||
|
||||
const scriptAnalysisOutputRules = `将分析目标能够对应到标准字段的结果写入标准字段；无法对应到标准字段的目标写入 custom_analysis；因原文信息不足或目标不适用于剧本而无法完成的目标写入 unmatched_objectives。不得为了填充标准字段而曲解分析目标。
|
||||
|
||||
必须只返回一个合法 JSON 对象，不得输出 Markdown、代码围栏、解释或额外文字。所有字段均须存在并严格符合以下结构：
|
||||
{
|
||||
"tags": {
|
||||
"drama_tags": ["标签"],
|
||||
"applicable_scene_tags": ["适用场景标签"]
|
||||
},
|
||||
"selling_point": {
|
||||
"summary": "一句话总结",
|
||||
"highlights": [
|
||||
{"type": "爽点|痛点|卖点|看点", "content": "具体内容"}
|
||||
]
|
||||
},
|
||||
"story_structure": {
|
||||
"synopsis": "短剧简介",
|
||||
"story_outline": ["按故事阶段组织的大纲"],
|
||||
"emotion_outline": [
|
||||
{"stage": "情绪阶段", "emotion": "主导情绪", "description": "情绪变化"}
|
||||
],
|
||||
"episode_outlines": [
|
||||
{"episode_no": 1, "title": "集标题", "outline": "本集大纲", "emotion": "本集情绪"}
|
||||
],
|
||||
"characters": [
|
||||
{"name": "人物名", "faction": "正派|反派|中立派", "biography": "人物小传", "motivation": "核心动机", "relationship": "关键人物关系"}
|
||||
]
|
||||
},
|
||||
"opening_hook": {
|
||||
"first_scene": "开篇第一个画面或场景",
|
||||
"first_line": "第一句台词",
|
||||
"attention_methods": ["悬念前置|激烈冲突|视觉奇观|其他明确手法"],
|
||||
"analysis": "为什么能或不能抓住观众",
|
||||
"optimization": "必要时给出优化建议"
|
||||
},
|
||||
"episode_hooks": [
|
||||
{"episode_no": 1, "hook_type": "钩子类型", "hook": "本集结尾或关键卡点钩子", "basis": "对应的原文依据或信息不足说明"}
|
||||
],
|
||||
"episode_10_paywall_hooks": [
|
||||
{"hook_type": "钩子类型", "content": "付费卡点候选", "appeal": "促使继续付费观看的原因"}
|
||||
],
|
||||
"custom_analysis": [
|
||||
{"objective": "无法对应标准字段的分析目标", "conclusion": "分析结论", "evidence": ["剧本原文依据"]}
|
||||
],
|
||||
"unmatched_objectives": [
|
||||
{"objective": "无法完成的分析目标", "reason": "无法完成的原因"}
|
||||
]
|
||||
}
|
||||
|
||||
仅当分析目标要求逐集钩子时，episode_hooks 必须按 episode_no 从1到10排列且恰好返回10项；仅当分析目标要求第10集付费卡点时，episode_10_paywall_hooks 返回3至5项。其他情况下对应数组返回空数组。`
|
||||
|
||||
func scriptAnalysisSystemPrompt(analysisGoal string) string {
|
||||
goal := strings.TrimSpace(analysisGoal)
|
||||
if goal == "" {
|
||||
return scriptAnalysisOutputRules
|
||||
}
|
||||
return goal + "\n\n" + scriptAnalysisOutputRules
|
||||
}
|
||||
@@ -0,0 +1,19 @@
|
||||
package worker
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestScriptAnalysisSystemPromptKeepsGoalAndFixedSchema(t *testing.T) {
|
||||
prompt := scriptAnalysisSystemPrompt("你是一名分析师。\n\n分析目标：\n分析拍摄预算")
|
||||
if !strings.Contains(prompt, "分析目标：\n分析拍摄预算") {
|
||||
t.Fatal("prompt does not contain the analysis goal")
|
||||
}
|
||||
if strings.Count(prompt, "你是一名分析师") != 1 {
|
||||
t.Fatal("prompt unexpectedly duplicates the database prompt")
|
||||
}
|
||||
if !strings.Contains(prompt, `"custom_analysis"`) || !strings.Contains(prompt, `"unmatched_objectives"`) {
|
||||
t.Fatal("prompt does not contain the fixed fallback schema")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,38 @@
|
||||
package worker
|
||||
|
||||
import "testing"
|
||||
|
||||
func TestValidateScriptAnalysisResult(t *testing.T) {
|
||||
var result scriptAnalysisResult
|
||||
result.SellingPoint.Summary = "核心卖点"
|
||||
result.StoryStructure.Synopsis = "剧情梗概"
|
||||
for episode := 1; episode <= 10; episode++ {
|
||||
result.EpisodeHooks = append(result.EpisodeHooks, scriptEpisodeHook{EpisodeNo: episode})
|
||||
}
|
||||
for index := 0; index < 3; index++ {
|
||||
result.Episode10PaywallHooks = append(result.Episode10PaywallHooks, scriptPaywallHook{Content: "候选"})
|
||||
}
|
||||
result.StoryStructure.Characters = []scriptAnalysisCharacter{{Name: "甲", Faction: "正派"}, {Name: "乙", Faction: "反派"}, {Name: "丙", Faction: "中立派"}}
|
||||
if err := validateScriptAnalysisResult(result); err != nil {
|
||||
t.Fatalf("expected valid analysis result: %v", err)
|
||||
}
|
||||
|
||||
result.EpisodeHooks[9].EpisodeNo = 9
|
||||
if err := validateScriptAnalysisResult(result); err == nil {
|
||||
t.Fatal("expected invalid episode hook ordering")
|
||||
}
|
||||
}
|
||||
|
||||
func TestValidateScriptAnalysisResultRejectsInvalidFaction(t *testing.T) {
|
||||
var result scriptAnalysisResult
|
||||
result.SellingPoint.Summary = "核心卖点"
|
||||
result.StoryStructure.Synopsis = "剧情梗概"
|
||||
for episode := 1; episode <= 10; episode++ {
|
||||
result.EpisodeHooks = append(result.EpisodeHooks, scriptEpisodeHook{EpisodeNo: episode})
|
||||
}
|
||||
result.Episode10PaywallHooks = make([]scriptPaywallHook, 3)
|
||||
result.StoryStructure.Characters = []scriptAnalysisCharacter{{Name: "甲", Faction: "未知"}}
|
||||
if err := validateScriptAnalysisResult(result); err == nil {
|
||||
t.Fatal("expected invalid character faction")
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user