diff --git a/short_drama.db b/short_drama.db index 0b2fb53..9caef56 100644 Binary files a/short_drama.db and b/short_drama.db differ diff --git a/shortdrama/agent/chat_model.go b/shortdrama/agent/chat_model.go index 9b8a8de..b85d945 100644 --- a/shortdrama/agent/chat_model.go +++ b/shortdrama/agent/chat_model.go @@ -16,22 +16,33 @@ type ModelConfig struct { BaseURL string // API地址 MaxTokens int // 最大Token数 Temperature float32 // 温度参数 - ImageModel string // 图片模型名 } // context keys type ctxKey string const ( - ctxKeyAPIKey ctxKey = "api_key" - ctxKeyImageModel ctxKey = "image_model" - ctxKeyBaseURL ctxKey = "base_url" + ctxKeyAPIKey ctxKey = "api_key" + ctxKeyBaseURL ctxKey = "base_url" + ctxKeyDramaId ctxKey = "drama_id" ) +// WithDramaId 将短剧ID注入 context，供工具函数读取场景图片 +func WithDramaId(ctx context.Context, dramaId int64) context.Context { + return context.WithValue(ctx, ctxKeyDramaId, dramaId) +} + +// GetDramaId 从 context 获取短剧ID +func GetDramaId(ctx context.Context) int64 { + if v, ok := ctx.Value(ctxKeyDramaId).(int64); ok { + return v + } + return 0 +} + // WithModelConfig 将模型配置注入 context，供工具函数读取 func WithModelConfig(ctx context.Context, cfg *ModelConfig) context.Context { ctx = context.WithValue(ctx, ctxKeyAPIKey, cfg.APIKey) - ctx = context.WithValue(ctx, ctxKeyImageModel, cfg.ImageModel) ctx = context.WithValue(ctx, ctxKeyBaseURL, cfg.BaseURL) return ctx } @@ -44,14 +55,6 @@ func GetAPIKey(ctx context.Context) string { return "" } -// GetImageModel 从 context 获取图片模型名 -func GetImageModel(ctx context.Context) string { - if v, ok := ctx.Value(ctxKeyImageModel).(string); ok && v != "" { - return v - } - return "" -} - // GetBaseURL 从 context 获取 API 地址 func GetBaseURL(ctx context.Context) string { if v, ok := ctx.Value(ctxKeyBaseURL).(string); ok && v != "" { diff --git a/shortdrama/agent/tools.go b/shortdrama/agent/tools.go index 746c5ed..44c1710 100644 --- a/shortdrama/agent/tools.go +++ b/shortdrama/agent/tools.go @@ -1,16 +1,18 @@ package agent import ( - "bytes" "context" + "encoding/base64" "encoding/json" "fmt" - "io" - "net/http" + "os" + "path/filepath" "strings" - "time" + + "video-factory/shortdrama/dao" "github.com/cloudwego/eino/schema" + "github.com/gogf/gf/v2/frame/g" ) // ToolInfo 工具定义 @@ -263,7 +265,7 @@ func analyzeScriptForEpisodeTool() *ToolInfo { func generateSceneImageTool() *ToolInfo { return &ToolInfo{ Name: "generate_scene_image", - Description: "根据场景的画面描述生成场景图片，返回图片的base64编码数据", + Description: "根据场景的画面描述获取场景图片，返回图片的base64编码数据。优先从已有场景库中读取图片，不存在时返回空。", Parameters: map[string]any{ "type": "object", "properties": map[string]any{ @@ -293,9 +295,33 @@ func generateSceneImageTool() *ToolInfo { sceneIdx, _ := args["scene_index"].(float64) prompt := fmt.Sprintf("画面描述：%s，风格：%s", visualDesc, style) - imgBase64, err := generateRealImage(ctx, prompt) - if err != nil { - return "", fmt.Errorf("生成场景图片失败: %w", err) + _ = prompt // 保留供日志使用 + + // 从 DB 场景表中查找已有场景图片 + dramaId := GetDramaId(ctx) + var imgBase64 string + if dramaId > 0 && visualDesc != "" { + scenes, err := dao.Scene.ListByDrama(ctx, dramaId) + if err == nil { + for _, sc := range scenes { + if sc.ImagePath == "" { + continue + } + // 匹配场景描述（双向包含匹配） + if strings.Contains(visualDesc, sc.Description) || strings.Contains(sc.Description, visualDesc) { + b64, err := readImageFileAsBase64(sc.ImagePath) + if err == nil { + imgBase64 = b64 + g.Log().Infof(ctx, "从场景库读取图片: %s (场景: %s)", sc.ImagePath, sc.Name) + } + break + } + } + } + } + + if imgBase64 == "" { + g.Log().Infof(ctx, "场景库中无匹配图片，返回空(episode=%d, scene=%d)", int(episodeIdx), int(sceneIdx)) } result, _ := json.Marshal(map[string]any{ @@ -329,75 +355,20 @@ func cleanEpisodeTitle(title string) string { return title } -func generateRealImage(ctx context.Context, prompt string) (string, error) { - imageModel := GetImageModel(ctx) - if imageModel == "" { - // 图片模型未配置，返回空字符串（不报错），避免 agent 反复重试 - return "", nil - } - - // 调用通义万相生成图片 - apiKey := GetAPIKey(ctx) - url := "https://dashscope.aliyuncs.com/api/v1/services/aigc/multimodal-generation/generation" - - body := map[string]any{ - "model": imageModel, - "input": map[string]any{ - "messages": []map[string]any{ - { - "role": "user", - "content": []map[string]string{ - {"type": "text", "text": prompt}, - }, - }, - }, - }, - "parameters": map[string]any{ - "size": "1024*1364", - "n": 1, - "watermark": false, - }, - } - - payload, _ := json.Marshal(body) - httpClient := &http.Client{Timeout: 60 * time.Second} - req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewBuffer(payload)) +func readImageFileAsBase64(path string) (string, error) { + data, err := os.ReadFile(path) if err != nil { return "", err } - req.Header.Set("Authorization", "Bearer "+apiKey) - req.Header.Set("Content-Type", "application/json") - - resp, err := httpClient.Do(req) - if err != nil { - return "", err + ext := strings.ToLower(filepath.Ext(path)) + mime := "image/png" + switch ext { + case ".jpg", ".jpeg": + mime = "image/jpeg" + case ".gif": + mime = "image/gif" + case ".webp": + mime = "image/webp" } - defer resp.Body.Close() - - data, _ := io.ReadAll(resp.Body) - - var result struct { - Output struct { - Choices []struct { - Message struct { - Content []struct { - Image string `json:"image"` - } `json:"content"` - } `json:"message"` - } `json:"choices"` - } `json:"workspace"` - Code string `json:"code"` - } - - err = json.Unmarshal(data, &result) - if err != nil || len(result.Output.Choices) == 0 || result.Code != "" { - return "", fmt.Errorf("生成图片失败: %s", string(data)) - } - - imgBase64 := result.Output.Choices[0].Message.Content[0].Image - if imgBase64 == "" { - return "", fmt.Errorf("图片内容为空") - } - - return imgBase64, nil + return "data:" + mime + ";base64," + base64.StdEncoding.EncodeToString(data), nil } diff --git a/shortdrama/service/drama_service.go b/shortdrama/service/drama_service.go index 2c89df4..ceabacc 100644 --- a/shortdrama/service/drama_service.go +++ b/shortdrama/service/drama_service.go @@ -209,6 +209,8 @@ func (s *dramaService) GenerateEpisode(ctx context.Context, dramaId, epId int64, APIKey: modelCfg.ChatApiKey, BaseURL: modelCfg.ChatBaseUrl, }) + genCtx = agent.WithDramaId(genCtx, d.Id) + genCtx = agent.WithDramaId(genCtx, d.Id) segDurs := calcSegDurs(d.EpisodeDuration, modelCfg) numSegments := len(segDurs) @@ -690,6 +692,7 @@ func (s *dramaService) FeedbackSegment(ctx context.Context, taskId int64, feedba APIKey: modelCfg.ChatApiKey, BaseURL: modelCfg.ChatBaseUrl, }) + genCtx = agent.WithDramaId(genCtx, d.Id) g.Log().Infof(genCtx, "第%d集第%d段根据反馈重新生成", ep.Index, task.SegmentIdx+1) if err := s.generateOneSegment(genCtx, d, ep, taskId, task.SegmentIdx, segDur, feedback); err != nil { g.Log().Errorf(genCtx, "第%d集第%d段重新生成失败: %v", ep.Index, task.SegmentIdx+1, err) @@ -759,40 +762,8 @@ func (s *dramaService) GetEpisodePollStatus(ctx context.Context, epId int64) (*d return v.Val().(*dto.EpisodePollRes), nil } - // 缓存未命中，查询 DB - tasks, err := dao.GenerationTask.ListByEpisode(ctx, epId) - if err != nil || len(tasks) == 0 { - return nil, nil - } - - status := "generating" - var errMsg string - var currentTaskId int64 - allDone := true - for _, t := range tasks { - if t.Status == "failed" { - status = "failed" - errMsg = t.ErrorMessage - break - } - if t.Status == "review" && currentTaskId == 0 { - currentTaskId = t.Id - status = "review" - } - if t.Status != "completed" { - allDone = false - } - } - if allDone { - status = "completed" - } - - return &dto.EpisodePollRes{ - Status: status, - ErrorMessage: errMsg, - Tasks: tasks, - CurrentTaskId: currentTaskId, - }, nil + // 缓存未命中 — 不查 DB，让前端继续轮询等待缓存写入 + return nil, nil } // StartVideoPoller 启动后台视频轮询器 @@ -873,12 +844,12 @@ func (s *dramaService) pollPendingVideos(ctx context.Context) { } else { g.Log().Warningf(ctx, "轮询器: 任务 %d 第%d段视频任务失败: %v（保留 task_id，不会重新提交消耗额度）", task.Id, task.SegmentIdx+1, err) } + } else if strings.Contains(err.Error(), "RUNNING") { + g.Log().Debugf(ctx, "轮询器: 任务 %d 第%d段视频正在生成中，继续等待...", task.Id, task.SegmentIdx+1) + } else { + g.Log().Warningf(ctx, "轮询器: 任务 %d 第%d段视频查询异常: %v", task.Id, task.SegmentIdx+1, err) } continue - if videoURL == "" { - g.Log().Debugf(ctx, "轮询器: 任务 %d 第%d段视频正在生成中，继续等待...", task.Id, task.SegmentIdx+1) - continue - } } // 视频就绪 — 下载到本地 @@ -950,7 +921,12 @@ func (s *dramaService) pollPendingVideos(ctx context.Context) { epUpdates[task.EpisodeId] = true } - // 更新有变化的 episode poll cache + // 刷新所有 generating 任务的 episode 缓存（不管状态有无变化，确保前端轮询不走 DB） + for _, t := range tasks { + epUpdates[t.EpisodeId] = true + } + + // 更新 episode poll cache for epId := range epUpdates { if ts, e := dao.GenerationTask.ListByEpisode(ctx, epId); e == nil { setPollCache(ctx, epId, ts)