503 lines
18 KiB
Go
503 lines
18 KiB
Go
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
|
|
}
|