red-future/video-factory
3
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

1

Browse Source
This commit is contained in:
admin
2026-07-01 15:50:13 +08:00
parent fba5dc8e2a
commit 503545fba6
4 changed files with 77 additions and 127 deletions
Download Patch File Download Diff File Expand all files Collapse all files
short_drama.db
BIN
View File
Binary file not shown.
shortdrama/agent/chat_model.go
+16 -13
View File
@@ -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 != "" {
shortdrama/agent/tools.go
+46 -75
View File
@@ -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
}
shortdrama/service/drama_service.go
+15 -39
View File
@@ -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)