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 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<.*?`) 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" }