298 lines
8.0 KiB
Go
298 lines
8.0 KiB
Go
package model
|
|
|
|
import (
|
|
"bytes"
|
|
"encoding/json"
|
|
"strings"
|
|
)
|
|
|
|
// AgentOutput Agent 最终输出的 JSON 结构
|
|
type AgentOutput struct {
|
|
Title string `json:"title"`
|
|
TotalEpisodes int `json:"total_episodes"`
|
|
TotalDuration int `json:"total_duration"`
|
|
Episodes []AgentEpisode `json:"episodes"`
|
|
Characters []AgentCharacter `json:"characters"`
|
|
}
|
|
|
|
// AgentEpisode 剧集部分
|
|
type AgentEpisode struct {
|
|
Index int `json:"index"`
|
|
Title string `json:"title"`
|
|
Duration int `json:"duration"`
|
|
Scenes []AgentScene `json:"scenes"`
|
|
}
|
|
|
|
// AgentScene 场景部分
|
|
type AgentScene 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"`
|
|
ImageUrl string `json:"imageUrl"`
|
|
}
|
|
|
|
// AgentCharacter 演员部分
|
|
type AgentCharacter struct {
|
|
Name string `json:"name"`
|
|
Description string `json:"description"`
|
|
ImageUrl string `json:"imageUrl"`
|
|
}
|
|
|
|
// extractJSON 从可能含 markdown 代码块标记或额外说明文字的文本中提取 JSON 字符串
|
|
func extractJSON(s string) string {
|
|
// 尝试直接解析
|
|
if json.Valid([]byte(s)) {
|
|
return s
|
|
}
|
|
// 尝试从 ```json ... ``` 代码块中提取（即便 JSON 不完整，也返回给 repairJSON 修复）
|
|
const jsonPrefix = "```json"
|
|
const codeFence = "```"
|
|
if idx := strings.Index(s, jsonPrefix); idx >= 0 {
|
|
start := idx + len(jsonPrefix)
|
|
if end := strings.Index(s[start:], codeFence); end >= 0 {
|
|
return strings.TrimSpace(s[start : start+end])
|
|
}
|
|
}
|
|
// 尝试从 ``` 代码块中提取
|
|
if idx := strings.Index(s, codeFence); idx >= 0 {
|
|
start := idx + len(codeFence)
|
|
if end := strings.Index(s[start:], codeFence); end >= 0 {
|
|
trimmed := strings.TrimSpace(s[start : start+end])
|
|
if json.Valid([]byte(trimmed)) {
|
|
return trimmed
|
|
}
|
|
}
|
|
}
|
|
// 最后手段：查找第一个 { 和最后一个 } 截取 JSON
|
|
if start := strings.Index(s, "{"); start >= 0 {
|
|
if end := strings.LastIndex(s, "}"); end > start {
|
|
candidate := s[start : end+1]
|
|
if json.Valid([]byte(candidate)) {
|
|
return candidate
|
|
}
|
|
}
|
|
}
|
|
return s
|
|
}
|
|
|
|
// ParseAgentOutput 解析 Agent 输出的 JSON 字符串为 SegmentOutput
|
|
// 如果解析失败，返回包含原始文本的 SegmentOutput
|
|
func ParseAgentOutput(jsonStr string, segIdx int) *SegmentOutput {
|
|
out := &SegmentOutput{
|
|
Index: segIdx,
|
|
TextOutput: jsonStr,
|
|
}
|
|
|
|
extracted := extractJSON(jsonStr)
|
|
var agentOut AgentOutput
|
|
if err := json.Unmarshal([]byte(extracted), &agentOut); err != nil {
|
|
// JSON 解析失败时尝试修复常见错误（如对话中的未转义引号）
|
|
if fixed := repairJSON(extracted); fixed != extracted {
|
|
_ = json.Unmarshal([]byte(fixed), &agentOut)
|
|
} else if len(extracted) > 0 && extracted[0] == '{' {
|
|
// 尝试修复整个 textOutput 中的 JSON
|
|
fixed2 := repairJSON(jsonStr)
|
|
extracted2 := extractJSON(fixed2)
|
|
if extracted2 != jsonStr {
|
|
_ = json.Unmarshal([]byte(extracted2), &agentOut)
|
|
}
|
|
}
|
|
if agentOut.Title == "" {
|
|
return out
|
|
}
|
|
}
|
|
|
|
// 提取演员
|
|
for _, c := range agentOut.Characters {
|
|
out.Characters = append(out.Characters, SegmentCharacter{
|
|
Name: c.Name,
|
|
Description: c.Description,
|
|
})
|
|
}
|
|
|
|
// 提取场景（优先从 episodes[0].scenes 取，兼容多种 JSON 结构）
|
|
if len(agentOut.Episodes) > 0 {
|
|
for _, s := range agentOut.Episodes[0].Scenes {
|
|
out.Scenes = append(out.Scenes, SegmentScene{
|
|
Index: s.Index,
|
|
Description: s.Description,
|
|
Lines: s.Lines,
|
|
Duration: s.Duration,
|
|
Characters: s.Characters,
|
|
})
|
|
}
|
|
}
|
|
if len(out.Scenes) == 0 {
|
|
// 兼容顶层 scenes 或嵌套在其他字段中的 scenes
|
|
var flat struct {
|
|
Scenes []AgentScene `json:"scenes"`
|
|
Segment *struct {
|
|
Scenes []AgentScene `json:"scenes"`
|
|
} `json:"segment"`
|
|
Episodes []struct {
|
|
Scenes []AgentScene `json:"scenes"`
|
|
} `json:"episodes"`
|
|
Data *struct {
|
|
Scenes []AgentScene `json:"scenes"`
|
|
} `json:"data"`
|
|
}
|
|
if err := json.Unmarshal([]byte(extracted), &flat); err == nil {
|
|
switch {
|
|
case len(flat.Scenes) > 0:
|
|
out.Scenes = extractAgentScenes(flat.Scenes)
|
|
case flat.Segment != nil && len(flat.Segment.Scenes) > 0:
|
|
out.Scenes = extractAgentScenes(flat.Segment.Scenes)
|
|
case len(flat.Episodes) > 0 && len(flat.Episodes[0].Scenes) > 0:
|
|
out.Scenes = extractAgentScenes(flat.Episodes[0].Scenes)
|
|
case flat.Data != nil && len(flat.Data.Scenes) > 0:
|
|
out.Scenes = extractAgentScenes(flat.Data.Scenes)
|
|
}
|
|
}
|
|
}
|
|
|
|
return out
|
|
}
|
|
|
|
// AgentExtractedImages 从 Agent 原始输出中提取的图片数据（图片以 base64 形式返回，由调用方负责保存到文件）
|
|
type AgentExtractedImages struct {
|
|
CharacterImages map[string]string // name → base64
|
|
SceneImages []AgentSceneImage // 按出场顺序排列（第一集）
|
|
}
|
|
|
|
// AgentSceneImage 单张场景图片数据
|
|
type AgentSceneImage struct {
|
|
Description string // 场景描述，用于库文件命名
|
|
Base64 string // base64 图片数据
|
|
}
|
|
|
|
// ExtractAgentImages 从 Agent 原始输出 JSON 中提取角色和场景图片数据，
|
|
// 与 ParseAgentOutput 独立，不依赖 SegmentOutput 结构
|
|
func ExtractAgentImages(jsonStr string) *AgentExtractedImages {
|
|
result := &AgentExtractedImages{
|
|
CharacterImages: make(map[string]string),
|
|
}
|
|
|
|
extracted := extractJSON(jsonStr)
|
|
var agentOut AgentOutput
|
|
if err := json.Unmarshal([]byte(extracted), &agentOut); err != nil {
|
|
if fixed := repairJSON(extracted); fixed != extracted {
|
|
_ = json.Unmarshal([]byte(fixed), &agentOut)
|
|
}
|
|
}
|
|
if agentOut.Title == "" {
|
|
return result
|
|
}
|
|
|
|
for _, c := range agentOut.Characters {
|
|
if c.Name != "" {
|
|
b64 := stripImagePrefix(c.ImageUrl)
|
|
if b64 != "" {
|
|
result.CharacterImages[c.Name] = b64
|
|
}
|
|
}
|
|
}
|
|
|
|
if len(agentOut.Episodes) > 0 {
|
|
for _, s := range agentOut.Episodes[0].Scenes {
|
|
b64 := stripImagePrefix(s.ImageUrl)
|
|
result.SceneImages = append(result.SceneImages, AgentSceneImage{
|
|
Description: s.Description,
|
|
Base64: b64,
|
|
})
|
|
}
|
|
}
|
|
|
|
return result
|
|
}
|
|
|
|
// stripImagePrefix 去掉 "data:image/png;base64," 等前缀
|
|
func stripImagePrefix(s string) string {
|
|
for _, prefix := range []string{
|
|
"data:image/png;base64,",
|
|
"data:image/jpeg;base64,",
|
|
"data:image/webp;base64,",
|
|
"data:image/",
|
|
} {
|
|
if len(s) > len(prefix) && s[:len(prefix)] == prefix {
|
|
return s[len(prefix):]
|
|
}
|
|
}
|
|
return s
|
|
}
|
|
|
|
// repairJSON 尝试修复 AI 输出 JSON 中的常见错误：字符串值内未转义的引号。
|
|
// AI 常在 dialogue/description 字段中使用 " 代替中文引号，导致 JSON 解析失败。
|
|
func repairJSON(s string) string {
|
|
if len(s) == 0 {
|
|
return s
|
|
}
|
|
|
|
var buf bytes.Buffer
|
|
buf.Grow(len(s) + 32)
|
|
|
|
inString := false
|
|
// 按 rune 处理以正确支持 UTF-8
|
|
runes := []rune(s)
|
|
n := len(runes)
|
|
|
|
for i := 0; i < n; i++ {
|
|
ch := runes[i]
|
|
|
|
if ch == '\\' && inString {
|
|
buf.WriteRune(ch)
|
|
if i+1 < n {
|
|
i++
|
|
buf.WriteRune(runes[i])
|
|
}
|
|
continue
|
|
}
|
|
|
|
if ch == '"' {
|
|
if !inString {
|
|
inString = true
|
|
buf.WriteRune(ch)
|
|
continue
|
|
}
|
|
|
|
// 在字符串内遇到 "，判断是否为结束符
|
|
// 向后查找下一个非空白字符
|
|
j := i + 1
|
|
for j < n && (runes[j] == ' ' || runes[j] == '\t' || runes[j] == '\n' || runes[j] == '\r') {
|
|
j++
|
|
}
|
|
|
|
if j < n && (runes[j] == ',' || runes[j] == ']' || runes[j] == '}' || runes[j] == ':') {
|
|
// 是 JSON 结束引号
|
|
inString = false
|
|
buf.WriteRune(ch)
|
|
} else {
|
|
// 是内容中的引号，转义
|
|
buf.WriteString(`\"`)
|
|
}
|
|
continue
|
|
}
|
|
|
|
buf.WriteRune(ch)
|
|
}
|
|
|
|
return buf.String()
|
|
}
|
|
|
|
// extractAgentScenes 将 AgentScene 切片转换为 SegmentScene 切片
|
|
func extractAgentScenes(src []AgentScene) []SegmentScene {
|
|
dst := make([]SegmentScene, len(src))
|
|
for i, s := range src {
|
|
dst[i] = SegmentScene{
|
|
Index: s.Index,
|
|
Description: s.Description,
|
|
Lines: s.Lines,
|
|
Duration: s.Duration,
|
|
Characters: s.Characters,
|
|
}
|
|
}
|
|
return dst
|
|
}
|