375 lines
10 KiB
Go
375 lines
10 KiB
Go
package agent
|
||
|
||
import (
|
||
"context"
|
||
"encoding/base64"
|
||
"encoding/json"
|
||
"fmt"
|
||
"os"
|
||
"path/filepath"
|
||
"strings"
|
||
|
||
"video-factory/shortdrama/dao"
|
||
|
||
"github.com/cloudwego/eino/schema"
|
||
"github.com/gogf/gf/v2/frame/g"
|
||
)
|
||
|
||
// ToolInfo 工具定义
|
||
type ToolInfo struct {
|
||
Name string
|
||
Description string
|
||
Parameters map[string]any
|
||
Func func(ctx context.Context, args map[string]any) (string, error)
|
||
}
|
||
|
||
// ToEinoToolInfo 转换为 Eino 的 ToolInfo 格式
|
||
func (t *ToolInfo) ToEinoToolInfo() *schema.ToolInfo {
|
||
params := make(map[string]*schema.ParameterInfo)
|
||
if paramsMap, ok := t.Parameters["properties"].(map[string]any); ok {
|
||
for key, val := range paramsMap {
|
||
if prop, ok := val.(map[string]any); ok {
|
||
desc, _ := prop["description"].(string)
|
||
pType, _ := prop["type"].(string)
|
||
params[key] = &schema.ParameterInfo{
|
||
Type: schema.DataType(pType),
|
||
Desc: desc,
|
||
Required: false,
|
||
}
|
||
}
|
||
}
|
||
}
|
||
if required, ok := t.Parameters["required"].([]any); ok {
|
||
for _, r := range required {
|
||
if rStr, ok := r.(string); ok {
|
||
if p, exists := params[rStr]; exists {
|
||
p.Required = true
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
return &schema.ToolInfo{
|
||
Name: t.Name,
|
||
Desc: t.Description,
|
||
ParamsOneOf: schema.NewParamsOneOfByParams(params),
|
||
}
|
||
}
|
||
|
||
// GetTools 获取 ReAct Agent 可用的所有工具
|
||
func GetTools() []*ToolInfo {
|
||
return []*ToolInfo{
|
||
parseScriptTool(),
|
||
analyzeScriptForEpisodeTool(),
|
||
generateSceneImageTool(),
|
||
}
|
||
}
|
||
|
||
// ==================== Tool 1: 解析剧本 ====================
|
||
|
||
func parseScriptTool() *ToolInfo {
|
||
return &ToolInfo{
|
||
Name: "parse_script",
|
||
Description: "将原始剧本文本解析为结构化的剧集列表，支持用 --- 分隔的多集剧本",
|
||
Parameters: map[string]any{
|
||
"type": "object",
|
||
"properties": map[string]any{
|
||
"raw_script": map[string]any{
|
||
"type": "string",
|
||
"description": "原始剧本文本，多集用 --- 分隔",
|
||
},
|
||
},
|
||
"required": []string{"raw_script"},
|
||
},
|
||
Func: func(ctx context.Context, args map[string]any) (string, error) {
|
||
rawScript, _ := args["raw_script"].(string)
|
||
if rawScript == "" {
|
||
return "", fmt.Errorf("剧本内容不能为空")
|
||
}
|
||
|
||
// 按 --- 分割多集
|
||
episodeTexts := strings.Split(rawScript, "---")
|
||
type episodeInfo struct {
|
||
Index int `json:"index"`
|
||
Title string `json:"title"`
|
||
Script string `json:"script"`
|
||
}
|
||
var episodes []episodeInfo
|
||
|
||
for i, text := range episodeTexts {
|
||
text = strings.TrimSpace(text)
|
||
if text == "" {
|
||
continue
|
||
}
|
||
lines := strings.SplitN(text, "\n", 2)
|
||
title := strings.TrimSpace(lines[0])
|
||
// 去掉可能的序号前缀如 "第1集"、"第一集"、"Episode 1" 等
|
||
title = cleanEpisodeTitle(title)
|
||
content := ""
|
||
if len(lines) > 1 {
|
||
content = strings.TrimSpace(lines[1])
|
||
} else {
|
||
content = title
|
||
}
|
||
episodes = append(episodes, episodeInfo{
|
||
Index: i + 1,
|
||
Title: title,
|
||
Script: content,
|
||
})
|
||
}
|
||
|
||
result, _ := json.Marshal(map[string]any{
|
||
"episodes": episodes,
|
||
"total_episodes": len(episodes),
|
||
})
|
||
return string(result), nil
|
||
},
|
||
}
|
||
}
|
||
|
||
// ==================== Tool 2: 分析单集剧本 ====================
|
||
|
||
func analyzeScriptForEpisodeTool() *ToolInfo {
|
||
return &ToolInfo{
|
||
Name: "analyze_script_for_episode",
|
||
Description: "分析单集剧本内容，根据时长将剧本拆分为多个场景，识别出场演员及画面描述",
|
||
Parameters: map[string]any{
|
||
"type": "object",
|
||
"properties": map[string]any{
|
||
"episode_index": map[string]any{
|
||
"type": "integer",
|
||
"description": "剧集索引（从1开始）",
|
||
},
|
||
"episode_title": map[string]any{
|
||
"type": "string",
|
||
"description": "本集标题",
|
||
},
|
||
"script_content": map[string]any{
|
||
"type": "string",
|
||
"description": "本集剧本内容",
|
||
},
|
||
"duration": map[string]any{
|
||
"type": "integer",
|
||
"description": "本集总时长（秒）",
|
||
},
|
||
"characters": map[string]any{
|
||
"type": "string",
|
||
"description": "演员列表JSON，格式：[{\"name\":\"演员名\",\"description\":\"演员描述\"}]",
|
||
},
|
||
},
|
||
"required": []string{"episode_index", "script_content", "duration", "characters"},
|
||
},
|
||
Func: func(ctx context.Context, args map[string]any) (string, error) {
|
||
scriptContent, _ := args["script_content"].(string)
|
||
durationFloat, _ := args["duration"].(float64)
|
||
duration := int(durationFloat)
|
||
|
||
if scriptContent == "" {
|
||
return "", fmt.Errorf("剧本内容不能为空")
|
||
}
|
||
if duration <= 0 {
|
||
duration = 60 // 默认60秒
|
||
}
|
||
|
||
episodeIndex, _ := args["episode_index"].(float64)
|
||
title, _ := args["episode_title"].(string)
|
||
|
||
// 按空行或场景标记分割场景
|
||
sceneTexts := strings.Split(scriptContent, "\n\n")
|
||
type sceneInfo struct {
|
||
Index int `json:"index"`
|
||
Description string `json:"description"`
|
||
Lines string `json:"lines"`
|
||
Duration int `json:"duration"`
|
||
Characters []string `json:"characters"`
|
||
VisualDesc string `json:"visualDesc"`
|
||
}
|
||
var scenes []sceneInfo
|
||
|
||
totalScenes := len(sceneTexts)
|
||
if totalScenes == 0 {
|
||
totalScenes = 1
|
||
sceneTexts = []string{scriptContent}
|
||
}
|
||
|
||
// 推测出场演员
|
||
var characters []string
|
||
if charsRaw, ok := args["characters"].(string); ok && charsRaw != "" {
|
||
var chars []struct {
|
||
Name string `json:"name"`
|
||
}
|
||
json.Unmarshal([]byte(charsRaw), &chars)
|
||
for _, c := range chars {
|
||
characters = append(characters, c.Name)
|
||
}
|
||
}
|
||
|
||
perSceneDuration := duration / totalScenes
|
||
remainder := duration % totalScenes
|
||
|
||
for i, text := range sceneTexts {
|
||
text = strings.TrimSpace(text)
|
||
if text == "" {
|
||
continue
|
||
}
|
||
|
||
sceneDur := perSceneDuration
|
||
if i < remainder {
|
||
sceneDur++
|
||
}
|
||
|
||
// 提取第一行作为场景描述
|
||
lines := strings.SplitN(text, "\n", 2)
|
||
desc := strings.TrimSpace(lines[0])
|
||
content := ""
|
||
if len(lines) > 1 {
|
||
content = strings.TrimSpace(lines[1])
|
||
} else {
|
||
content = desc
|
||
}
|
||
|
||
// 匹配出场演员
|
||
var sceneChars []string
|
||
for _, c := range characters {
|
||
if strings.Contains(text, c) {
|
||
sceneChars = append(sceneChars, c)
|
||
}
|
||
}
|
||
|
||
visualDesc := fmt.Sprintf("场景%d：%s，画面风格根据剧本内容自动生成", i+1, desc)
|
||
|
||
scenes = append(scenes, sceneInfo{
|
||
Index: i + 1,
|
||
Description: desc,
|
||
Lines: content,
|
||
Duration: sceneDur,
|
||
Characters: sceneChars,
|
||
VisualDesc: visualDesc,
|
||
})
|
||
}
|
||
|
||
result, _ := json.Marshal(map[string]any{
|
||
"episode_index": int(episodeIndex),
|
||
"episode_title": title,
|
||
"total_scenes": len(scenes),
|
||
"total_duration": duration,
|
||
"scenes": scenes,
|
||
})
|
||
return string(result), nil
|
||
},
|
||
}
|
||
}
|
||
|
||
// ==================== Tool 4: 生成场景图 ====================
|
||
|
||
func generateSceneImageTool() *ToolInfo {
|
||
return &ToolInfo{
|
||
Name: "generate_scene_image",
|
||
Description: "根据场景的画面描述获取场景图片，返回图片的base64编码数据。优先从已有场景库中读取图片，不存在时返回空。",
|
||
Parameters: map[string]any{
|
||
"type": "object",
|
||
"properties": map[string]any{
|
||
"episode_index": map[string]any{
|
||
"type": "integer",
|
||
"description": "剧集索引",
|
||
},
|
||
"scene_index": map[string]any{
|
||
"type": "integer",
|
||
"description": "场景索引",
|
||
},
|
||
"visual_description": map[string]any{
|
||
"type": "string",
|
||
"description": "画面描述（场景设定、演员动作、镜头角度等）",
|
||
},
|
||
"style": map[string]any{
|
||
"type": "string",
|
||
"description": "整体风格",
|
||
},
|
||
},
|
||
"required": []string{"visual_description", "style"},
|
||
},
|
||
Func: func(ctx context.Context, args map[string]any) (string, error) {
|
||
visualDesc, _ := args["visual_description"].(string)
|
||
style, _ := args["style"].(string)
|
||
episodeIdx, _ := args["episode_index"].(float64)
|
||
sceneIdx, _ := args["scene_index"].(float64)
|
||
|
||
prompt := fmt.Sprintf("画面描述：%s，风格：%s", visualDesc, style)
|
||
_ = 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{
|
||
"episode_index": int(episodeIdx),
|
||
"scene_index": int(sceneIdx),
|
||
"image_base64": imgBase64,
|
||
"status": "success",
|
||
})
|
||
return string(result), nil
|
||
},
|
||
}
|
||
}
|
||
|
||
// ==================== 工具函数 ====================
|
||
|
||
func cleanEpisodeTitle(title string) string {
|
||
prefixes := []string{"第", "Episode", "episode", "EP"}
|
||
for _, p := range prefixes {
|
||
if strings.HasPrefix(title, p) {
|
||
// 去掉序号前缀后取标题部分
|
||
parts := strings.SplitN(title, " ", 2)
|
||
if len(parts) > 1 {
|
||
return parts[1]
|
||
}
|
||
parts = strings.SplitN(title, "：", 2)
|
||
if len(parts) > 1 {
|
||
return parts[1]
|
||
}
|
||
}
|
||
}
|
||
return title
|
||
}
|
||
|
||
func readImageFileAsBase64(path string) (string, error) {
|
||
data, err := os.ReadFile(path)
|
||
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"
|
||
}
|
||
return "data:" + mime + ";base64," + base64.StdEncoding.EncodeToString(data), nil
|
||
}
|