666 lines
26 KiB
Go
666 lines
26 KiB
Go
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"
|
||
}
|